Original: E QUIVARIANT N ORMALIZING F LOWS FOR P OINT P ROCESSES AND S ETS
Revised: E QUIVARIANT N ORMALIZING F LOWS FOR P OINT P ROCESSES AND S ETS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: A point process describes how random sets of exchangeable points are generated.
Revised: A point process describes how random sets of exchangeable points are generated.

Original: The points usually inﬂuence the positions of each other via attractive and repulsive forces.
Revised: The points usually inﬂuence the positions of each other via attractive and repulsive forces.

Original: To model this behavior, it is enough to transform the samples from the uniform process with a sufﬁciently complex equivariant function.
Revised: To model this behavior, it is enough to transform the samples from the uniform process with a sufﬁciently complex equivariant function.

Original: However, learning the parameters of the resulting process is challenging since the likelihood is hard to estimate and often intractable.
Revised: However, learning the parameters of the resulting process is challenging since the likelihood is hard to estimate and often intractable.

Original: This leads us to our proposed model –
Revised: This leads us to our proposed model –

Original: C ONFET .
Revised: C ONFET .

Original: Based on continuous normalizing ﬂows, it allows arbitrary interactions between points while having tractable likelihood.
Revised: Based on continuous normalizing ﬂows, it allows arbitrary interactions between points while having tractable likelihood.

Original: Experiments on various real and synthetic datasets show the improved performance of our new scalable approach.
Revised: Experiments on various real and synthetic datasets show the improved performance of our new scalable approach.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Many domains contain unordered data with a variable number of elements.
Revised: Many domains contain unordered data with a variable number of elements.

Original: The lack of ordering, also known as exchangeability, can be found in locations of cellular stations, locations of trees in a forest, point clouds, items in a shopping cart etc.
Revised: The lack of ordering, also known as exchangeability, can be found in locations of cellular stations, locations of trees in a forest, point clouds, items in a shopping cart etc.

Original: This kind of data is represented with sets that are randomly generated from some underlying process that we wish to uncover.
Revised: This kind of data is represented with sets that are randomly generated from some underlying process that we wish to uncover.

Original: We choose to model this with spatial point processes, generative models whose realizations are sets of points.
Revised: We choose to model this with spatial point processes, generative models whose realizations are sets of points.

Original: Perhaps the simplest non-trivial model is an inhomogeneous Poisson process.
Revised: Perhaps the simplest non-trivial model is an inhomogeneous Poisson process.

Original: The locations of the points are assumed to be generated i.i.d.
Revised: The locations of the points are assumed to be generated i.i.d.

Original: from some density (Chiu et al., 2013).
Revised: from some density (Chiu et al., 2013).

Original: By simply modeling this density we can evaluate the likelihood and draw samples.
Revised: By simply modeling this density we can evaluate the likelihood and draw samples.

Original: We can do this easily with normalizing ﬂows (Germain et al., 2015).
Revised: We can do this easily with normalizing ﬂows (Germain et al., 2015).

Original: The process is then deﬁned with a transformation of samples from a simple distribution to samples in the target distribution.
Revised: The process is then deﬁned with a transformation of samples from a simple distribution to samples in the target distribution.

Original: However, the i.i.d.
Revised: However, the i.i.d.

Original: property is often wrong because the presence of one object will inﬂuence the distribution of the others.
Revised: property is often wrong because the presence of one object will inﬂuence the distribution of the others.

Original: For example, short trees grow near each other, but are inhibited by taller trees (Ogata & Tanemura, 1985).
Revised: For example, short trees grow near each other, but are inhibited by taller trees (Ogata & Tanemura, 1985).

Original: As an example we can generate points on the interval (0 , 1) in the following way: we ﬁrst ﬂip a coin to decide whether to sample inside or outside of the interval (1 / 4 , 3 / 4) ; then sample two points x 1 , x 2 uniformly on the chosen subset.
Revised: As an example we can generate points on the interval (0 , 1) in the following way: we ﬁrst ﬂip a coin to decide whether to sample inside or outside of the interval (1 / 4 , 3 / 4) ; then sample two points x 1 , x 2 uniformly on the chosen subset.

Original: Although the marginals p ( x 1 ) and p ( x 2 ) are uniformly distributed, knowing the position of one point gives us information about the other.
Revised: Although the marginals p ( x 1 ) and p ( x 2 ) are uniformly distributed, knowing the position of one point gives us information about the other.

Original: Therefore, we should not model this process as if the points are independent, but model the joint distribution p ( x 1 , x 2 ) , with a constraint that p is symmetric to permutations (Figure 1).
Revised: Therefore, we should not model this process as if the points are independent, but model the joint distribution p ( x 1 , x 2 ) , with a constraint that p is symmetric to permutations (Figure 1).

Original: Unfortunately, when we include interactions between the points, the problem becomes signiﬁcantly harder because the likelihood is intractable (Besag, 1975).
Revised: Unfortunately, when we include interactions between the points, the problem becomes signiﬁcantly harder because the likelihood is intractable (Besag, 1975).

Original: [Figure]
Revised: [Figure]

Original: [Equation]
Revised: [Equation]

Original: [Figure]
Revised: [Figure]

Original: In this work, we present a way to solve this issue by using continuous normalizing ﬂows that allow unrestricted transformations of points, with interactions between them.
Revised: In this work, we present a way to solve this issue by using continuous normalizing ﬂows that allow unrestricted transformations of points, with interactions between them.

Original: This way, a very hard problem of likelihood estimation suddenly becomes tractable.
Revised: This way, a very hard problem of likelihood estimation suddenly becomes tractable.

Original: Our approach transforms simple processes into complex ones, by transforming their samples with expressive functions (Figure 1).
Revised: Our approach transforms simple processes into complex ones, by transforming their samples with expressive functions (Figure 1).

Original: Our main contributions are the following:
Revised: Our main contributions are the following:

Original: • We reinterpret and unify existing techniques for modeling point processes and exchangeable data with normalizing ﬂows.
Revised: • We reinterpret and unify existing techniques for modeling point processes and exchangeable data with normalizing ﬂows.

Original: • We propose a new generative model (C ONFET ) that allows interactions between points and can be trained with maximum likelihood.
Revised: • We propose a new generative model (C ONFET ) that allows interactions between points and can be trained with maximum likelihood.

Original: The extensive experiments show that it outperforms other approaches, while remaining efﬁcient and scalable.
Revised: The extensive experiments show that it outperforms other approaches, while remaining efﬁcient and scalable.

Original:  P OINT PROCESSES AND RANDOM SETS
Revised:  P OINT PROCESSES AND RANDOM SETS

Original: Realizations of a ﬁnite point process on a bounded region B ⊂ R
Revised: Realizations of a ﬁnite point process on a bounded region B ⊂ R

Original: d
Revised: d

Original: are ﬁnite sets of points X
Revised: are ﬁnite sets of points X

Original: = { x 1 , . .
Revised: = { x 1 , . .

Original: , x n } , x i ∈ B .
Revised: , x n } , x i ∈ B .

Original: A point process is simple if no two points fall at exactly the same place.
Revised: A point process is simple if no two points fall at exactly the same place.

Original: In practice, point processes are usually both ﬁnite and simple.
Revised: In practice, point processes are usually both ﬁnite and simple.

Original: One way to construct a general point process is by deﬁning a discrete distribution p ( n ) for the number of points and a symmetric probability density p
Revised: One way to construct a general point process is by deﬁning a discrete distribution p ( n ) for the number of points and a symmetric probability density p

Original: ( X ) on B n for their locations (Daley &
Revised: ( X ) on B n for their locations (Daley &

Original: Vere-Jones, 2007).
Revised: Vere-Jones, 2007).

Original: Symmetry requirement comes from the fact that the probability of a sequence ( x π (1) , . . .
Revised: The symmetry requirement comes from the fact that the probability of a sequence ( x π (1) , . . .

Original: , x π ( n ) ) is the same for any permutation of the elements π — the points are exchangeable.
Revised: , x π ( n ) ) is the same for any permutation of the elements π — the points are exchangeable.

Original: Not knowing the order of the points can be solved trivially by averaging any density over all n ! permutations.
Revised: Not knowing the order of the points can be solved trivially by averaging any density over all n ! permutations.

Original: Another approach is to impose a canonical order, e.g. sorting the points by one of the dimensions.
Revised: Another approach is to impose a canonical order, e.g. sorting the points by one of the dimensions.

Original: Then the probability of observing a set of exchangeable elements is deﬁned w.r.t.
Revised: Then the probability of observing a set of exchangeable elements is deﬁned w.r.t.

Original: the joint probability of an order statistic x (1) < · · · < x ( n ) (Casella & Berger, 2002):
Revised: the joint probability of an order statistic x (1) < · · · < x ( n ) (Casella & Berger, 2002):

Original: [Equation]
Revised: [Equation]

Original: A traditional way to deﬁne a point process is with an intensity function that assigns a non-negative value to every subset of B , corresponding to a number of points we expect to see there (Møller &
Revised: A traditional way to deﬁne a point process is with an intensity function that assigns a non-negative value to every subset of B , corresponding to a number of points we expect to see there (Møller &

Original: Waagepetersen, 2003).
Revised: Waagepetersen, 2003).

Original: An example is a homogeneous Poisson process with constant intensity λ .
Revised: An example is a homogeneous Poisson process with constant intensity λ .

Original: To generate a new realization, we ﬁrst sample n ∼ Pois ( λ ) , and then sample n points uniformly on B .
Revised: To generate a new realization, we ﬁrst sample n ∼ Pois ( λ ) , and then sample n points uniformly on B .

Original: If we deﬁne the intensity as a function of position λ ( x ) , we get an inhomogeneous Poisson process which is equivalent to deﬁning the un-normalized probability density function on B .
Revised: If we deﬁne the intensity as a function of position λ ( x ) , we get an inhomogeneous Poisson process which is equivalent to deﬁning the un-normalized probability density function on B .

Original: Now n follows
Revised: Now n follows

Original: Pois (Λ) , where Λ is the total intensity (cid:82) λ ( x ) d x .
Revised: Pois (Λ) , where Λ is the total intensity (cid:82) λ ( x ) d x .

Original: We get the density at a location x by normalizing the intensity p ( x ) =
Revised: We get the density at a location x by normalizing the intensity p ( x ) =

Original: λ ( x ) / Λ .
Revised: λ ( x ) / Λ .

Original: Combining the distribution of the number of points with the distribution of their locations gives us a well known formula for the likelihood of an inhomogeneous
Revised: Combining the distribution of the number of points with the distribution of their locations gives us a well known formula for the likelihood of an inhomogeneous

