Neural Conservation Laws: A Divergence-Free Perspective
Abstract
We investigate the parameterization of deep neural networks that are by designdivergence-free.
To motivate the importance of divergence-free models, we observethat the continuity equation, a fundamental conservation law, can be characterizedthrough a divergence-free vector field.
We hence propose an approach to buildingdivergence-free neural networks through the concept of differential forms, andwith the aid of automatic differentiation, realize two practical constructions withdiffering trade offs.
Furthermore, we show these models are universal and canbe used to represent any divergence-free vector field if given sufficient capacity.
As a result, we can parameterize density and vector fields that always satisfy thecontinuity equation simply by design, foregoing the need for expensive numericalsimulation.
We experimentally validate our approaches on neural network-basedsolutions to fluid equations and on learning dynamical optimal transport maps.
Introduction
Modern successes in deep learning are often attributed to the expressiveness of black-box neuralnetworks.
These models are known to be universal function approximators [Hornik et al., 1989] – butthis flexibility comes at a cost.
In contrast to other parametric function approximators such as Finite Elements [Schroeder and Lube, 2017], it is hard to design networks that have baked-in constraintsat a structural level.
Existing approaches often resort to penalty methods for constraints—but theseincrease the cost of training and can produce inaccuracies in downstream applications when theconstraints are not perfectly satisfied.
Theoretical analysis of soft-constrained models also becomesmore difficult, as constraints cannot be assumed to be exactly satisfied.
On the other hand, enforcinghard constraints on the architecture can be difficult, and even once enforced, it is often unclearwhether the model remains sufficiently expressive within the constrained function class.
In this work, we discuss an approach to directly bake in two constraints into deep neural networks: (i)having a divergence of zero, and (ii) satisfying the continuity equation.
One of our key insights is thatthe former directly leads into the latter, so the first portion of the paper focuses on divergence-freevector fields.
These represent a special class of vector fields which have widespread use in thephysical sciences.
In computational fluid dynamics, divergence-free vector fields are often usedto model incompressible fluid interactions formalized by the Euler or Navier-Stokes equations.
In 29 R 3 we know that the curl of a vector field has a divergence of zero, which has seen many uses ingraphics simulations ( e.g. Eisenberger et al.
Lurking behind this fact is the generalizationof divergence and curl through differential forms [Cartan, 1899], and the powerful identity d 2 = 0 .
We first explore this generalization, then discuss two constructions derived from it for parmaterizingdivergence-free vector fields.
While this approach has been partially discussed previously [Barbarosie,2011, Kelliher, 2021], it is not extensively known and has not been validated in the machine learning community. [Equation] v = [Equation]
A ) eq.
[Equation] 4) v Concretely, we derive two approaches—visualized in Figure 1—to transform sufficiently smoothneural networks into divergence-free neural networks, which are made efficient by automatic differentiation and recent advances in practical composable transformations [Bradbury et al., 2018].
Furthermore, we show that our approaches can theoretically represent any divergence-free vectorfield.
We combine these new modeling tools with a novel observation that solutions of the continuityequation—a partial differential equation describing the relation between a flow and its transformeddensity—can be characterized jointly as a divergence-free vector field.
As a result, we can parameterize neural networks that, by design, always satisfy the continuity equation, which we coin Neural
Conservation Laws .
While prior works either resorted to penalizing errors [Raissi et al., 2019] ornumerically simulating the density given a flow [Chen et al., 2018], baking this directly into themodel as a constraint allows us to forego extra penalty terms and expensive numerical simulations.
Constructing divergence-free vector fields
We will use the notation of differential forms in R n for deriving the divergence-free and universalityproperties of our vector field constructions.
We provide a concise overview of differential forms in
Appendix A. For a more extensive introduction see e.g. , Do Carmo [1998], Morita [2001].
Withoutthis formalism, it is difficult to show either of these two properties; however, readers who wish toskip the derivations can jump to the matrix formulations in equation 4 and equation 7.
Let us denote A k ( R n ) as the space of differential k -forms, d : A k ( R n )
→ A k +1 ( R n ) as the exteriorderivative, and ⋆ : A k ( R n )
→ A n − k ( R n ) as the hodge star operator that maps each k -form to an( n - k )-form.
Then we can express the divergence div ( v ) as the composition d ⋆ v .
To parameterize a divergence-free vector field v : R n → R n , we note that from the properties ofdifferential forms, taking an arbitrary ( n - 2 )-form µ ∈ A n − 2 ( R n ) we have that
[Equation]
=
d ( dµ
[-i ∧ dx i = 
-]
[Equation]
(1) and therefore the vector field (1-form) v defined via ⋆v = dµ is divergence free.
Applying the hodgestar to both sides we get that v =
( − 1) n − 1 ⋆ dµ is divergence-free. The same is true without theconstant ( − 1) n − 1 , so it suffices to take
[Equation]
We write the parameterization v = ⋆dµ explicitly.
Since a basis for A n − 2 ( R n ) can be chosen to be⋆ ( dx i ∧ dx j ) for all i <
j we can write µ = (cid:80) i<j µ ij ⋆ ( dx i ∧ dx j ) .
Now a direct calculation showsthat up to a constant sign
[Equation]
This formula leads to a simple matrix formulation: If we let A : R n → R n × n be the anti-symmetricmatrix-valued function where A ij = µ ij for i > j and A ij = − µ ji then the divergence free vectorfield v = ⋆dµ can be written as taking row-wise divergence of A , i.e. ,
[Equation]
However, this requires parameterizing O ( n 2 ) functions. A more compact representation, which startsfrom a vector field, can also be derived.
The idea behind this second construction is to model µinstead as µ = δν , where ν ∈ A n − 1 ( R n ) .
Putting this together with equation 2 we get that
[Equation]
is a divergence free vector field.
To provide equation 5 in matrix formulation we first write ν ∈
A n − 1 ( R n ) in the ( n - 1 )-form basis, i.e. ν = (cid:80) ni
=
ν i ⋆ dx i .
Then a direct calculation provides up toa constant sign
[Equation]
This means that, given an arbitrary vector field b : R n → R n , and denoting J b as the Jacobian of b ,we can construct A as
[Equation]
By using automatic differentiation and recent advances in composable function transforms, we cancompute this expression and then its row-wise divergence (equation 4).
To summarize, we have two constructions for divergence free vector fields v :
Matrix-field : (equations 2 and 4) v is represented using an anti-symmetric matrix field A .
Vector-field : (equations 5 and 4+7) v is represented using a vector field b . [Figure]
As we show in the next section, these two approaches are maximallyexpressive ( i.e. universal), meaning they can approximate arbitrarysmooth divergence-free vector fields.
However, empirically thesetwo constructions can exhibit different practical trade-offs.
Thematrix-field construction has a computational advantage as it requiresone less Jacobian computation, though it requires more components— O ( n ) vs. O ( n 2 ) —to represent than the vector-field construction.
This generally isn’t a concern as all components can be computed inparallel. However, the vector-field construction can make it easy tobake in additional constraints; an example of this is discussed in Section 7.1, where a non-negativityconstraint is imposed for modeling continuous-time probability density functions. In Figure 2 weshow wallclock times for evaluating the divergence-free vector field based on our two constructions. Both exhibit quadratic scaling with the number of dimensions due to the row-wise divergence inequation 4, while the vector-field construction has an extra Jacobian computation.
2.1 Universality
Both the matrix-field and vector-field representations are universal, i.e. they can model arbitrarydivergence-free vector fields.
The main tool in the proof is the Hodge decomposition theorem [Morita,2001, Berger, 2003].
For simplicity we will be working on a periodic piece of R n , namely the torus 103 T =
[ − M, M ] n / ∼ where ∼ means identifying opposite edges of the n -dimensional cube, and M > 0 is arbitrary.
Vector fields with compact support can always be encapsulated in T with M > 0sufficiently large.
Since T is locally identical to R n , all the previous definitions and constructionshold.
Theorem 2.1.
The matrix and vector-field representations are universal in T , possibly only missinga constant vector field.
A formal proof of this result is in Appendix B.
Neural Conservation Laws
We now discuss a key aspect of our work, which is parameterizing solutions of scalar conservationlaws.
Conservation laws are a focal point of mathematical physics and have seen applications inmachine learning, with the most well known examples being conserved scalar quantities often referredto as density, energy, or mass.
[-Since we assumed v is div-free, d ⋆ v = 
 and we get dδτ = 0 . Since dδτ-]Formally, a conservation law can be expressed as a first-order PDEwritten in divergence form as ∂ρ∂t + div ( j ) = 0 , where j is known as the flux, and the div is taken overspatial variables.
In the case of the continuity equation, there is a velocity field u which describes theflow and the flux is equal to the density times the velocity field:
[Equation]
(8) where ρ : R n → R and u : R n → R n .
One can then interpret the equation to mean that utransports ρ (0 , · ) to ρ ( t, · ) continuously —without teleporting, creating or destroying mass. Such aninterpretation plays a key role in physics simulations as well as the dynamic formulation of optimaltransport [Benamou and Brenier, 2000].
In machine learning, the continuity equation appears whenusing continuous normalizing flows [Chen et al., 2018]—which also have been used to approximatesolutions to dynamical optimal transport.
[Finlay et al., 2020, Tong et al., 2020, Onken et al., 2021] These, however, only model the velocity u and rely on numerical simulation to solve for the densityρ , which can be costly and time-consuming.
Instead, we observe that equation 8 can be expressed as a divergence-free field by augmenting thespatial dimensions with the time variable, resulting in a vector field v that as input ( t, x ) and outputs( ρ, ρu ) .
Then equation 8 is equivalent to
[Equation]
where the divergence operator is now taken with respect to the joint system ( t, x ) , i.e. ∂∂t + (cid:80) ni =1 ∂∂x
[Equation] We thus propose modeling solutions of conservation laws by parameterizing the divergence-freevector field v .
This allows us to enforce the continuity equation at an architecture level.
Specifically,we parameterize a divergence-free vector field v and set v 1 = ρ and v 2: n +1 = ρu , allowing us torecover the vector field as u = v 2: n +1 ρ assuming ρ ̸ = 0 .
This allows us to forego computationallyexpensive procedures for approximating or simulating ρ given the flow.
Related Works
Deep learning Existing approaches to enforcing constraints in deep neural networks can induceconstraints on the derivatives, such as convexity [Amos et al., 2017] or Lipschitz continuity
[Miyatoet al., 2018].
More complicated formulations involve solving numerical problems such as usingsolutions of convex optimization problems [Amos and Kolter, 2017], solutions of fixed-pointsiterations [Bai et al., 2019], and solutions of ordinary differential equations [Chen et al., 2018].
Thesemodels can help provide more efficient approaches or alternative algorithms for handling downstreamapplications such as constructing flexible density models [Chen et al., 2019, Lu et al., 2021] orapproximating optimal transport paths [Tong et al., 2020, Makkuva et al., 2020, Onken et al., 2021].
However, in many cases, there is a need to solve a numerical problem in which the solution mayonly be approximated up to some numerical accuracy; for instance, the need to compute the densityflowing through a vector field under the continuity equation [Chen et al., 2018].
Applications of differential forms Differential forms and more generally, differential geometry,have been previously applied in manifold learning—see e.g. Arvanitidis et al. [2021] and Bronsteinet al. [2017] for an in-depth overview.
Most of the applications thus far have been restricted to 2 or 3dimensions—either using identities like div ◦ curl = 0 in 3D for fluid simulations [Rao et al., 2020],or for learning geometric invariances in 2D images or 3D space [Gerken et al., 2021, Li et al., 2021].
Neural approximations to PDE solutions
As a demonstration of our method, we apply it to neural-based PDE simulations of fluid dynamics.
First, we apply it to a problem consisting of modelling idealized fluid flow on the flat Torus, then, to showcase the ability to handle boundaries, another in the open unit ball B ⊆ R 3 .
While these are toyexamples, they demonstrate the value of our method in comparison to existing approaches—namelythat we can exactly satisfy the continuity equation and preserve exact mass.
The Euler equations of incompressible flow The incompressible Euler equations [Feynman et al.
,1989] form an idealized model of inviscid fluid flow, governed by the system of partial differentialequations
[Equation]
in three unknowns: the fluid velocity u ( t, x ) ∈ R 3 , the pressure p ( t, x ) , and the fluid density ρ ( t, x ) .
While the fluid velocity and density are usually given at t = 0 , the initial pressure is not required.
Typically, on a bounded domain Ω ⊆ R n , these are supplemented by the free-slip boundary conditionand initial conditions
[Equation]
The density ρ plays a critical role since in addition to being a conserved quantity, it influences thedynamics of the fluid evolution over time.
In numerical simulations, satisfying the continuity equationas closely as possible is desirable since equations in equation 10 are coupled.
Error in the densityfeeds into error in the velocity and then back into the density over time.
In the finite element literature,a great deal of effort has been put towards developing conservative schemes that preserve mass (orenergy in the more general compressible case)—see Guermond and Quartapelle [2000] and Almgrenet al.
[1998] for an overview.
But since the application of physics informed neural networks (PINNs)to fluid problems is much newer, conservative constraints have only been incorporated as penaltyterms into the loss [Mao et al., 2020, Jin et al., 2021].
5.1 Physics informed neural networks
Physics Informed Neural Networks (PINNs; Raissi et al.
[2019, 2017]) have recently received renewedattention as an application of deep neural networks.
While using neural networks as approximatesolutions to PDE had been explored before Lagaris et al.
[1998], modern advances in automaticdifferentiation algorithms have made the application to much more complex problems feasible [Raissiet al., 2019].
The “physics” in the name is derived from the incorporation of physical terms into theloss function, which consist of adding the squared residual norm of a PDE.
For example, to train aneural network ϕ
=
[ ρ, p, u ] to satisfy the Euler equations, the standard choice of loss to fit to is
[Equation]
where γ = ( γ F , γ I , γ div γ Cont , γ G ) denotes suitable coefficients (hyperparameter).
The loss term L Gensures fluid does not pass through boundaries, when they are present.
Similar approaches were takenin [Mao et al., 2020] and [Jagtap et al., 2020] for related equations.
While schemes of this nature arevery easy to implement, they have the drawback that since PDE terms are only penalized and notstrictly enforced, one cannot make guarantees as to the properties of the solution.
To showcase the ability of our method to model conservation laws, we will parameterize the densityand vector field as v =
[ ρ, ρu ] , as detailed in Section 3.
This means we can omit the term L Cont asdescribed in Section 5.1 from the training loss.
The divergence penalty, L div remains when modelingincompressible fluids, since u is not necessarily itself divergence-free.
In order to stablize training,we can modify the loss terms L F , L G , L I to avoid division by ρ .
This is detailed in Appendix B.2.
5.2 Incompressible variable density inside the 3D unit ball
We first construct a simple unit within a bounded domain, specifically, we will consider the Eulerequations inside B (0 , 1) ⊆
R 3 , with the initial conditions
[Equation]
t = 0.
t = 0.
t = 0.
[Figure]
t = 0.
t = 0. t = 0. [Figure] As a baseline, we compare against a parameterization that uses the divergence-free vector field onlyfor the incompressible constraint.
This is equivalent to the curl approach described in Rao et al.[2020]; another neural network is used to parameterize ρ .
In contrast, our generalized formulationof divergence-free fields allow us to model the continuity equation is this case as it requires fourdimensional inputs.
We parameterize our method using v =
[ ρ, ρu ] , as detailed in Section 3, with thevector-field approach Section 2.
The results for training both networks in the described manner are shown in Figure 3.
Surprisingly,even though both architectures achieve similar training loss after a comparable number of iterations(see Appendix C), the curl approach fails to advect the density accurately, whereas our conservationapproach evolves the density as it should given the initial conditions.
5.3 Incompressible variable density flow on T 2
As a more involved experiment, we model the Euler equations on the flat 2-Torus T 2 .
This wasparameterized as a multilayer perceptron pre-composed with the periodic embedding ι : [0 , 1] 2 → T [-
-]2ι ( x, y ) =
[cos(2 πx ) , sin(2 πx ) , cos(2 πy ) , sin(2 πy )]
The continuity equation is enforced exactly, using the method detailed in Section 3, so the L Cont termis omitted from L total .
The initial velocity and density {+areρ 
+}
( x, y ) = ( z 1 + z 3 ) 2 +v 0 ( x, y ) = [ e z 3 , e z 1 / 2] where
z = θ ( x, y, 1) .
Here we used the matrix-field parameterization as referred to in Section 2, asit is empirically more flexible in low dimensions.
We also include the harmonic component from
Hodge decomposition, which is just a constant vector field, to ensure we have full universality.
Fullmethodology for the experiment is detailed in Appendix C.
The results of evolving the flow for timet ∈ [0 , 1 / 3 ] are shown in Figure 4, as well as a comparison with a reference finite element solutiongenerated by the conservative splitting method from [Guermond and Quartapelle, 2000].
In general,our model fits the initial conditions well and results in fluid-like movement of the density, howeversmall approximation errors can accumulate from simulating over a longer time horizon.
To highlightthis, we also plotted the best result we were able to achieve with a standard PINN.
Although the PINN made progress minimizing the loss, was unable to properly model the dynamics of the fluidevolution. While this failure can likely be corrected by tuning the γ constant (see Section 5.1), wewere unable to find a combination that fit the equation correctly. By contrast, our model worked withthe first configuration we tried.
Learning the Hodge decomposition
With our divergence-free neural network, we can also learn the Hodge decomposition of a vectorfield itself.
Informally, the Hodge decomposition states that any vector field can be expressed as
[Figure]
PINN
[Figure]
the combination of a gradient field ( i.e. gradient of a scalar function) and a divergence-free field.
To motivate, there are many scenarios where the neural network should be a gradient field at theoptimum, such as score matching [Song and Ermon, 2019], iterative denoising [Hurault et al.
, 2022],amortized optimization [Xue et al., 2020], and meta-learned optimization [Andrychowicz et al., 2016].
Even though the target is a gradient field, many works have reported that it is easier to work withunconstrained vector field.
Learning the Hodge decomposition can then correspond to removing anyextra divergence-free, i.e. rotational, components of the vector field.
For this experiment, we construct a synthetic scenario where we have access to the ground truth
Hodge decomposition, in order to validate the accuracy of our method. As such, we construct thetarget field ˆ v as [Equation] (13) where η is a fixed divergence-free neural network, and w is a fixed scalar network.
We use the sameembedding from Section 5.3, but extended to T n , for n ∈ { 25 , 50 , 100 } .
We then parameterizeanother divergence-free model v θ which has a different architecture than the target η .
We proposeusing the following objective for training: [Equation] In the notation of differential forms, this corresponds to ∥ δ (ˆ v − v θ ) ∥ F , which when zero impliesv θ =
ˆ v + δµ , where µ is an arbitrary n -form.
However, since v θ is divergence free, it must holdδµ = −∇ w , i.e. the left-over portion of ˆ v − v θ is a gradient field.
This follows because the componentsof the Hodge decomposition are orthogonal [Do Carmo, 1998].
[Figure]
For evaluation, we plot the squared L 2 norm to the ground truthη against wall clock time.
While these dimensions are not highby machine learning standards, they are significantly beyondthose achievable using grid based discretizations of divergencefree fields.
The results are shown in Figure 5.
While largerdimensions are slower, we are still able to recover the groundtruth in a reasonable amount of time in dimensions up to 100.
Dynamical optimal transport
The calculation of optimal transport (OT) maps between two densities of interest p 0 and p 1 servesas a powerful tool for computing distances between distributions.
Moreover, the inclusion of
OT into machine learning has inspired multiple works on building efficient or stable generativemodels [Arjovsky et al., 2017, Tong et al., 2020, Makkuva et al., 2020, Huang et al., 2020, Finlayet al., 2020, Rout et al., 2021, Rozen et al., 2021].
In particular, we are interested in the dynamical formulation of Benamou and Brenier [2000], which is a specific instance where the transport map ischaracterized through a vector field.
Specifically, we wish to solve the optimization problem
[Equation]
subject to the boundary constraints, ρ (0 , · )
= p 0 ( · ) and ρ (1 , · )
=
p 1 ( · ) , as well as remainingnon-negative ρ ≥ 0 on the domain and satisfying the continuity equation ∂ρ∂t + div ( ρu ) = 0 .
In many existing cases, the models for estimating either ρ and equation 15 are completely decoupledfrom the generative model or transport map [Arjovsky et al., 2017, Makkuva et al., 2020], or insome cases ρ would be computed through numerical simulation and equation 15 estimated throughsamples [Chen et al., 2018, Finlay et al., 2020, Tong et al., 2020].
These approaches may not be idealas they require an auxiliary process to satisfy or approximate the continuity equation.
In contrast, our NCL approach allows us to have access to a pair of ρ and u that always satisfies thecontinuity equation, without the need for any numerical simulation.
As such, we only need to matchthe boundary conditions, and will automatically have access to a vector field u that transports massfrom ρ (0 , · ) to ρ (1 , · ) .
In order to find an optimal vector field, we can then optimize u with respect tothe objective in equation 15.
In this setting, it is important that our density is non-negative for the entire domain and within thetime interval between 0 and 1.
Instead of adding another loss function to satisfy this condition, weshow how we can easily bake this condition into our model through the vector-field parameterization.
Recall for the vector-field parameterization of the continuity equation, we have a vector field b θ ( y )with y =
[ t, x ] , for t ∈ [0 , 1] , x ∈ R n .
This results in the parameterization [Equation] for i = 1 , . . . , n + [Equation]
We now make the adjustment that ∂b 2: n +1 ∂t = 0 .
This condition does not reduce the expressivity of theformulation, as ρ can still represent any density function.
Interestingly, this modification does implythat ρ = ∇ 2 b 1 with ∇ 2 being the Laplacian operator, so that satisfying the boundary conditions isequivalent to solving a time-dependent Poisson’s equation ∇ 2 b 1 ( t, · ) =
p t .
Furthermore, the class of subharmonic functions is exactly the set of functions g that satisfies∇ 2 g (0 , · ) ≥ 0 .
In one dimension, this corresponds to the class of convex functions but is a strictgeneralization of convex functions in higher dimensions.
With this in mind, we propose using theparameterization
[Equation]
where ϕ ( z ) =
14 π (cid:0) log(
∥ z ∥ 2 ) + E 1 ( ∥ z ∥ 2 / 4 ) (cid:1) (17) with E 1 being the exponential integral function, and w k , a k ∈ R + and b k ∈ R n are free functionsof t .
This a weighted sum of generalized linear models, where the choice of nonlinear function ϕ ischosen because it is the exact solution to the Poisson’s equation for a standard normal distributionwhen n = 2 , i.e. ∇ 2 ϕ
=
N (0 , I ) , while for n ≥ 2 this remains a subharmonic function.
As such,this can be seen as a generalization of a mixture model, which are often referenced to be universaldensity approximators [Goodfellow et al., 2016], albeit requiring many mixture components. 7.2 Experiments
We experiment with pairs of 2D densities (visualized in Figure 6) and fit a pair of ( ρ, u ) to equation 15while satisfying boundary conditions.
Specifically, we train with the loss
[Equation]
where ˜ p i is a mixture between p i and a uniform density over a sufficiently large area, for i = 0 , 1 ,and λ is a hyperparameter.
We use a batch size of 256 and set K =128 (from equation 17).
Qualitativeresults are displayed in Figure 6, which include interpolations between pairs of densities.
[Figure]
[Figure] Furthermore, we can use the trained models to estimate the Wasserstein distance, by transformingsamples from p 0 through the learned vector field to time t =1.
We do this for 5000 samples andcompare the estimated value to (i) one estimated by that of a neural network trained through aminimax optimization formulation [Makkuva et al., 2020], and (ii) one estimated by that of a discrete OT solver [Bonneel et al., 2011] based on kernel density approximation and interfaced through thepot library [Flamary et al., 2021].
Estimated squared Wasserstein distance as a function of wallclock time are shown in Figure 7.
Wesee that our estimated values roughly agree with the discrete OT algorithm.
However, the baselineminimax approach consistently underestimates the optimal transport distance.
This is an issue withthe minimax formulation not being able to cover all of the modes, and has been commonly observedin the literature ( e.g. Salimans et al. [2016], Che et al. [2017], Srivastava et al. Moreover,in order to stabilize minimax optimization in practice, Makkuva et al.
[2020] made use of carefuloptimization tuning such as reducing momentum
[Radford et al., 2016].
In contrast, our NCLapproach is a simple optimization problem.
Conclusion
We proposed two constructions for building divergence-free neural networks from an original unconstrained smooth neural network.
This resulted in two practical approaches where either an originalmatrix field or vector field is parameterized.
We found both to be useful in practice; the matrix fieldapproach is more flexible in lower dimensions while the vector field approach allows the additionof extra constraints such as a positive constraint.
We combine these methods with the insight thatthe continuity equation itself is a divergence-free vector field.
This resulted in a method for buildingdeep neural networks that are constrained to always satisfy the continuity equation, which we referto as Neural Conservation Laws.
Currently, these models are difficult to scale due to extensive useof automatic differentiation for computing divergence of neural networks.
It may be possible tocombine our approach with parameterizations that provide cheap divergence computations [Chen and
Duvenaud, 2019].
Nevertheless, the development of efficient batching and automatic differentiationtools is an active research area, so we expect these tools to improve as development progresses.