Original: Poisson process (Daley & Vere-Jones, 2007, eq.
Revised: Poisson process (Daley & Vere-Jones, 2007, eq.

Original: [Equation]
Revised: [Equation]

Original: Instead of modeling λ ( x ) , we can model p ( x ) directly to avoid estimating the integral, without losing generality (Yuan et al., 2020).
Revised: Instead of modeling λ ( x ) , we can model p ( x ) directly to avoid estimating the integral, without losing generality (Yuan et al., 2020).

Original: This shift in the perspective from intensity to probability density function allows us to utilize rich existing methods from density estimation.
Revised: This shift in the perspective from intensity to probability density function allows us to utilize rich existing methods from density estimation.

Original: An extension of an inhomogeneous process that allows interactions between points deﬁnes the conditional intensity λ ( x i | X ) (Papangelou, 1974).
Revised: An extension of an inhomogeneous process that allows interactions between points deﬁnes the conditional intensity λ ( x i | X ) (Papangelou, 1974).

Original: This may be interpreted as the conditional probability of having a point at x i given the rest of the process coincides with X .
Revised: This may be interpreted as the conditional probability of having a point at x i given the rest of the process coincides with X .

Original: The likelihood is not tractable anymore so previous works used pseudolikelihood instead, replacing λ ( x ) with λ ( x | X ) in
Revised: The likelihood is not tractable anymore so previous works used pseudolikelihood instead, replacing λ ( x ) with λ ( x | X ) in

Original: Eq. 2 (Besag, 1975; Baddeley & Turner, 2000).
Revised: Eq. 2 (Besag, 1975; Baddeley & Turner, 2000).

Original: One example of such a process is a clustering process (Neyman & Scott, 1958) that generates the points in two steps.
Revised: One example of such a process is a clustering process (Neyman & Scott, 1958) that generates the points in two steps.

Original: First, we sample the cluster centers from a Poisson process, then we sample the ﬁnal points from normal distributions centered around cluster positions.
Revised: First, we sample the cluster centers from a Poisson process, then we sample the ﬁnal points from normal distributions centered around cluster positions.

Original: In contrast, a repulsion process (Matérn, 2013) generates initial points from a uniform process and removes those that have neighbors inside radius R .
Revised: In contrast, a repulsion process (Matérn, 2013) generates initial points from a uniform process and removes those that have neighbors inside radius R .

Original: A different perspective on point processes is viewing them as random sets (Baddeley et al., 2006).
Revised: A different perspective on point processes is viewing them as random sets (Baddeley et al., 2006).

Original: Recent approaches for modeling sets (Korshunova et al., 2018; Bender et al., 2020) invoke de
Revised: Recent approaches for modeling sets (Korshunova et al., 2018; Bender et al., 2020) invoke de

Original: Finetti’s theorem (De Finetti, 1937) which states that the probability of an inﬁnite exchangeable sequence is a mixture of i.i.d.
Revised: Finetti’s theorem (De Finetti, 1937) which states that the probability of an inﬁnite exchangeable sequence is a mixture of i.i.d.

Original: processes:
Revised: processes:

Original: [Equation]
Revised: [Equation]

Original: here written for a ﬁnite subset (O’Neill, 2009, Theorem 1).
Revised: here written for a ﬁnite subset (O’Neill, 2009, Theorem 1).

Original: We can use Eq.
Revised: We can use Eq.

Original: 3 to construct a latent variable model where points are independent of each other given z .
Revised: 3 to construct a latent variable model where points are independent of each other given z .

Original: Since the integral is intractable, we will have to resort to approximate inference.
Revised: Since the integral is intractable, we will have to resort to approximate inference.

Original: Equations 1, 2 and 3 give us different ways to obtain the parameters θ of a point process by maximizing the likelihood p θ ( X ) .
Revised: Equations 1, 2 and 3 give us different ways to obtain the parameters θ of a point process by maximizing the likelihood p θ ( X ) .

Original: However, the approaches that arise from them are equivalent in the sense they all model the same symmetric density.
Revised: However, the approaches that arise from them are equivalent in the sense they all model the same symmetric density.

Original: Therefore, omitting terms related to cardinality (Vo et al., 2018), like p ( n ) , does not matter when comparing them.
Revised: Therefore, omitting terms related to cardinality (Vo et al., 2018), like p ( n ) , does not matter when comparing them.

Original: Further discussion is in  A.1.
Revised: Further discussion is in Appendix A.1.

Original: Our goal is to have a generative model deﬁned with θ that can produce sets as realizations.
Revised: Our goal is to have a generative model deﬁned with θ that can produce sets as realizations.

Original: One way to generate realizations of different point processes is to transform the points drawn from a uniform Poisson process with an invertible function.
Revised: One way to generate realizations of different point processes is to transform the points drawn from a uniform Poisson process with an invertible function.

Original: The only requirement is that the process remains locally ﬁnite (Baddeley et al., 2006, Section 1.7).
Revised: The only requirement is that the process remains locally ﬁnite (Baddeley et al., 2006, Section 1.7).

Original: To keep track of the likelihood we will use the change of variables formula from the normalizing ﬂow framework.
Revised: To keep track of the likelihood we will use the change of variables formula from the normalizing ﬂow framework.

Original:  U NIFYING POINT PROCESS MODELS WITH NORMALIZING FLOWS
Revised:  U NIFYING POINT PROCESS MODELS WITH NORMALIZING FLOWS

Original: In the following section we review the models for point processes and random sets that use likelihoods listed in Section 2.
Revised: In the following section we review the models for point processes and random sets that use likelihoods listed in Section 2.

Original: These models can be found fully or partially in the literature on point processes and modeling exchangeable data.
Revised: These models can be found fully or partially in the literature on point processes and modeling exchangeable data.

Original: We unify them here under the common umbrella of normalizing ﬂows.
Revised: We unify them here under the common umbrella of normalizing ﬂows.

Original: Their limitations form the foundation for our proposed solution in Section 4.
Revised: Their limitations form the foundation for our proposed solution in Section 4.

Original: A normalizing ﬂow is a generative model that deﬁnes a complex distribution as a series of invertible smooth transformations of the initial random variable (Germain et al., 2015).
Revised: A normalizing ﬂow is a generative model that deﬁnes a complex distribution as a series of invertible smooth transformations of the initial random variable (Germain et al., 2015).

Original: That means, if we apply a function f : R d → R d to the random variable z ∼ q ( z ) , where f is invertible and differentiable, we can get the log-density of x = f ( z ) by calculating the change of variables formula:
Revised: That means, if we apply a function f : R d → R d to the random variable z ∼ q ( z ) , where f is invertible and differentiable, we can get the log-density of x = f ( z ) by calculating the change of variables formula:

Original: [Equation]
Revised: [Equation]

Original: To sample from p ( x )
Revised: To sample from p ( x )

Original: we ﬁrst sample z from q ( z ) , then apply the forward transformation z
Revised: we ﬁrst sample z from q ( z ) , then apply the forward transformation z

Original: (cid:55)→ x .
Revised: (cid:55)→ x .

Original: To estimate density p ( x ) for a given sample, we apply the inverse transformation x (cid:55)→ z and use the above formula.
Revised: To estimate density p ( x ) for a given sample, we apply the inverse transformation x (cid:55)→ z and use the above formula.

Original: We can also use a composition of functions f ( x ) =
Revised: We can also use a composition of functions f ( x ) =

Original: ( f 1 ◦ · · · ◦ f k )( z ) to deﬁne more complex distributions, using Eq. 4 at every step.
Revised: ( f 1 ◦ · · · ◦ f k )( z ) to deﬁne more complex distributions, using Eq. 4 at every step.

Original: The main challenge of normalizing ﬂow models is deﬁning the function f .
Revised: The main challenge of normalizing ﬂow models is deﬁning the function f .

Original: First, we want to have a way to efﬁciently calculate the inverse f − 1 .
Revised: First, we want to have a way to efﬁciently calculate the inverse f − 1 .

Original: Second, calculating the determinant of the Jacobian becomes prohibitively expensive as the dimension d grows.
Revised: Second, calculating the determinant of the Jacobian becomes prohibitively expensive as the dimension d grows.

Original: We often use f with a special Jacobian form, e.g. the determinant of a lower triangular matrix is simply the product of its diagonal entries.
Revised: We often use f with a special Jacobian form, e.g. the determinant of a lower triangular matrix is simply the product of its diagonal entries.

Original: In all of our models we parametrize f − 1 : B n → B n that transforms the input set X to a set
Revised: In all of our models we parametrize f − 1 : B n → B n that transforms the input set X to a set

Original: Z , with its corresponding base density.
Revised: Z , with its corresponding base density.

Original: To make sure p ( X ) is symmetric we have to satisfy two conditions.
Revised: To make sure p ( X ) is symmetric we have to satisfy two conditions.

Original: The base density q ( Z ) should be permutation invariant.
Revised: The base density q ( Z ) should be permutation invariant.

Original: The mapping f − 1 should be permutation equivariant, meaning any permutation of the input permutes the output in the same way (Papamakarios et al., 2019, Section 5.6, Lemma 1).
Revised: The mapping f − 1 should be permutation equivariant, meaning any permutation of the input permutes the output in the same way (Papamakarios et al., 2019, Section 5.6, Lemma 1).

Original: We will see different ways to enforce this, along with how to handle varying input size n .
Revised: We will see different ways to enforce this, along with how to handle varying input size n .

Original: A detailed implementation description is in  B.1.
Revised: A detailed implementation description is in Appendix B.1.

Original: Inhomogeneous Poisson process.
Revised: Inhomogeneous Poisson process.

Original: Given an observed set X we would like to use Eq.
Revised: Given an observed set X we would like to use Eq.

Original: 2 for maximum likelihood training.
Revised: 2 for maximum likelihood training.

Original: Since the model assumes the independence between points we can deﬁne a normalizing ﬂow on B , transform each point independently and get p ( X ) =
Revised: Since the model assumes the independence between points we can deﬁne a normalizing ﬂow on B , transform each point independently and get p ( X ) =

Original: (cid:81) i p ( x i ) with Eq.
Revised: (cid:81) i p ( x i ) with Eq.

Original: Invariance is achieved trivially.
Revised: Invariance is achieved trivially.

Original: In the experiments, the ﬂow is parametrized with coupling layers (Dinh et al., 2017) that have tractable inverse and determinant.
Revised: In the experiments, the ﬂow is parametrized with coupling layers (Dinh et al., 2017) that have tractable inverse and determinant.

Original: Inside of coupling layers, elementwise functions are used, in particular splines with K knots (Durkan et al., 2019).
Revised: Inside of coupling layers, elementwise functions are used, in particular splines with K knots (Durkan et al., 2019).

Original: Autoregressive model with canonical ordering.
Revised: Autoregressive model with canonical ordering.

Original: The idea to order the points by some arbitrary dimension allows us to use Eq.
Revised: The idea to order the points by some arbitrary dimension allows us to use Eq.

Original: 1 for training.
Revised: 1 for training.

Original: The ﬂow acts on densities over B n .
Revised: The ﬂow acts on densities over B n .

Original: Since the ordering of the elements is known, a common approach is to use an autoregressive transformations that mimic the conditional probability formula by conditioning each element on all the previous, z i = f − 1 θ ( x i | x 1 , . . .
Revised: Since the ordering of the elements is known, a common approach is to use an autoregressive transformations that mimic the conditional probability formula by conditioning each element on all the previous, z i = f − 1 θ ( x i | x 1 , . . .

Original: , x i − 1 ) (Kingma et al., 2016).
Revised: , x i − 1 ) (Kingma et al., 2016).

Original: The ordered sequence of points ( x 1 , . . .
Revised: The ordered sequence of points ( x 1 , . . .

Original: , x n ) is passed to f − 1 θ , parametrized with a recurrent neural network.
Revised: , x n ) is passed to f − 1 θ , parametrized with a recurrent neural network.

Original: This scales to different set sizes n .
Revised: This scales to different set sizes n .

Original: Stacking multiple such layers together with set coupling gives a model like in Bender et al.
Revised: Stacking multiple such layers together with set coupling gives a model like in Bender et al.

Original: Variational autoencoder and exchangeability.
Revised: Variational autoencoder and exchangeability.

Original: Another way to ensure exchangeability is to introduce a latent variable z such that the points x i are conditionally independent given z .
Revised: Another way to ensure exchangeability is to introduce a latent variable z such that the points x i are conditionally independent given z .

Original: Further, we want to maximize the log-likelihood log p ( X ) =
Revised: Further, we want to maximize the log-likelihood log p ( X ) =

Original: log (cid:82) p ( X, z ) d z (Eq. 3)
Revised: log (cid:82) p ( X, z ) d z (Eq. 3)

Original: (Yuan et al., 2020).
Revised: (Yuan et al., 2020).

Original: However, the true posterior p ( z | X ) is intractable so we approximate it with a variational distribution
Revised: However, the true posterior p ( z | X ) is intractable so we approximate it with a variational distribution

Original: q ( z | X ) .
Revised: q ( z | X ) .

Original: We can now maximize the likelihood by maximizing the evidence lower bound:
Revised: We can now maximize the likelihood by maximizing the evidence lower bound:

Original: [Equation]
Revised: [Equation]

Original: Using amortized inference, parameters of q ( z | X ) are deﬁned as a function of X , and taking gradients w.r.t.
Revised: Using amortized inference, parameters of q ( z | X ) are deﬁned as a function of X , and taking gradients w.r.t.

Original: the samples is enabled with the reparametrization trick (Kingma & Welling, 2014).
Revised: the samples is enabled with the reparametrization trick (Kingma & Welling, 2014).

Original: We use a standard posterior approximation — a factorized normal distribution N ( µ ( X ) , σ ( X )) .
Revised: We use a standard posterior approximation — a factorized normal distribution N ( µ ( X ) , σ ( X )) .

Original: To ensure the functions µ and σ are permutation invariant, we use deep sets (Zaheer et al., 2017), or set transformers (Lee et al., 2019).
Revised: To ensure the functions µ and σ are permutation invariant, we use deep sets (Zaheer et al., 2017), or set transformers (Lee et al., 2019).

Original: The last thing we still need to calculate is log p ( X | z ) .
Revised: The last thing we still need to calculate is log p ( X | z ) .

Original: Each point x i ∈
Revised: Each point x i ∈

Original: X is conditionally independent given z .
Revised: X is conditionally independent given z .

Original: In practice, this means we have an inhomogeneous
Revised: In practice, this means we have an inhomogeneous

Original: Poisson process, but conditioned on a latent variable, i.e. a normalizing ﬂow whose parameters are a function of z .
Revised: Poisson process, but conditioned on a latent variable, i.e. a normalizing ﬂow whose parameters are a function of z .

Original: To get a new realization we ﬁrst sample z , then sample x i from a normalizing ﬂow conditioned on z .
Revised: To get a new realization we ﬁrst sample z , then sample x i from a normalizing ﬂow conditioned on z .

Original: To make a fair comparison with other methods we implement an importance weighted autoencoder with a tighter likelihood lower bound (Burda et al., 2015).
Revised: To make a fair comparison with other methods we implement an importance weighted autoencoder with a tighter likelihood lower bound (Burda et al., 2015).

Original:  C ONFET M ODEL
Revised:  C ONFET M ODEL

Original: The models from the previous section trade-off between ﬂexibility, having no special assumptions on the data structure, and modeling exact likelihood.
Revised: The models from the previous section trade-off between ﬂexibility, having no special assumptions on the data structure, and modeling exact likelihood.

Original: Here, we introduce our new model, a continuous normalizing ﬂow with equivariant transformations (C ONFET ), that satisﬁes all of the prerequisites.
Revised: Here, we introduce our new model, a continuous normalizing ﬂow with equivariant transformations (C ONFET ), that satisﬁes all of the prerequisites.

Original: If we want to have an expressive generative model that can both sample and estimate likelihood, we want to use a normalizing ﬂow.
Revised: If we want to have an expressive generative model that can both sample and estimate likelihood, we want to use a normalizing ﬂow.

Original: However, having interactions between points leads in most cases to intractable inverses and determinants.
Revised: However, having interactions between points leads in most cases to intractable inverses and determinants.

Original: This is not the case when using continuous normalizing ﬂows.
Revised: This is not the case when using continuous normalizing ﬂows.

Original: A CNF transforms an initial sample z ∼ q ( z ) to a sample in the target distribution p ( x ) with a differential equation f ( z ( t ) , t ) =
Revised: A CNF transforms an initial sample z ∼ q ( z ) to a sample in the target distribution p ( x ) with a differential equation f ( z ( t ) , t ) =

Original: ∂
Revised: ∂

Original: z ( t ) /∂t .
Revised: z ( t ) /∂t .

Original: The change in log-density is (Chen et al., 2018):
Revised: The change in log-density is (Chen et al., 2018):

Original: [Equation]
Revised: [Equation]

Original: (6) and the ﬁnal log-density is obtained by integrating across time, e.g., using a black-box ODE solver:
Revised: (6) and the ﬁnal log-density is obtained by integrating across time, e.g., using a black-box ODE solver:

Original: [Equation]
Revised: [Equation]

Original: The inverse is obtained by integrating in the opposite direction, t 1 to t 0 .
Revised: The inverse is obtained by integrating in the opposite direction, t 1 to t 0 .

Original: The only constraint on f , besides permutation equivariance, is that it needs to be, along with its derivatives, Lipschitz continuous which is easily satisﬁed by using smooth activations in a neural network.
Revised: The only constraint on f , besides permutation equivariance, is that it needs to be, along with its derivatives, Lipschitz continuous which is easily satisﬁed by using smooth activations in a neural network.

Original: The beneﬁt of this approach is that we are restricting f mildly, meaning the dimensions can interact with each other arbitrarily.
Revised: The beneﬁt of this approach is that we are restricting f mildly, meaning the dimensions can interact with each other arbitrarily.

Original: We utilize this to model interactions between points in a point process.
Revised: We utilize this to model interactions between points in a point process.

Original: Equivariant transformation.
Revised: Equivariant transformation.

Original: A general way to deﬁne an equivariant map f : B n → B n is to deﬁne an invariant function g : B n → R , and let f ( X ) =
Revised: A general way to deﬁne an equivariant map f : B n → B n is to deﬁne an invariant function g : B n → R , and let f ( X ) =

Original: ∇ X g ( X ) (Papamakarios et al., 2019).
Revised: ∇ X g ( X ) (Papamakarios et al., 2019).

Original: However, this is computationally costly and numerically unstable (Köhler et al., 2020).
Revised: However, this is computationally costly and numerically unstable (Köhler et al., 2020).

Original: We can instead use an explicitly deﬁned equivariant layer (Zaheer et al., 2017; Maron et al., 2020):
Revised: We can instead use an explicitly deﬁned equivariant layer (Zaheer et al., 2017; Maron et al., 2020):

Original: [Equation]
Revised: [Equation]

Original: where g, h : B → B are linear functions.
Revised: where g, h : B → B are linear functions.

Original: We stack multiple layers with activations in between.
Revised: We stack multiple layers with activations in between.

Original: An alternative equivariant mapping is deﬁned with a (multihead) self-attention layer (Vaswani et al., 2017).
Revised: An alternative equivariant mapping is deﬁned with a (multihead) self-attention layer (Vaswani et al., 2017).

Original: Lee et al.
Revised: Lee et al.

Original: (2019) propose an extension for large sets that uses a small number of learnable inducing points m (cid:28) n which reduces the computational cost.
Revised: (2019) propose an extension for large sets that uses a small number of learnable inducing points m (cid:28) n which reduces the computational cost.

Original: The detailed discussion is in B.2. x
Revised: The detailed discussion is in B.2. x

Original: [Figure]
Revised: [Figure]

Original: (cid:40) x(cid:40)
Revised: (cid:40) x(cid:40)

Original: Dimensionwise
Revised: Dimensionwise

Original: Within point
Revised: Within point

Original: Between points
Revised: Between points

Original: Full
Revised: Full

Original: [Figure]
Revised: [Figure]

Original: Computing the trace in Eq.
Revised: Computing the trace in Eq.

Original: 6 during training is expensive since it scales quadratically with the number of points O ( n 2 d 2 ) .
Revised: 6 during training is expensive since it scales quadratically with the number of points O ( n 2 d 2 ) .

Original: We propose three different solutions to this: 1) using a trace estimator, 2) having a closed-form expression, and 3) decoupling the network to cheaply compute the trace.
Revised: We propose three different solutions to this: 1) using a trace estimator, 2) having a closed-form expression, and 3) decoupling the network to cheaply compute the trace.

Original: In the following, we make a distinction between the models based on which of these methods they use.
Revised: In the following, we make a distinction between the models based on which of these methods they use.

Original: C ONFET -S TOCHASTIC .
Revised: C ONFET -S TOCHASTIC .

Original: We use Hutchinson’s trace estimator (Hutchinson, 1989) during training as proposed by Grathwohl et al.
Revised: We use Hutchinson’s trace estimator (Hutchinson, 1989) during training as proposed by Grathwohl et al.

Original: This reduces the cost to O ( nd ) .
Revised: This reduces the cost to O ( nd ) .

Original: However, it introduces additional variance into training making it potentially unstable.
Revised: However, it introduces additional variance into training making it potentially unstable.

Original: C ONFET -F IXED .
Revised: C ONFET -F IXED .

Original: Another approach is having a closed-form result for the trace.
Revised: Another approach is having a closed-form result for the trace.

Original: Notice how Eq.decouples the mapping into a pointwise transformation and an inﬂuence from other points.
Revised: Notice how Eq.decouples the mapping into a pointwise transformation and an inﬂuence from other points.

Original: Then a simple function with a trace equal to 0 is f
Revised: Then a simple function with a trace equal to 0 is f

Original: ( X ) i
Revised: ( X ) i

Original: =
Revised: =

Original: (cid:80) j (cid:54) =
Revised: (cid:80) j (cid:54) =

Original: i h ( x j ) .
Revised: i h ( x j ) .

Original: This connects to the early models in traditional normalizing ﬂows (Dinh et al., 2015).
Revised: This connects to the early models in traditional normalizing ﬂows (Dinh et al., 2015).

Original: In this case, the CNF is a single ﬂow layer and we add additional elementwise coupling layers to model the inhomogeneous part of the process.
Revised: In this case, the CNF is a single ﬂow layer and we add additional elementwise coupling layers to model the inhomogeneous part of the process.

Original: C ONFET -E XACT .
Revised: C ONFET -E XACT .

Original: We extend the decoupling idea further to model arbitrarily complex architectures with cheap exact trace calculation.
Revised: We extend the decoupling idea further to model arbitrarily complex architectures with cheap exact trace calculation.

Original: Chen & Duvenaud (2019) show how a neural network can be viewed as a combination of per dimension mappings x
Revised: Chen & Duvenaud (2019) show how a neural network can be viewed as a combination of per dimension mappings x

Original: i (cid:55)→ y i and interactions between the dimensions x j (cid:55)→ y
Revised: i (cid:55)→ y i and interactions between the dimensions x j (cid:55)→ y

Original: i .
Revised: i .

Original: They separate them explicitly into a transformer τ and conditioner c , e.g. y i
Revised: They separate them explicitly into a transformer τ and conditioner c , e.g. y i

Original: =
Revised: =

Original: τ ( x i , c ( x j )) , without losing expressiveness.
Revised: τ ( x i , c ( x j )) , without losing expressiveness.

Original: This factors the Jacobian into a diagonal ( ∂τ i /∂x i ) and hollow matrix ( ∂τ i /∂x j ), meaning we only need to calculate the diagonal to obtain the trace.
Revised: This factors the Jacobian into a diagonal ( ∂τ i /∂x i ) and hollow matrix ( ∂τ i /∂x j ), meaning we only need to calculate the diagonal to obtain the trace.

Original: This can be done cheaply with modern deep learning frameworks.
Revised: This can be done cheaply with modern deep learning frameworks.

Original: For example, the interaction between the points in Eq.
Revised: For example, the interaction between the points in Eq.

Original: 8 detaches itself naturally from the computation graph and we decouple the pointwise transformation g with a masked neural network (Germain et al., 2015).
Revised: 8 detaches itself naturally from the computation graph and we decouple the pointwise transformation g with a masked neural network (Germain et al., 2015).

Original: We take this idea further by exploiting it hierarchically (Figure 2): after the dimensionwise mapping, we ﬁrst impose interactions within a point, subsequently interactions between points.
Revised: We take this idea further by exploiting it hierarchically (Figure 2): after the dimensionwise mapping, we ﬁrst impose interactions within a point, subsequently interactions between points.

Original: To summarize, one dimension of one point x i ∈ X is transformed conditioned on all the other points in the set and all the other dimensions in the point.
Revised: To summarize, one dimension of one point x i ∈ X is transformed conditioned on all the other points in the set and all the other dimensions in the point.

Original: This decoupling allows us to obtain the exact trace in O ( nd ) .
Revised: This decoupling allows us to obtain the exact trace in O ( nd ) .

Original: Implementation details for Eq.
Revised: Implementation details for Eq.

Original: 8 and attention network are in  B.4.
Revised: 8 and attention network are in Appendix B.4.

Original: Training.
Revised: Training.

Original: We can either use a ﬁxed-step or adaptive ODE solver for Eq.
Revised: We can either use a ﬁxed-step or adaptive ODE solver for Eq.

Original: The trade-off is between having a bounded number of solver steps and a bounded error.
Revised: The trade-off is between having a bounded number of solver steps and a bounded error.

Original: To regularize ODE dynamics we can apply weight decay on the parameters of f or penalize the Jacobian Frobenius norm (Finlay et al., 2020).
Revised: To regularize ODE dynamics we can apply weight decay on the parameters of f or penalize the Jacobian Frobenius norm (Finlay et al., 2020).

Original: With this, we can often resort to ﬁxed-step solvers that offer signiﬁcant speed improvements.
Revised: With this, we can often resort to ﬁxed-step solvers that offer signiﬁcant speed improvements.

Original: Instead of backpropagating through solver steps, we use a memory efﬁcient adjoint method to obtain the gradients (Farrell et al., 2013; Chen et al., 2018).
Revised: Instead of backpropagating through solver steps, we use a memory efﬁcient adjoint method to obtain the gradients (Farrell et al., 2013; Chen et al., 2018).

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Point processes.
Revised: Point processes.

Original: Cox process (Cox, 1955) is a doubly stochastic Poisson process that deﬁnes intensity as a realization of a random ﬁeld, e.g. a Gaussian process.
Revised: Cox process (Cox, 1955) is a doubly stochastic Poisson process that deﬁnes intensity as a realization of a random ﬁeld, e.g. a Gaussian process.

Original: We can think of this as a two step procedure where we ﬁrst sample the intensity function and then sample points based on this intensity.
Revised: We can think of this as a two step procedure where we ﬁrst sample the intensity function and then sample points based on this intensity.

Original: A Neyman-Scott process (Neyman & Scott, 1958) generates parents from a Poisson process, and children around parents.
Revised: A Neyman-Scott process (Neyman & Scott, 1958) generates parents from a Poisson process, and children around parents.

Original: A more speciﬁc example is a Matérn cluster process (Matérn, 2013) that places children in balls centered around parents.
Revised: A more speciﬁc example is a Matérn cluster process (Matérn, 2013) that places children in balls centered around parents.

Original: A variation where children follow a normal distribution is called a modiﬁed Thomas process (Thomas, 1949).
Revised: A variation where children follow a normal distribution is called a modiﬁed Thomas process (Thomas, 1949).

Original: Explicit interactions can be deﬁned with a Gibbs process by assigning energy values to point conﬁgurations, stemming from statistical physics (Ruelle, 1969).
Revised: Explicit interactions can be deﬁned with a Gibbs process by assigning energy values to point conﬁgurations, stemming from statistical physics (Ruelle, 1969).

Original: However, this comes at a cost of calculating the normalization constant to ﬁnd densities.
Revised: However, this comes at a cost of calculating the normalization constant to ﬁnd densities.

Original: Besides using pseudolikelihood (Besag, 1975) to ﬁt them, other works used logistic composite likelihood (Clyde & Strauss, 1991) and Monte Carlo methods (Huang & Ogata, 1999).
Revised: Besides using pseudolikelihood (Besag, 1975) to ﬁt them, other works used logistic composite likelihood (Clyde & Strauss, 1991) and Monte Carlo methods (Huang & Ogata, 1999).

Original: Even though we do not specify the random ﬁeld explicitly, the randomness from the base density is allowing us to model these processes with equivariant transformations.
Revised: Even though we do not specify the random ﬁeld explicitly, the randomness from the base density is allowing us to model these processes with equivariant transformations.

Original: We avoid calculating the normalization constant by using CNFs that give us likelihood directly at no ﬂexibility trade-off.
Revised: We avoid calculating the normalization constant by using CNFs that give us likelihood directly at no ﬂexibility trade-off.

Original: Continuous normalizing ﬂows.
Revised: Continuous normalizing ﬂows.

Original: Neural ODEs were proposed recently (Lu et al., 2018; Chen et al., 2018) as a continuous equivalent of the widely adopted residual architecture (He et al., 2016).
Revised: Neural ODEs were proposed recently (Lu et al., 2018; Chen et al., 2018) as a continuous equivalent of the widely adopted residual architecture (He et al., 2016).

Original: By taking advantage of its properties we can construct a continuous version of a normalizing ﬂow.
Revised: By taking advantage of its properties we can construct a continuous version of a normalizing ﬂow.

Original: Grathwohl et al.
Revised: Grathwohl et al.

Original: (2019) demonstrate the beneﬁts of this approach by using a stochastic trace estimator that allows them to scale to bigger datasets.
Revised: (2019) demonstrate the beneﬁts of this approach by using a stochastic trace estimator that allows them to scale to bigger datasets.

Original: Having free-form Jacobian proved to be useful especially in our use case where we require it from equivariant transformations (beyond invariance through independence).
Revised: Having free-form Jacobian proved to be useful especially in our use case where we require it from equivariant transformations (beyond invariance through independence).

Original: Many recent works focused on improving the neural ODEs, including better representation power (Dupont et al., 2019), improved training (Gholami et al., 2019; Zhuang et al., 2020), regularization of the solver dynamics (Finlay et al., 2020; Kelly et al., 2020) etc.
Revised: Many recent works focused on improving the neural ODEs, including better representation power (Dupont et al., 2019), improved training (Gholami et al., 2019; Zhuang et al., 2020), regularization of the solver dynamics (Finlay et al., 2020; Kelly et al., 2020) etc.

Original: The idea of using CNFs together with equivariant functions has been proposed before.
Revised: The idea of using CNFs together with equivariant functions has been proposed before.

Original: Köhler et al.
Revised: Köhler et al.

Original: (2020) use it in the context of multi-body physical systems by modeling the Gibbs distribution with a simple equivariant function, calculating the forces between points proportional to their distance.
Revised: (2020) use it in the context of multi-body physical systems by modeling the Gibbs distribution with a simple equivariant function, calculating the forces between points proportional to their distance.

Original: In particular, they use a transformation based on a Gaussian kernel which allows them to have a closed-form trace solution.
Revised: In particular, they use a transformation based on a Gaussian kernel which allows them to have a closed-form trace solution.

Original: In Section 4 we offer ways to use more expressive equivariant functions with closed-form trace that do not rely on a predeﬁned kernel.
Revised: In Section 4 we offer ways to use more expressive equivariant functions with closed-form trace that do not rely on a predeﬁned kernel.

Original: We also show how this can be applied for modeling point processes with exact likelihood — an important, previously unsolved problem.
Revised: We also show how this can be applied for modeling point processes with exact likelihood — an important, previously unsolved problem.

Original: A concurrent preprint by Li et al.
Revised: A concurrent work by Li et al.

Original: (2020) proposes modeling sets similar to our S TOCHASTIC model.
Revised: (2020) proposes modeling sets similar to our S TOCHASTIC model.

Original: Our work takes a general approach with further variations, and is more efﬁcient and scalable.
Revised: Our work takes a general approach with further variations, and is more efﬁcient and scalable.

Original: Toth et al.
Revised: Toth et al.

Original: (2020) propose modeling physics systems by transforming theinitial samples with a
Revised: (2020) propose modeling physics systems by transforming the

Original: Hamiltonian dynamics.

Original: The method is similar to CNFs in general, but it preserves volume, like our

Original: C ONFET -F IXED model.

Original: Training requires variational approximation.

Original: Modeling sets.
Revised: Modeling sets.

Original: Zaheer et al.
Revised: Zaheer et al.

Original: ; Qi et al.
Revised: ; Qi et al.

Original: (2017) achieve permutation equivariance similar to Eq.
Revised: (2017) achieve permutation equivariance similar to Eq.

Original: They both show the universal approximation property of such a network.
Revised: They both show the universal approximation property of such a network.

Original: Same property holds for the attention based network (Lee et al., 2019).
Revised: Same property holds for the attention based network (Lee et al., 2019).

Original: Wagstaff et al.
Revised: Wagstaff et al.

Original: (2019) study the theoretical limitations of invariant functions with sum aggregation.
Revised: (2019) study the theoretical limitations of invariant functions with sum aggregation.

Original: Other aggregation schemes include max-pooling,
Revised: Other aggregation schemes include max-pooling,

Original: Janossy pooling (Murphy et al., 2019), featurewise sort pooling (Zhang et al., 2020), constructing the adjacency matrix and using graph methods (Grover & Leskovec, 2016) etc.
Revised: Janossy pooling (Murphy et al., 2019), featurewise sort pooling (Zhang et al., 2020), constructing the adjacency matrix and using graph methods (Grover & Leskovec, 2016) etc.

Original: We can use any of these methods in our model as well.
Revised: We can use any of these methods in our model as well.

Original: Korshunova et al.
Revised: Korshunova et al.

Original: (2018) present a model that transforms every point independently and has exchangeable Student-t distribution for base density.
Revised: (2018) present a model that transforms every point independently and has exchangeable Student-t distribution for base density.

Original: Bender et al.
Revised: Bender et al.

Original: (2020) extend this by having an autoregressive transformation and train on ordered sequences using Eq. 1.
Revised: (2020) extend this by having an autoregressive transformation and train on ordered sequences using Eq. 1.

Original: Our autoregressive model is very similar to this formulation.
Revised: Our autoregressive model is very similar to this formulation.

Original: Other autoregressive models for set-like data include graph generation on breath-ﬁrst search ordering (You et al., 2018); and ordering on z-axis for meshes (Nash et al., 2020) and point clouds (Sun et al., 2020).
Revised: Other autoregressive models for set-like data include graph generation on breath-ﬁrst search ordering (You et al., 2018); and ordering on z-axis for meshes (Nash et al., 2020) and point clouds (Sun et al., 2020).

Original: Neural models for temporal point processes know the explicit order and usually use RNNs to deﬁne density (Du et al., 2016; Mei & Eisner, 2017).
Revised: Neural models for temporal point processes know the explicit order and usually use RNNs to deﬁne density (Du et al., 2016; Mei & Eisner, 2017).

Original: In contrast to autoregressive methods, Yang et al.
Revised: In contrast to autoregressive methods, Yang et al.

Original: (2019) model point clouds in a similar way to our variational autoencoder, with a CNF for posterior and a shared CNF for points.
Revised: (2019) model point clouds in a similar way to our variational autoencoder, with a CNF for posterior and a shared CNF for points.

Original: Yuan et al.
Revised: Yuan et al.

Original: (2020) model spatial point processes with a variational autoencoder.
Revised: (2020) model spatial point processes with a variational autoencoder.

Original: Instead of intensity, they also use density to avoid calculating the integral in Eq. 2 and deﬁne it with a nonparametric kernel.
Revised: Instead of intensity, they also use density to avoid calculating the integral in Eq. 2 and deﬁne it with a nonparametric kernel.

Original: We generalize this by using normalizing ﬂows as expressive density models.
Revised: We generalize this by using normalizing ﬂows as expressive density models.

Original: Finally, they demonstrate an application of spatial point processes to recommender systems by embedding the non-spatial data with a graph neural network, which is something our model can support as well.
Revised: Finally, they demonstrate an application of spatial point processes to recommender systems by embedding the non-spatial data with a graph neural network, which is something our model can support as well.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: In the following section we show how our model compares to competitors presented in Sectionon various synthetic datasets that are both inhomogeneous and with clustering behavior.
Revised: In the following section we show how our model compares to competitors presented in Sectionon various synthetic datasets that are both inhomogeneous and with clustering behavior.

Original: We also use real-world datasets that provide realistic location modeling scenarios and are commonly used in the literature.
Revised: We also use real-world datasets that provide realistic location modeling scenarios and are commonly used in the literature.

Original: We also demonstrate the sampling capabilities of our models and compare their scalability and efﬁciency directly.
Revised: We also demonstrate the sampling capabilities of our models and compare their scalability and efﬁciency directly.

Original: [Figure]
Revised: Groundtruth

Original: Samples
Revised: Samples

Original: Thomas
Revised: Thomas

Original: Checkins
Revised: Checkins

Original: [Figure]
Revised: [Figure]

Original: NY

Original: [Figure]
Revised: [Figure]

Original: Autoregressive
Revised: Autoregressive

Original: [Figure]
Revised: [Figure]

Original: C
Revised: C

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: -E
Revised: -E

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: TOCHASTIC
Revised: TOCHASTIC

Original: Check
Revised: Check

Original: [Table]
Revised: [Table]

Original: NY -
Revised: NY -

Original: [Table]
Revised: [Table]

Original: Crimes
Revised: Crimes

Original: [Table]
Revised: [Table]

Original: Matérn
Revised: Matérn

Original: [Table]
Revised: [Table]

Original: Mixture
Revised: Mixture

Original: [Table]
Revised: [Table]

Original: Thomas
Revised: Thomas

Original: [Table]
Revised: [Table]

Original: Uniform
Revised: Uniform

Original: [Table]
Revised: [Table]

Original: 6. D ATASETS
Revised: 6. D ATASETS

Original: Synthetic data.
Revised: Synthetic data.

Original: For all datasets we simulate 1000 realizations on (0 , 1) × (0 , 1) area.
Revised: For all datasets we simulate 1000 realizations on (0 , 1) × (0 , 1) area.

Original: Thomas dataset is simulated by ﬁrst sampling m ∼ Pois(3) parents uniformly, and then for each of them we sample n i ∼ Pois (5) points from normal distribution centered around the parents, with diagonal covariance with value 0 .
Revised: Thomas dataset is simulated by ﬁrst sampling m ∼ Pois(3) parents uniformly, and then for each of them we sample n i ∼ Pois (5) points from normal distribution centered around the parents, with diagonal covariance with value 0 .

Original: Matérn generates the parents the same way, but children are uniformly sampled on a circle with R = 0 .
Revised: Matérn generates the parents the same way, but children are uniformly sampled on a circle with R = 0 .

Original: Mixture is an inhomogeneous process that generates n ∼ Pois (64) points from a mixture of 3 normal distributions with means (0 . 3 , 0 . 3) , (0 . 5 , 0 . 7) , (0 . 7 , 0 . 3) and diagonal covariance 0 .
Revised: Mixture is an inhomogeneous process that generates n ∼ Pois (64) points from a mixture of 3 normal distributions with means (0 . 3 , 0 . 3) , (0 . 5 , 0 . 7) , (0 . 7 , 0 . 3) and diagonal covariance 0 .

Original: Uniform is a mixture of two uniform processes, as described in Section 1.
Revised: Uniform is a mixture of two uniform processes, as described in Section 1.

Original: Real-world data.
Revised: Real-world data.

Original: Check-ins NY 1 is a dataset of locations collected from social network users.
Revised: Check-ins NY 1 is a dataset of locations collected from social network users.

Original: We took all the data points from an area in New York and users that have more than 10 and less than 100 locations giving us 1938 unique users.
Revised: We took all the data points from an area in New York and users that have more than 10 and less than 100 locations giving us 1938 unique users.

Original: We consider a set of all locations for a user as a single realization.
Revised: We consider a set of all locations for a user as a single realization.

Original: We also construct a smaller dataset for a different city (Check-ins Paris) that hasunique users.
Revised: We also construct a smaller dataset for a different city (Check-ins Paris) that hasunique users.

Original: Crimes 2 dataset contains daily records of locations and types of crimes that occurred in Portland.
Revised: Crimes 2 dataset contains daily records of locations and types of crimes that occurred in Portland.

Original: Each day contains between 298 and 736 points with 480 on average.
Revised: Each day contains between 298 and 736 points with 480 on average.

Original: 6. M ODELS
Revised: 6. M ODELS

Original: Inhomogeneous Poisson process (IHP ) is implemented as a density estimator p ( X ) =
Revised: Inhomogeneous Poisson process (IHP ) is implemented as a density estimator p ( X ) =

Original: (cid:81) i p ( x i ) since it assumes independence between the points.
Revised: (cid:81) i p ( x i ) since it assumes independence between the points.

Original: We use a normalizing ﬂow with spline coupling layers.
Revised: We use a normalizing ﬂow with spline coupling layers.

Original: Autoregressive model inputs points sorted on the ﬁrst dimension, with Eq.
Revised: Autoregressive model inputs points sorted on the ﬁrst dimension, with Eq.

Original: 1 used for training.
Revised: 1 used for training.

Original: We stack autoregressive and set coupling layers (Bender et al., 2020).
Revised: We stack autoregressive and set coupling layers (Bender et al., 2020).

Original: Importance weighted autoencoder (IWAE), as described in Section 3 uses either deep sets or attention layers as an invariant function that outputs posterior parameters.
Revised: Importance weighted autoencoder (IWAE), as described in Section 3 uses either deep sets or attention layers as an invariant function that outputs posterior parameters.

Original: In both we use max-pooling for aggregation.
Revised: In both we use max-pooling for aggregation.

Original: Between attention layers we apply layer normalization (Ba et al., 2016).
Revised: Between attention layers we apply layer normalization (Ba et al., 2016).

Original: Term p ( X | z ) is an IHP conditioned on z .
Revised: Term p ( X | z ) is an IHP conditioned on z .

Original: We use the same hyperparameter ranges.
Revised: We use the same hyperparameter ranges.

Original: During training we use 5 samples, to estimate the test likelihood (Eq. 3) we use 5000 samples (Burda et al., 2015).
Revised: During training we use 5 samples, to estimate the test likelihood (Eq. 3) we use 5000 samples (Burda et al., 2015).

Original: Models from Sectionare implemented as a CNF with equivariant transformations that are either deep set or attention layers.
Revised: Models from Sectionare implemented as a CNF with equivariant transformations that are either deep set or attention layers.

Original: We usually use dopri5 solver (Dormand & Prince, 1980).
Revised: We usually use dopri5 solver (Dormand & Prince, 1980).

Original: C
Revised: C

Original: ONFET -F IXED uses Runge-Kutta solver with 10 steps.
Revised: ONFET -F IXED uses Runge-Kutta solver with 10 steps.

Original: Further details on hyperparameter search and ablation studies are in  D.
Revised: Further details on hyperparameters and ablation studies are in Appendix D.

Original: Datasets are split into training, validation and test sets (60%-20%-20%).
Revised: Datasets are split into training, validation and test sets (60%-20%-20%).

Original: We train with early stopping after validation loss does not improve for 20 consecutive epochs.
Revised: We train with early stopping after validation loss does not improve for 20 consecutive epochs.

Original: We use mini-batches of size 64 and Adam optimizer with learning rate of 10 − 3 (Kingma & Ba, 2015).
Revised: We use mini-batches of size 64 and Adam optimizer with learning rate of 10 − 3 (Kingma & Ba, 2015).

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: 6. R ESULTS
Revised: 6. R ESULTS

Original: The loss we use is per-point negative log-likelihood, i [Figure] e. L ( X ) 
Revised: The loss we use is per-point negative log-likelihood, i.e. L ( X ) =

Original: = − log p ( X ) /n .
Revised: − log p ( X ) /n .

Original: We report the mean and variance of the loss on a held-out test set averaged over 5 runs for models selected based on their validation set performance.
Revised: We report the mean and variance of the loss on a held-out test set averaged over 5 runs for models selected based on their validation set performance.

Original: The results are in Table 1.
Revised: The results are in Table 1.

Original: Our models achieve the best scores on most of the datasets.
Revised: Our models achieve the best scores on most of the datasets.

Original: It is worth noting that C ONFET -F
Revised: It is worth noting that C ONFET -F

Original: IXED is often on par with others even though it has the simplest architecture.
Revised: IXED is often on par with others even though it has the simplest architecture.

Original: Figure 4 shows how we can use our model trained on Checkins-NY dataset to detect the out-of-distribution samples, in this case from Checkins-Paris.
Revised: Figure 4 shows how we can use our model trained on Checkins-NY dataset to detect the out-of-distribution samples, in this case from Checkins-Paris.

Revised: We also note our model outperforms (Köhler et al., 2020) by a large margin (see Appendix C.2).

Original: Sampling quality.
Revised: Sampling quality.

Original: We provide qualitative samples for C ONFET -E XACT model on Thomas and
Revised: We provide qualitative samples for C ONFET -E XACT model on Thomas and

Original: Checkins-NY dataset in Figure 3.
Revised: Checkins-NY dataset in Figure 3.

Original: We show it is able to learn the clustering process, and in the second example additionally learns the underlying city topology that restricts the samples, e.g. to appear only on the peninsula and not in the water.
Revised: We show it is able to learn the clustering process, and in the second example additionally learns the underlying city topology that restricts the samples, e.g. to appear only on the peninsula and not in the water.

Original: Additional samples for all the models are in
Revised: Additional samples for all the models are in

Original: C.2.
Revised: Appendix C.3.

Original: We also assess the quantitative quality of the samples on real-world datasets by comparing the distributions of in-between point distances in real data and model samples.
Revised: We also assess the quantitative quality of the samples on real-world datasets by comparing the distributions of in-between point distances in real data and model samples.

Original: The results in Tablesuggest our models produce samples of higher quality.
Revised: The results in Table 2 suggest our models produce samples of higher quality.

Original: C ONFET model comparison.
Revised: C ONFET model comparison.

Original: In Section 4 we provided three different C ONFET models based on how they calculate the trace.
Revised: In Section 4 we provided three different C ONFET models based on how they calculate the trace.

Original: We have seen that C ONFET -E XACT and C ONFET -S
Revised: We have seen that C ONFET -E XACT and C ONFET -S

Original: TOCHASTIC perform the best.
Revised: TOCHASTIC perform the best.

Original: Now we compare them directly.
Revised: Now we compare them directly.

Original: Figure 5 contains the results.
Revised: Figure 5 contains the results.

Original: First, we test the speed of trace calculation by timing it with different set sizes.
Revised: First, we test the speed of trace calculation by timing it with different set sizes.

Original: We use an identical architecture for both models, further details are in C.3.
Revised: We use an identical architecture for both models, further details are in Appendix C.4.

Original: We can see that C ONFET -E XACT is faster although both have the same theoretical cost O ( nd ) .
Revised: We can see that C ONFET -E XACT is faster although both have the same theoretical cost O ( nd ) .

Original: The naive exact trace calculation proves to be unusable in practice.
Revised: The naive exact trace calculation proves to be unusable in practice.

Original: Second, we test the training stability on Mixture dataset and notice that both converge to the same value but C ONFET -S TOCHASTIC has very noisy loss curve.
Revised: Second, we test the training stability on Mixture dataset and notice that both converge to the same value but C ONFET -S TOCHASTIC has very noisy loss curve.

Original: Finally, we plot the number of function evaluations of the dopri5 solver trained on Matérn dataset during training.
Revised: Finally, we plot the number of function evaluations of the dopri5 solver trained on Matérn dataset during training.

Original: Here again,
Revised: Here again,

Original: C ONFET -E XACT achieves the better result.
Revised: C ONFET -E XACT achieves the better result.

Original: We conclude that C ONFET can be used in most cases to model sets and point processes.
Revised: We conclude that C ONFET can be used in most cases to model sets and point processes.

Original: If we want better scalability while retaining the same performance we should use C ONFET -E XACT .
Revised: If we want better scalability while retaining the same performance we should use C ONFET -E XACT .

Original: In case we model simple interactions or inhomogeneous data C
Revised: In case we model simple interactions or inhomogeneous data C

Original: ONFET -F
Revised: ONFET -F

Original: IXED will work well.
Revised: IXED will work well.

Revised: [Figure]

Original:  C ONCLUSION AND FUTURE WORK
Revised: [Figure] 6. S TOCHASTIC INTENSITY DIRECTLY FROM C ONFET MODEL Since the density p ( X ) in C ONFET model is inﬂuenced by the interactions between the points, it will not be stationary but will depend on the locations in X (unlike e.g., IHP model). Therefore, we can draw one density realization on the region B by sampling X ∼ p

Revised: ( X ) .

Revised: The density can also be easily evaluated at non-observed points x / ∈ X .

Revised: Figure 6 shows different density realizations and the corresponding samples from a model trained on Check-ins NY data.

Revised: This connects to the doubly stochastic processes, in particular Cox process (Cox, 1955).

Revised: A popular way to deﬁne a Cox process is by deﬁning the random ﬁeld with a Gaussian process (Rasmussen & Williams, 2006), and drawing an intensity function as a random realization.

Revised: To ensure the intensity is always non-negative, the random function f ∼ GP can be squared (Lloyd et al., 2015) or we can apply a sigmoid function (Adams et al., 2009).

Revised: Figure 7 shows how a Gaussian process modulated poisson process (GP-PP) (Lloyd et al., 2015) ﬁts the coal mining disaster dataset with 191 points collected from 1851 to 1962.

Revised: We additionally plot the empirical intensity measure estimated with a Gaussian kernel and notice that GP-PP overestimates the point counts on some parts of the data.

Revised: Finally, we plot 10 intensity function realizations from our model.

Revised: C ONFET closely matches the underlying intensity and shows uncertainty in parts where it observed less data.

Revised: 6. M ODELING POINT CLOUDS AND DIGITS

Revised: To further test the capabilities of our model we use two traditional set modeling problems: point clouds and discretized images.

Revised: We use airplane samples from a point cloud dataset (Wu et al., 2015), and uniformly sample 512 spatial points to reduce the input size.

Revised: We stack 8 continuous normalizing ﬂow transformations and train using maximum likelihood.

Revised: For digits dataset (LeCun et al., 1998) we randomly sample 50 white pixels in each image and record their 2-D positions.

Revised: In

Revised: Figure 8 we see samples from our model compared to the real data.

Revised: We conclude our model is able to capture the shapes which shows that it can be used for complex set data.

Revised:  C ONCLUSION AND FUTURE WORK

Original: In this work we presented a novel model for point processes and random sets that uses normalizing ﬂows with equivariant functions.
Revised: In this work we presented a novel model for point processes and random sets that uses normalizing ﬂows with equivariant functions.

Original: Throughout the experiments we showed that it is able to model complex processes with interactions between the points while having tractable likelihood and straightforward sampling.
Revised: Throughout the experiments we showed that it is able to model complex processes with interactions between the points while having tractable likelihood and straightforward sampling.

Original: The performance of our model is better than all the competitors.
Revised: The performance of our model is better than all the competitors.

Original: We additionally provide different versions that are more scalable while retaining the same properties of the base model.
Revised: We additionally provide different versions that are more scalable while retaining the same properties of the base model.

Original: Our model relies in particular on two types of deep learning models: (continuous) normalizing ﬂows and equivariant neural networks.
Revised: Our model relies in particular on two types of deep learning models: (continuous) normalizing ﬂows and equivariant neural networks.

Original: By utilizing new techniques developed in these two subﬁelds we hope our method will perform even better in the future.
Revised: By utilizing new techniques developed in these two subﬁelds we hope our method will perform even better in the future.

Original: A T HEORETICAL BACKGROUND A. C ARDINALITY IN LIKELIHOOD
Revised: A T HEORETICAL BACKGROUND A. C ARDINALITY IN LIKELIHOOD

Original: In this section we discuss how to deﬁne the likelihood for point processes.
Revised: In this section we discuss how to deﬁne the likelihood for point processes.

Original: Let’s take an inhomogeneous Poisson process as an example.
Revised: Let’s take an inhomogeneous Poisson process as an example.

Original: It deﬁnes the density on B and the probability of observing a set X is p ( X )
Revised: It deﬁnes the density on B and the probability of observing a set X is p ( X )

Original: =
Revised: =

Original: p ( x 1 , . . .
Revised: p ( x 1 , . . .

Original: , x n ) p ( n ) .
Revised: , x n ) p ( n ) .

Original: In our experiments we considered only the ﬁrst part, ignoring the cardinality of the set.
Revised: In our experiments we considered only the ﬁrst part, ignoring the cardinality of the set.

Original: This leads to some unexpected behavior.
Revised: This leads to some unexpected behavior.

Original: For example, we will almost always get higher probability for sets with a smaller cardinality because, e.g., p > p 2 , when p < 1 , and vice versa.
Revised: For example, we will almost always get higher probability for sets with a smaller cardinality because, e.g., p > p 2 , when p < 1 , and vice versa.

Original: This is also pointed out by Vo et al.
Revised: This is also pointed out by Vo et al.

Original: They propose the following likelihood:
Revised: They propose the following likelihood:

Original: [Equation]
Revised: [Equation]

Original: where U is the unit of hyper-volume to ensure consistent results over different measurement units.
Revised: where U is the unit of hyper-volume to ensure consistent results over different measurement units.

Original: However, our approach is not completely wrong for the following reasons.
Revised: However, our approach is not completely wrong for the following reasons.

Original: First, we consistently use the same metric across all the datasets and models.
Revised: First, we consistently use the same metric across all the datasets and models.

Original: As we have shown in Section 2, a point process can be deconstructed into modeling cardinality p ( n ) and modeling point locations.
Revised: As we have shown in Section 2, a point process can be deconstructed into modeling cardinality p ( n ) and modeling point locations.

Original: Thus, all of the models can be trivially extended to include p ( n ) and the relative differences in the likelihood would stay the same since the additional term would be a constant.
Revised: Thus, all of the models can be trivially extended to include p ( n ) and the relative differences in the likelihood would stay the same since the additional term would be a constant.

Original: Second, we use per-point likelihood which tells us how much mass is allocated on each point on average.
Revised: Second, we use per-point likelihood which tells us how much mass is allocated on each point on average.

Original: This, however, ignores cardinality in a different way.
Revised: This, however, ignores cardinality in a different way.

Original: Finally, we use this to directly compare to other related work done on modeling sets.
Revised: Finally, we use this to directly compare to other related work done on modeling sets.

Original: To summarize, we can compare models with only p ( x 1 , . . .
Revised: To summarize, we can compare models with only p ( x 1 , . . .

Original: , x n ) but need to include p ( n ) when comparing instances inside the dataset.
Revised: , x n ) but need to include p ( n ) when comparing instances inside the dataset.

Original: Figure 6 illustrates the difference perfectly.
Revised: Figure 9 illustrates the difference perfectly.

Original: Having a ground truth model Pois (64) of on inhomogeneous Poisson process means that the sample with n = 64 elements should have higher probability than n = 70 .
Revised: Having a ground truth model Pois (64) of on inhomogeneous Poisson process means that the sample with n = 64 elements should have higher probability than n = 70 .

Original: But since the per-point likelihood is larger than 1 , this is not the case.
Revised: But since the per-point likelihood is larger than 1 , this is not the case.

Original: By including the p ( n ) term, we ﬂip the difference and the probability is now correctly assigned.
Revised: By including the p ( n ) term, we ﬂip the difference and the probability is now correctly assigned.

Original: When using just per-point likelihood, the cardinality is completely ignored and sets have the same probability.
Revised: When using just per-point likelihood, the cardinality is completely ignored and sets have the same probability.

Original: [Figure]
Revised: [Figure]

Original: A. E QUIVARIANT DRIFT GIVES INVARIANT DENSITY
Revised: A. E QUIVARIANT DRIFT GIVES INVARIANT DENSITY

Original: Theorem 1 (Equivariant ﬂows) .
Revised: Theorem 1 (Equivariant ﬂows) .

Original: (Adapted from (Papamakarios et al., 2019, Section 5.6, Lemma 1))
Revised: (Adapted from (Papamakarios et al., 2019, Section 5.6, Lemma 1))

Original: Let p ( X ) be the density function of a ﬂow-based model with transformation f : R n × d → R n × d and base density q ( Z ) .
Revised: Let p ( X ) be the density function of a ﬂow-based model with transformation f : R n × d → R n × d and base density q ( Z ) .

Original: If f is equivariant with respect to Γ and q is invariant with respect to Γ , then p ( X ) is invariant with respect to Γ .
Revised: If f is equivariant with respect to Γ and q is invariant with respect to Γ , then p ( X ) is invariant with respect to Γ .

Original: Proof.
Revised: Proof.

Original: See (Papamakarios et al., 2019).
Revised: See (Papamakarios et al., 2019).

Original: Theorem 1 is satisﬁed for permutation group Γ if we, e.g., use normal or uniform base distribution, and an equivariant transformation.
Revised: Theorem 1 is satisﬁed for permutation group Γ if we, e.g., use normal or uniform base distribution, and an equivariant transformation.

Original: Theorem 2 (Equivariant ordinary differential equation) .
Revised: Theorem 2 (Equivariant ordinary differential equation) .

Original: Let f ( X ) , f : R n × d → R n × d be the solution of the ordinary differential equation dZ ( t ) /dt
Revised: Let f ( X ) , f : R n × d → R n × d be the solution of the ordinary differential equation dZ ( t ) /dt

Original: =
Revised: =

Original: g ( Z ( t ) , t ) on I =
Revised: g ( Z ( t ) , t ) on I =

Original: [ t 0 , t 1 ] with an initial condition Z ( t 0 ) =
Revised: [ t 0 , t 1 ] with an initial condition Z ( t 0 ) =

Original: X .
Revised: X .

Original: If g is equivariant with respect to Γ then so is f .
Revised: If g is equivariant with respect to Γ then so is f .

Original: Proof.
Revised: Proof.

Original: γf ( X ) = γf
Revised: γf ( X ) = γf

Original: [Equation]
Revised: [Equation]

Revised: [Equation]

Original:  =
Revised: ) =

Original: γZ
Revised: [Equation] [Equation] +

Original: [Equation]
Revised: [Equation]

Original: t 0 )
Revised: t 1 t

Revised: [Equation]

Revised: [Equation]

Original: + γ (cid:82) t 1 t
Revised: t ) , t ) dt = γZ [Equation] t

Original: [Equation]
Revised: [Equation]

Original: ( Z ( t ) , t ) dt = γZ ( t 0 ) + (cid:82) t 1 t 0 g ( γZ ( t ) , t )
Revised: + [Equation] t 1 t [Equation]

Original:  dt = f ( γZ ( t 0 ))
Revised: ( γZ ( t ) , t ) dt = f ( γZ ( t 0 ))

Original: =
Revised: =

Original: f ( γX ) , γ ∈ Γ .
Revised: f ( γX ) , γ ∈ Γ .

Original: B I MPLEMENTATION
Revised: B I MPLEMENTATION

Original: In the following we describe in detail the implementation of the normalizing ﬂows, in particular the traditional ones for baseline models.
Revised: In the following we describe in detail the implementation of the normalizing ﬂows, in particular the traditional ones for baseline models.

Original: We later describe attention, self-attention and induced attention for sets.
Revised: We later describe attention, self-attention and induced attention for sets.

Original: We provide the background and implementation of efﬁcient trace calculation for equivariant transformations.
Revised: We provide the background and implementation of efﬁcient trace calculation for equivariant transformations.

Original: We additionally discuss mini-batch training and C ONFET architecture.
Revised: We additionally discuss mini-batch training and C ONFET architecture.

Original: B. N ORMALIZING FLOW LAYERS
Revised: B. N ORMALIZING FLOW LAYERS

Original: In the following we review the traditional normalizing ﬂows, i.e. different ways to deﬁne f : R d → R d with the tractable inverse and the determinant of the Jacobian.
Revised: In the following we review the traditional normalizing ﬂows, i.e. different ways to deﬁne f : R d → R d with the tractable inverse and the determinant of the Jacobian.

Original: We denote random variables z , x ∈ R d coming from the base and target density q and p , respectively.
Revised: We denote random variables z , x ∈ R d coming from the base and target density q and p , respectively.

Original: The forward map z (cid:55)→ x is parametrized with f , and the inverse map x
Revised: The forward map z (cid:55)→ x is parametrized with f , and the inverse map x

Original: (cid:55)→ z with f − 1 .
Revised: (cid:55)→ z with f − 1 .

Original: In all of the models, the analytical inverse exists.
Revised: In all of the models, the analytical inverse exists.

Original: In practice we often focus on parametrizing the inverse map, allowing density estimation and maximum likelihood training.
Revised: In practice we often focus on parametrizing the inverse map, allowing density estimation and maximum likelihood training.

Original: Coupling layer (Dinh et al., 2015; 2017) deﬁnes the transformation of one part of the data point x conditioned on the rest.
Revised: Coupling layer (Dinh et al., 2015; 2017) deﬁnes the transformation of one part of the data point x conditioned on the rest.

Original: If x ∈ R d , we take ﬁrst k dimensions and copy them, x 1: k = z 1: k .
Revised: If x ∈ R d , we take ﬁrst k dimensions and copy them, x 1: k = z 1: k .

Original: Rest go through an afﬁne transformation with coefﬁcients as a function of z 1: k :
Revised: Rest go through an afﬁne transformation with coefﬁcients as a function of z 1: k :

Original: [Equation]
Revised: [Equation]

Original: where s, t : R k
Revised: where s, t : R k

Original: → R d − k are unrestricted neural networks.
Revised: → R d − k are unrestricted neural networks.

Original: Notice the Jacobian has a special form where the only non-zero elements are on the diagonal and in one block under the diagonal.
Revised: Notice the Jacobian has a special form where the only non-zero elements are on the diagonal and in one block under the diagonal.

Original: The determinant is then the product of the diagonal entries.
Revised: The determinant is then the product of the diagonal entries.

Original: Inverse is obtained by noticing z 1: k = x 1: k , and then z k +1:
Revised: Inverse is obtained by noticing z 1: k = x 1: k , and then z k +1:

Original: d = ( z k +1:
Revised: d = ( z k +1:

Original: d − t ( x 1: k ))
Revised: d − t ( x 1: k ))

Original: (cid:12) exp( − s ( x 1: k )) .
Revised: (cid:12) exp( − s ( x 1: k )) .

Original: The computation complexity is the same in both the sampling and density estimation directions.
Revised: The computation complexity is the same in both the sampling and density estimation directions.

Original: The function described in Eq.
Revised: The function described in Eq.

Original: 10 is called an afﬁne transformation.
Revised: 10 is called an afﬁne transformation.

Original: Instead, we can use more expressive functions such as rational-quadratic splines.
Revised: Instead, we can use more expressive functions such as rational-quadratic splines.

Original: Details of implementation can be found in (Durkan et al., 2019) 3 .
Revised: Details of implementation can be found in (Durkan et al., 2019) 3 .

Original: It deﬁnes a monotonic function with K bins where each bin is a rational-quadratic function.
Revised: It deﬁnes a monotonic function with K bins where each bin is a rational-quadratic function.

Original: Increasing K increases the expressiveness.
Revised: Increasing K increases the expressiveness.

Original: The inverse and determinant are easy to obtain.
Revised: The inverse and determinant are easy to obtain.

Original: In the context of coupling layers, one part of the data deﬁnes the spline parameters that transform the rest.
Revised: In the context of coupling layers, one part of the data deﬁnes the spline parameters that transform the rest.

Original: Autoregressive layer transforms the i th dimension based on all the previous values x : i .
Revised: Autoregressive layer transforms the i th dimension based on all the previous values x : i .

Original: We can implement this by processing a sequence ( z 1 , . . .
Revised: We can implement this by processing a sequence ( z 1 , . . .

Original: , z d ) with a neural network, in our case with a recurrent neural network that outputs the parameters of an afﬁne transformation similar to the coupling layer.
Revised: , z d ) with a neural network, in our case with a recurrent neural network that outputs the parameters of an afﬁne transformation similar to the coupling layer.

Original: The Jacobian is now a lower triangular matrix.
Revised: The Jacobian is now a lower triangular matrix.

Original: When calculating the forward direction, we know all the values of z , and can use efﬁcient parallel implementations making the calculation fast.
Revised: When calculating the forward direction, we know all the values of z , and can use efﬁcient parallel implementations making the calculation fast.

Original: However, when inverting element x i we need to know all of the previously inverted values z : i .
Revised: However, when inverting element x i we need to know all of the previously inverted values z : i .

Original: This means that autoregressive layer is inherently slow and sequential in one direction.
Revised: This means that autoregressive layer is inherently slow and sequential in one direction.

Original: Therefore, parametrizing x (cid:55)→ z direction is preferred for maximum likelihood training (Papamakarios et al., 2017), and for fast sampling, we would parametrize the z (cid:55)→ x direction (Kingma et al., 2016).
Revised: Therefore, parametrizing x (cid:55)→ z direction is preferred for maximum likelihood training (Papamakarios et al., 2017), and for fast sampling, we would parametrize the z (cid:55)→ x direction (Kingma et al., 2016).

Original: Combining layers is implemented as a composition of functions f 1 ◦ · · · ◦ f k .
Revised: Combining layers is implemented as a composition of functions f 1 ◦ · · · ◦ f k .

Original: In practice, for each transformation we know f and f − 1 together with the log-determinant of the Jacobian.
Revised: In practice, for each transformation we know f and f − 1 together with the log-determinant of the Jacobian.

Original: Therefore, for input z 0 we calculate z 1 = f 1 ( z 0 ) , . . .
Revised: Therefore, for input z 0 we calculate z 1 = f 1 ( z 0 ) , . . .

Original: , z k = f k ( z k − 1 ) .
Revised: , z k = f k ( z k − 1 ) .

Original: We accumulate the log-determinant at every step and calculate p ( z k ) using Eq. 4.
Revised: We accumulate the log-determinant at every step and calculate p ( z k ) using Eq. 4.

Original: The same principle applies for the inverse direction.
Revised: The same principle applies for the inverse direction.

Original: Why don’t we use CNFs in all models?
Revised: Why don’t we use CNFs in all models?

Original: We do not use the continuous normalizing ﬂows in models from Section 3 for two reasons.
Revised: We do not use the continuous normalizing ﬂows in models from Section 3 for two reasons.

Original: First, some models represent existing works in literature that we reimplement here, e.g. (Bender et al., 2020).
Revised: First, some models represent existing works in literature that we reimplement here, e.g. (Bender et al., 2020).

Original: Second, CNFs are not necessarily always better in terms of efﬁciency and ﬁnal performance given the task.
Revised: Second, CNFs are not necessarily always better in terms of efﬁciency and ﬁnal performance given the task.

Original: Both the variational autoencoder and inhomogeneous process are deﬁning the density on B without any constraints.
Revised: Both the variational autoencoder and inhomogeneous process are deﬁning the density on B without any constraints.

Original: Since invariance is implemented via independence, CNFs are here as good as any other normalizing ﬂow method.
Revised: Since invariance is implemented via independence, CNFs are here as good as any other normalizing ﬂow method.

Original: On the other hand, if we actually want to model interactions, like we do with C ONFET , we need to have a model that supports full Jacobian and CNFs are a perfect candidate.
Revised: On the other hand, if we actually want to model interactions, like we do with C ONFET , we need to have a model that supports full Jacobian and CNFs are a perfect candidate.

Original: B. S ELF - ATTENTION
Revised: B. S ELF - ATTENTION

Original: Attention originated in natural language processing (Vaswani et al., 2017).
Revised: Attention originated in natural language processing (Vaswani et al., 2017).

Original: Formally, given matrices
Revised: Formally, given matrices

Original: Q ∈ R n q × d k , K ∈ R n v × d k , V ∈ R n v × d v of queries, keys and values, respectively, the attention layer is deﬁned as:
Revised: Q ∈ R n q × d k , K ∈ R n v × d k , V ∈ R n v × d v of queries, keys and values, respectively, the attention layer is deﬁned as:

Original: [Equation]
Revised: [Equation]

Original: Vaswani et al.
Revised: Vaswani et al.

Original: (2017) introduce the multihead attention, an extension where we apply h attention layers in parallel, with learnable parameters that transform each matrix with a linear layer.
Revised: (2017) introduce the multihead attention, an extension where we apply h attention layers in parallel, with learnable parameters that transform each matrix with a linear layer.

Original: In our case, we are interested in a (multihead) self-attention layer that models explicit interactions between elements of sets:
Revised: In our case, we are interested in a (multihead) self-attention layer that models explicit interactions between elements of sets:

Original: SelfAttention
Revised: SelfAttention

Original: [Equation]
Revised: [Equation]

Original: X
Revised: X

Original: [Equation]
Revised: [Equation]

Original: R n × d .
Revised: R n × d .

Original: For large sets, the interaction matrix will be of size n × n .
Revised: For large sets, the interaction matrix will be of size n × n .

Original: To reduce the computation cost, Lee et al.
Revised: To reduce the computation cost, Lee et al.

Original: (2019) introduce the induced attention layer:
Revised: (2019) introduce the induced attention layer:

Original: InducedAttention
Revised: InducedAttention

Original: [Equation]
Revised: [Equation]

Original: where I ∈ R m × d is a set of learnable m points, m (cid:28) n .
Revised: where I ∈ R m × d is a set of learnable m points, m (cid:28) n .

Original: B. A TTENTION NETWORK
Revised: B. A TTENTION NETWORK

Original: The attention network is deﬁned as follows.
Revised: The attention network is deﬁned as follows.

Original: The input data of shape ( n, d ) is projected with a linear layer to ( n, h ) .
Revised: The input data of shape ( n, d ) is projected with a linear layer to ( n, h ) .

Original: We then apply self-attention layers H 

Revised: We then apply self-attention layers H

Original: = H 

Revised: = H

Original: + Attention( H 0 , H 0 , H 0 ) , where
Revised: + Attention( H 0 , H 0 , H 0 ) , where

Original: H 

Revised: H

Original: =
Revised: =

Original: X (we can also use InducedAttention m ).
Revised: X (we can also use InducedAttention m ).

Original: We apply optional layer normalization, and a single linear layer mapping proceeded with an optional residual layer.
Revised: We apply optional layer normalization, and a single linear layer mapping proceeded with an optional residual layer.

Original: After L such layers we obtain ﬁnal transformed set H L , and project it back to ( n, d ) with a linear layer.
Revised: After L such layers we obtain ﬁnal transformed set H L , and project it back to ( n, d ) with a linear layer.

Original: Using initial input and a single hidden layer neural network with sigmoid activation we obtain the gating set G .
Revised: Using initial input and a single hidden layer neural network with sigmoid activation we obtain the gating set G .

Original: The ﬁnal output of attention network is Y =
Revised: The ﬁnal output of attention network is Y =

Original: G · X + (1 − G ) · H L , a weighted combination of the initial input and the attention output.
Revised: G · X + (1 − G ) · H L , a weighted combination of the initial input and the attention output.

Original: B. E FFICIENT TRACE ESTIMATION
Revised: B. E FFICIENT TRACE ESTIMATION

Original: In this we discuss how to efﬁciently calculate the trace, without using stochastic methods.
Revised: In this we discuss how to efﬁciently calculate the trace, without using stochastic methods.

Original: We ﬁrst review the results from Chen & Duvenaud (2019).
Revised: We ﬁrst review the results from Chen & Duvenaud (2019).

Original: In general, if we have a neural network f : R d → R d we want to get the diagonal values of its Jacobian ∂f ( x ) i /∂x i .
Revised: In general, if we have a neural network f : R d → R d we want to get the diagonal values of its Jacobian ∂f ( x ) i /∂x i .

Original: A straightforward but expensive solution is to use built-in functionality of modern deep learning frameworks to calculate this sequentially for each i .
Revised: A straightforward but expensive solution is to use built-in functionality of modern deep learning frameworks to calculate this sequentially for each i .

Original: The cost is then O ( d 2 ) .
Revised: The cost is then O ( d 2 ) .

Original: Another way is to deﬁne f ( x ) i =
Revised: Another way is to deﬁne f ( x ) i =

Original: τ ( x i , h ) , where h is a d h -dimensional vector containing information about all the other values of input vector x , i.e. a conditioner network c that takes all x j , j (cid:54)
Revised: τ ( x i , h ) , where h is a d h -dimensional vector containing information about all the other values of input vector x , i.e. a conditioner network c that takes all x j , j (cid:54)

Original: =
Revised: =

Original: i as an input.
Revised: i as an input.

Original: Function τ maps R d h +1 → R .
Revised: Function τ maps R d h +1 → R .

Original: Now when we take derivative ∂f ( x ) i /∂x i it breaks into ∂τ
Revised: Now when we take derivative ∂f ( x ) i /∂x i it breaks into ∂τ

Original: i /x i and zero term from the conditioner.
Revised: i /x i and zero term from the conditioner.

Original: Therefore, we can detach the network c completely from the computation graph and just calculate ∂τ i /x i for all i in parallel.
Revised: Therefore, we can detach the network c completely from the computation graph and just calculate ∂τ i /x i for all i in parallel.

Original: When backpropagating we need to attach c back to allow learning of the parameters.
Revised: When backpropagating we need to attach c back to allow learning of the parameters.

Original: More details can be found in Chen & Duvenaud (2019).
Revised: More details can be found in Chen & Duvenaud (2019).

Original: If d h = d − 1 we recover the original network.
Revised: If d h = d − 1 we recover the original network.

Original: We often want to use much smaller values for efﬁciency.
Revised: We often want to use much smaller values for efﬁciency.

Original: In our case, we use sets of points.
Revised: In our case, we use sets of points.

Original: First step is to build an inhomogeneous transformation — a neural network that acts on individual points x i ∈ R d .
Revised: First step is to build an inhomogeneous transformation — a neural network that acts on individual points x i ∈ R d .

Original: This is the use case from above.
Revised: This is the use case from above.

Original: We build a conditioner as a masked neural network (Germain et al., 2015) that gives a hidden vector h ∈ R d h for each dimension in x i .
Revised: We build a conditioner as a masked neural network (Germain et al., 2015) that gives a hidden vector h ∈ R d h for each dimension in x i .

Original: We input the concatenation of h and a single dimension of x i to a neural network with output dimension of 1 .
Revised: We input the concatenation of h and a single dimension of x i to a neural network with output dimension of 1 .

Original: The trace is calculated as above.
Revised: The trace is calculated as above.

Original: Next, we discuss how to extend this with interactions between the points.
Revised: Next, we discuss how to extend this with interactions between the points.

Original: Sum-pooling in Eq. 8 is deﬁned as a sum of all the contributions from other points in the set.
Revised: Sum-pooling in Eq. 8 is deﬁned as a sum of all the contributions from other points in the set.

Original: We apply a function h : R d → R d to each point x i to get h i .
Revised: We apply a function h : R d → R d to each point x i to get h i .

Original: Now the contribution for element i can be simply written as (cid:80) j (cid:54) =
Revised: Now the contribution for element i can be simply written as (cid:80) j (cid:54) =

Original: i h
Revised: i h

Original: j = (cid:80)
Revised: j = (cid:80)

Original: j h j − h i .
Revised: j h j − h i .

Original: This way, we calculate (cid:80)
Revised: This way, we calculate (cid:80)

Original: j h j only once.
Revised: j h j only once.

Original: We can extend this trivially to mean-pooling by dividing everything with n − 1 .
Revised: We can extend this trivially to mean-pooling by dividing everything with n − 1 .

Original: For max-pooling we would need to get the highest and second highest values of dimensions in all h i and assign the highest to i if its value is smaller, otherwise assign the second highest value.
Revised: For max-pooling we would need to get the highest and second highest values of dimensions in all h i and assign the highest to i if its value is smaller, otherwise assign the second highest value.

Original: This way we simulate the max operation on the set excluding i th element.
Revised: This way we simulate the max operation on the set excluding i th element.

Original: Attention cannot be parallelized as easily.
Revised: Attention cannot be parallelized as easily.

Original: The issue is that a row in the similarity matrix is built by comparing current element with all the others.
Revised: The issue is that a row in the similarity matrix is built by comparing current element with all the others.

Original: It cannot be decoupled and masking out the diagonal is not enough.
Revised: It cannot be decoupled and masking out the diagonal is not enough.

Original: A pragmatic approach is to ﬁrst obtain the conditioner vectors with a masked neural network.
Revised: A pragmatic approach is to ﬁrst obtain the conditioner vectors with a masked neural network.

Original: Now, we can input them to any downstream network, including self-attention, without losing the decoupling property.
Revised: Now, we can input them to any downstream network, including self-attention, without losing the decoupling property.

Original: Another approach is to calculate attention with all the elements excluding i , and use max-pooling to get the embedding h
Revised: Another approach is to calculate attention with all the elements excluding i , and use max-pooling to get the embedding h

Original: i .
Revised: i .

Original: This is more expensive so we pursue the ﬁrst option.
Revised: This is more expensive so we pursue the ﬁrst option.

Original: B. T RAINING IN BATCHES
Revised: B. T RAINING IN BATCHES

Original: All of the described functions work with inputs of shape ( n, d ) .
Revised: All of the described functions work with inputs of shape ( n, d ) .

Original: We train in mini-batches, i.e. process multiple sets at a time.
Revised: We train in mini-batches, i.e. process multiple sets at a time.

Original: Since n varies between them, we pad them all to the biggest set with zeros and keep track of the original size.
Revised: Since n varies between them, we pad them all to the biggest set with zeros and keep track of the original size.

Original: This is important when calculating the loss since we use per-point negative log-likelihood.
Revised: This is important when calculating the loss since we use per-point negative log-likelihood.

Original: Also, when implementing interactions between the elements, we want to remove those that represent the masking.
Revised: Also, when implementing interactions between the elements, we want to remove those that represent the masking.

Original: For example, in the attention network, the elements interact with each other through the similarity matrix.
Revised: For example, in the attention network, the elements interact with each other through the similarity matrix.

Original: If most of our inputs are zero padding, this will give unexpected results after calculating softmax .
Revised: If most of our inputs are zero padding, this will give unexpected results after calculating softmax .

Original: Therefore, we zero-out the elements corresponding to the padding with inﬁnity-masking.
Revised: Therefore, we zero-out the elements corresponding to the padding with inﬁnity-masking.

Original: Same is done for other methods.
Revised: Same is done for other methods.

Original: In case we forget to mask-out the padding values before aggregation, the model will learn slower, get worse results, and in extreme cases not train at all
Revised: In case we forget to mask-out the padding values before aggregation, the model will learn slower, get worse results, and in extreme cases not train at all

Original: . B. C ONFET IMPLEMENTATION
Revised: . B. C ONFET IMPLEMENTATION

Original: We assume data comes from a region B .
Revised: We assume data comes from a region B .

Original: For x ∈ B , it holds that x i ∈ x , x i ∈ (0 , 1) , i = 1 , . . .
Revised: For x ∈ B , it holds that x i ∈ x , x i ∈ (0 , 1) , i = 1 , . . .

Original: , d .
Revised: , d .

Original: This means all of the points fall on the unit square.
Revised: This means all of the points fall on the unit square.

Original: We can achieve this by appropriately preprocessing the data.
Revised: We can achieve this by appropriately preprocessing the data.

Original: The introduced restriction is here to simplify the notation, without losing generality.
Revised: The introduced restriction is here to simplify the notation, without losing generality.

Original: The base density for all of the models is then form on B .
Revised: The base density for all of the models is then form on B .

Original: Since CNFs work on a real number domain we project the input data to R n × d with a logit transformation, apply CNF, and project it back with sigmoid transformation.
Revised: Since CNFs work on a real number domain we project the input data to R n × d with a logit transformation, apply CNF, and project it back with sigmoid transformation.

Original: The model is deﬁned as a normalizing ﬂow with three transformation, one of them being a CNF.
Revised: The model is deﬁned as a normalizing ﬂow with three transformation, one of them being a CNF.

Original: Logit is an inverse of sigmoid and their Jacobians can be calculated in closed-form.
Revised: Logit is an inverse of sigmoid and their Jacobians can be calculated in closed-form.

Original: We integrate an ODE from t 0 = 0 to t 1 = 1 .
Revised: We integrate an ODE from t 0 = 0 to t 1 = 1 .

Original: Although tuning these hyperparameters, or learning them during training is possible, we avoid doing that.
Revised: Although tuning these hyperparameters, or learning them during training is possible, we avoid doing that.

Original: If we change t 1 from 1 to 0 . 1 , the network can adapt the weights to get the same output.
Revised: If we change t 1 from 1 to 0 . 1 , the network can adapt the weights to get the same output.

Original: We extend the library provided by Chen et al.
Revised: We extend the library provided by Chen et al.

Original: (2018) 4 to allow permutation invariant densities.
Revised: (2018) 4 to allow permutation invariant densities.

Original: C E XPERIMENT DETAILS C. D ATASET PREPROCESSING AND GENERATION
Revised: C E XPERIMENT DETAILS C. D ATASET PREPROCESSING AND GENERATION

Original: For synthetic datasets we take care of the edge effects by simulating on a region larger than (0 , 1) × (0 , 1) , e.g. for Matérn clustering process with R = 0 .
Revised: For synthetic datasets we take care of the edge effects by simulating on a region larger than (0 , 1) × (0 , 1) , e.g. for Matérn clustering process with R = 0 .

Original: 1 we extend the simulation square by 0 .
Revised: 1 we extend the simulation square by 0 .

Original: 1 on all sides.
Revised: 1 on all sides.

Original: The points that fall out of the observed region are discarded.
Revised: The points that fall out of the observed region are discarded.

Original: Thomas and Matérn are simulated as it is described in the main text.
Revised: Thomas and Matérn are simulated as it is described in the main text.

Original: Checkins-NY and Checkins-Paris use data of user location logs that are recorded over the period of 2009 and 2010.
Revised: Checkins-NY and Checkins-Paris use data of user location logs that are recorded over the period of 2009 and 2010.

Original: All locations for a user are a single realization of the underlying process.
Revised: All locations for a user are a single realization of the underlying process.

Original: Initial data contains just under 6 . 5 million locations.
Revised: Initial data contains just under 6 . 5 million locations.

Original: We take the locations from two cities, New York, between coordinates (40 .
Revised: We take the locations from two cities, New York, between coordinates (40 .

Original: 6829 , − 74 . 0479) and (40 .
Revised: 6829 , − 74 . 0479) and (40 .

Original: 8820 , − 73 . 9067) , and Paris, between (48 .
Revised: 8820 , − 73 . 9067) , and Paris, between (48 .

Original: 5 , 2) and (49 , 3) , respectively.
Revised: 5 , 2) and (49 , 3) , respectively.

Original: Data contains very large sets, largest one having 904 elements.
Revised: Data contains very large sets, largest one having 904 elements.

Original: [Figure]
Revised: [Figure]

Revised: Mixture

Revised: Thomas

Revised: [Table]

Revised: Original model

Original:  Crimes dataset contains recordings of crimes that happened in the city of Portland from March2012 to December 31 2012, with 309 days in total.
Revised: [Table] With added inhomogeneous layers [Table] Crimes dataset contains recordings of crimes that happened in the city of Portland from March2012 to December 31 2012, with 309 days in total.

Original: It contains 146927 logs, each with a time, coordinate and type of case.
Revised: It contains 146927 logs, each with a time, coordinate and type of case.

Original: We aggregate all the occurrences in one day to get a single realization.
Revised: We aggregate all the occurrences in one day to get a single realization.

Original: This gives us very large sets, smallest one having 298 and largest 736 elements.
Revised: This gives us very large sets, smallest one having 298 and largest 736 elements.

Original: C. A DDITIONAL SAMPLES
Revised: C. C OMPARISON TO OTHER EQUIVARIANT FLOWS

Revised: In addition to the main experiments in Section 6, here we compare to the model by (Köhler et al., 2020) that also uses an equivariant ﬂow with closed-form trace. The results are shown in Table Their model is designed for physics simulations so it has to include rotation and translation symmetry.

Revised: They achieve this by transforming the points based on their distances using a simple Gaussian kernel.

Revised: Because of this limitation, it is unable to model the inhomogeneous part of the process as can be seen from the results.

Revised: To combat this we added additional spline coupling layers but the ﬁnal results are still much worse than our model.

Revised: We conclude that our equivariant transformation (e.g. Eq.

Revised: 8) is well suited for the problem of modeling sets and achieves better performance than all competing methods while having closed-form trace computation.

Revised: C. A DDITIONAL SAMPLES

Original: Figure 8 shows 6 samples drawn from models trained on Thomas dataset.
Revised: Figure 11 shows 6 samples drawn from models trained on Thomas dataset.

Original: In the ﬁrst row we can see ground truth samples.
Revised: In the ﬁrst row we can see ground truth samples.

Original: The clustering behavior is not captured by IHP model.
Revised: The clustering behavior is not captured by IHP model.

Original: Other models are able to generate convincing samples.
Revised: Other models are able to generate convincing samples.

Original: We used the same random seed for sampling the initial set
Revised: We used the same random seed for sampling the initial set

Original: Z ∼ p
Revised: Z ∼ p

Original: ( Z ) in all C ONFET models.
Revised: ( Z ) in all C ONFET models.

Original: Remarkably, the similarity of the output is prominent even though the architectures and parameters are different.
Revised: Remarkably, the similarity of the output is prominent even though the architectures and parameters are different.

Original: We conclude that C ONFET models learn similar distributions p ( X ) given the same dataset. 
Revised: We conclude that C ONFET models learn similar distributions p ( X ) given the same dataset. [Figure] [Figure] [Table]

Original: C. S PEED COMPARISON
Revised: S PEED COMPARISON

Original: In Figure 5 (Left) we compare the time to calculate the trace using three different methods: naive calculation, one step of stochastic estimator and using a cheap exact computation as described before.
Revised: In Figure 5 (Left) we compare the time to calculate the trace using three different methods: naive calculation, one step of stochastic estimator and using a cheap exact computation as described before.

Original: We use the same C ONFET -E XACT architecture for all three with input dimension of 2 , hidden dimension 64 , and a single attention layer.
Revised: We use the same C ONFET -E XACT architecture for all three with input dimension of 2 , hidden dimension 64 , and a single attention layer.

Original: We vary the number of input set elements from 2 4 to 2 8 .
Revised: We vary the number of input set elements from 2 4 to 2 8 .

Original: The experiment was repeated 5 times on the hardware described in C.5.
Revised: The experiment was repeated 5 times on the hardware described in C.6.

Original: C. T RAINING CURVES
Revised: C. T RAINING CURVES

Original: In Figure 5 (Middle) we overlay the validation losses recorded during training for C ONFET -E XACT and C
Revised: In Figure 5 (Middle) we overlay the validation losses recorded during training for C ONFET -E XACT and C

Original: ONFET -S TOCHASTIC .
Revised: ONFET -S TOCHASTIC .

Original: We picked two models that have the best ﬁnal validation loss after hyperparameter search.
Revised: We picked two models that have the best ﬁnal validation loss after hyperparameter search.

Original: We plot all 5 runs.
Revised: We plot all 5 runs.

Original: Figure 9 shows the validation loss change across training epochs for all our models on Mixture,
Revised: Figure 12 shows the validation loss change across training epochs for all our models on Mixture,

Original: Matérn and Thomas dataset.
Revised: Matérn and Thomas dataset.

Original: The chosen models are those that achieved the best average validation loss.
Revised: The chosen models are those that achieved the best average validation loss.

Original: The different number of epochs in different plots comes from the early stopping.
Revised: The different number of epochs in different plots comes from the early stopping.

Original: We can notice that C
Revised: We can notice that C

Original: ONFET -S
Revised: ONFET -S

Original: TOCHASTIC has more noisy training compared to C ONFET -E XACT and C ONFETF IXED because of its stochastic trace estimator.[Figure] [Figure] [Table]
Revised: TOCHASTIC has more noisy training compared to C ONFET -E XACT and C ONFETF IXED because of its stochastic trace estimator.

Original: [Table]

Original:  H ARDWARE
Revised: C. H ARDWARE

Original: We train all of our models on a single Nvidia GTX1080Ti GPU (12GB).
Revised: We train all of our models on a single Nvidia GTX1080Ti GPU (12GB).

Original: The machine has 256GB
Revised: The machine has 256GB

Original: RAM and an Intel Xeon E5-2630 v4 @ 2.20 GHz CPU.
Revised: RAM and an Intel Xeon E5-2630 v4 @ 2.20 GHz CPU.

Original: D A BLATION STUDIESD. H YPERPARAMETER SEARCH
Revised: D A BLATION STUDIES

Revised: D . H YPERPARAMETER SEARCH

Original: In all of the models we tried placing the weight decay of 10 − 3 on the weights.
Revised: In all of the models we tried placing the weight decay of 10 − 3 on the weights.

Original: This sometimes leads to more stable training and better performance.
Revised: This sometimes leads to more stable training and better performance.

Original: Using smaller learning rate ( 10 − 4 ) makes training slow without other beneﬁts.
Revised: Using smaller learning rate ( 10 − 4 ) makes training slow without other beneﬁts.

Original: We did not try learning rate schedulers.
Revised: We did not try learning rate schedulers.

Original: In IHP we use a normalizing ﬂow with L ∈ { 1 , 5 } spline coupling layers, each spline deﬁned with
Revised: In IHP we use a normalizing ﬂow with L ∈ { 1 , 5 } spline coupling layers, each spline deﬁned with

Original: K ∈ { 5 , 10 } knots.
Revised: K ∈ { 5 , 10 } knots.

Original: Autoregressive model stacks L a ∈ { 5 , 10 } autoregressive and L c ∈ { 5 , 10 } set coupling layers.
Revised: Autoregressive model stacks L a ∈ { 5 , 10 } autoregressive and L c ∈ { 5 , 10 } set coupling layers.

Original: When using attention in IWAE, besides the number of layers L ∈ { 1 , 5 } , we deﬁne the number of heads H ∈ { 1 , 8 } .
Revised: When using attention in IWAE, besides the number of layers L ∈ { 1 , 5 } , we deﬁne the number of heads H ∈ { 1 , 8 } .

Original: In all of the above, having more layers (bigger model) does not improve the result.
Revised: In all of the above, having more layers (bigger model) does not improve the result.

Original: Equivariant transformations are deﬁned with L ∈ { 2 , 5 } deep set or attention layers.
Revised: Equivariant transformations are deﬁned with L ∈ { 2 , 5 } deep set or attention layers.

Original: We use mean or max-pooling for aggregation.
Revised: We use mean or max-pooling for aggregation.

Original: We noticed that having 2 or 5 attention layers in C ONFETS TOCHASTIC does not inﬂuence the results much.
Revised: We noticed that having 2 or 5 attention layers in C ONFETS TOCHASTIC does not inﬂuence the results much.

Original: Because of that, using a bigger model is unnecessary for our datasets.
Revised: Because of that, using a bigger model is unnecessary for our datasets.

Original: We tried Frobenius norm regularization of an ODE (Finlay et al., 2020), but did not observe big changes in performance.
Revised: We tried Frobenius norm regularization of an ODE (Finlay et al., 2020), but did not observe big changes in performance.

Original: D. T HE ROLE OF ADDITIONAL LAYERS IN C ONFET -F IXED
Revised: D. T HE ROLE OF ADDITIONAL LAYERS IN C ONFET -F IXED

Original: The C ONFET -F IXED model has only interactions between the points in its drift function so it is unable to model the inhomogeneous part of the process.
Revised: The C ONFET -F IXED model has only interactions between the points in its drift function so it is unable to model the inhomogeneous part of the process.

Original: That is the reason why we add additional coupling layers, same as those in IHP, along with the CNF layer.
Revised: That is the reason why we add additional coupling layers, same as those in IHP, along with the CNF layer.

Original: Here, we compare the two versions of the model, one with 5 of these layers and other without any.
Revised: Here, we compare the two versions of the model, one with 5 of these layers and other without any.

Original: We additionally include IHP into  comparison.
Revised: We additionally include IHP into [Table] comparison.

Original: The results in Table 3 show that we need the layers to achieve competitive results.
Revised: The results in Table 4 show that we need the layers to achieve competitive results.

Original: They are on par with IHP on a synthetic inhomogeneous dataset.
Revised: They are on par with IHP on a synthetic inhomogeneous dataset.

Original: It performs much better on datasets with interactions.
Revised: It performs much better on datasets with interactions.

Original: A version with no coupling layers has a better test loss score than IHP on Thomas dataset.
Revised: A version with no coupling layers has a better test loss score than IHP on Thomas dataset.

Original: D. W HEN TO USE INDUCED ATTENTION ?
Revised: D. W HEN TO USE INDUCED ATTENTION ?

Original: Table 4 shows the comparison between using induced ( m = 8 ) and regular self-attention.
Revised: Table 5 shows the comparison between using induced ( m = 8 ) and regular self-attention.

Original: We can see that induced attention usually performs worse.
Revised: We can see that induced attention usually performs worse.

Original: We did not tune the hyperparameter m so the results can be different for the optimal value.
Revised: We did not tune the hyperparameter m so the results can be different for the optimal value.

Original: Having an extra hyperparameter is a drawback, however, the lower memory footprint of induced attention can help train models on sets with very large cardinality
Revised: Having an extra hyperparameter is a drawback, however, the lower memory footprint of induced attention can help train models on sets with very large cardinality.