Pathwise Derivatives for Multivariate Distributions

Martin Jankowiak⇤ Theofanis KaraletsosUber AI Labs, San Francisco, CA

Abstract

We exploit the link between the transportequation and derivatives of expectations toconstruct efficient pathwise gradient estima-tors for multivariate distributions. We focuson two main threads. First, we use null solu-tions of the transport equation to constructadaptive control variates that can be used toconstruct gradient estimators with reducedvariance. Second, we consider the case ofmultivariate mixture distributions. In par-ticular we show how to compute pathwisederivatives for mixtures of multivariate Nor-mal distributions with arbitrary means anddiagonal covariances. We demonstrate in avariety of experiments in the context of varia-tional inference that our gradient estimatorscan outperform other methods, especially inhigh dimensions.

1 Introduction

Stochastic optimization is a major component of manymachine learning algorithms. In some cases—for ex-ample in variational inference—the stochasticity arisesfrom an explicit, parameterized distribution q✓(z). Inthese cases the optimization problem can often be castas maximizing an objective

L = Eq✓(z) [f(z)] (1)

where f(z) is a (differentiable) test function and ✓is a vector of parameters.2 In order to maximize L

we would like to follow gradients r✓L. If z were a⇤Correspondence to: jankowiak@uber.com2Without loss of generality we assume that f(z) has no

explicit dependence on ✓, since gradients of f(z) are readilyhandled by bringing r✓ inside the expectation.

Proceedings of the 22nd International Conference on Ar-tificial Intelligence and Statistics (AISTATS) 2019, Naha,Okinawa, Japan. PMLR: Volume 89. Copyright 2019 bythe author(s).

deterministic function of ✓ we would simply apply thechain rule; for any component ✓↵ of ✓ we have

r✓↵f(z) = rzf ·@z

@✓↵(2)

What is the generalization of Eqn. 2 to the stochasticcase? As observed in [18], we can construct an unbiasedestimator for the gradient of Eqn. 1 provided we cansolve the following partial differential equation (thetransport a.k.a. continuity equation):

@

@✓↵q✓ +rz ·

�q✓v

✓↵�= 0 (3)

Here v✓↵ is a vector field defined on the sample spaceof q✓(z).3 Note that there is a vector field v✓↵ for eachcomponent ✓↵ of ✓. We can then form the followingpathwise gradient estimator:

r✓↵L = Eq✓(z)

⇥rzf · v✓↵

⇤(4)

Pathwise gradient estimators are particularly appeal-ing because, empirically, they generally exhibit lowervariance than alternatives.

That the gradient estimator in Eqn. 4 is unbiased fol-lows directly from the divergence theorem:

r✓↵L =

Zdzf(z)

@q✓(z)

@✓↵= �

Zdzf(z)rz ·

�q✓v

✓↵�

=

Z

Vdzrzf · (q✓v

✓↵)�

I

Sfq✓v

✓· n dS

= limS!1

⇢Z

Vdzrzf · (q✓v

✓↵)�

I

Sfq✓v

✓· n dS

�

=

Zdzq✓(z)rzf · v✓↵ = Eq✓(z)

⇥rzf · v✓↵

⇤

Here we have substituted for @@✓↵

q✓ using Eqn. 3 inthe first line, appealed to the divergence theorem4 inthe second line, and taken the surface to infinity in the

3More explicitly: at each point z in the sample spacethe vector field v✓↵ is specified by D components v✓↵

i ,i = 1, ..., D, where D is the dimension of the sample space.

4More explicitly: we apply the divergence theorem to thevector field fq✓v

✓ and expand the divergence in the volumeintegral using the product rule (cf. Green’s identities).

Pathwise Derivatives for Multivariate Distributions

third line. In the final step we assume that fq✓v✓ issufficiently well-behaved that we can drop the surfaceintegral as the surface goes to infinity.

Thus Eqn. 4 is a generalization of Eqn. 2 to the stochas-tic case, with the velocity field v✓↵ playing the role of@z@✓↵

in the deterministic case. In contrast to the deter-ministic case where @z

@✓↵is uniquely specified, Eqn. 3

generally admits an infinite dimensional space of solu-tions for v✓↵ .

The transport equation encodes an intuitive geometricpicture. As we vary ✓ we move q✓(z) along a curve inthe space of distributions over the sample space. Thiscurve corresponds to a time-varying cloud of particles;in this analogy the set of velocity fields {v✓↵}↵ describesthe (infinitesimal) displacements that particles undergoas ✓ is varied. Because each velocity field obeys thecorresponding transport equation, Eqn. 3, the totalprobability is conserved and the displaced particles aredistributed according to the displaced distribution.

In this work we are interested in the multivariate case,where the solution space to Eqn. 3 is very rich, andwhere alternative gradient estimators tend to sufferfrom large variance, especially in high dimensions. Ourcontributions are two-fold. In Sec. 3 we exploit thethe fact that the transport equation admits a largespace of solutions to construct gradient estimators thatare adapted to the test function f(z). In the contextof the reparameterization trick (see Sec. 2.3) our ap-proach is conceptually analogous to adaptively choosinga reparameterization T from a family of transforma-tions {T�(✏;✓)}, with the difference that the trans-formation is never explicitly constructed; rather it isimplicit as a solution to Eqn. 3. Second, in Sec. 4 weconstruct pathwise gradient estimators for mixtures ofNormal distributions with diagonal covariance matrices.We emphasize that the resulting gradient estimatorsare the only known unbiased pathwise gradient esti-mators for this class of mixture distributions whosecomputational complexity is linear in the dimensionof the sample space. As we show in Sec. 6 our gra-dient estimators often exhibit low variance and goodperformance as compared to alternative methods.

The rest of this paper is organized as follows. In Sec. 2we give an overview of stochastic gradient variationalinference (SGVI) and stochastic gradient estimators.In Sec. 3 we use parameterized families of solutions toEqn. 3 to construct adaptive pathwise gradient esti-mators. In Sec. 4 we present solutions to Eqn. 3 formultivariate mixture distributions. In Sec. 5 we placeour work in the context of recent research. In Sec. 6we validate our proposed techniques with a variety ofexperiments in the context of variational inference. InSec. 7 we discuss directions for future work.

2 Background

2.1 Stochastic Gradient Variational Inference

Given a probabilistic model p(x, z) with observationsx and latent random variables z, variational inferencerecasts posterior inference as an optimization problem.Specifically we define a family of variational distribu-tions q✓(z) parameterized by ✓ and seek to find a valueof ✓ that minimizes the KL divergence between q✓(z)and the (unknown) posterior p(z|x) [19, 29, 15, 45, 31].This is equivalent to maximizing the ELBO, defined as

ELBO = Eq✓(z) [log p(x, z)� log q✓(z)] (5)

This stochastic optimization problem will be the basicsetting for most of our experiments in Sec. 6.

2.2 Score Function Estimator

The score function estimator, also referred to as re-inforce [12, 44, 10], provides a simple and broadlyapplicable recipe for estimating gradients, with thesimplest variant given by

r✓↵L = Eq✓(z) [f(z)r✓↵ log q✓(z)]

Although the score function estimator is very general(e.g. it can be applied to discrete random variables)it typically suffers from high variance, although thiscan be mitigated with the use of variance reductiontechniques such as Rao-Blackwellization [5] and controlvariates [36].

2.3 Reparameterization Trick

The reparameterization trick (RT) is not as broadlyapplicable as the score function estimator, but it gener-ally exhibits lower variance [30, 39, 22, 11, 42, 34]. Itis applicable to continuous random variables whoseprobability density q✓(z) can be reparameterizedsuch that we can rewrite expectations Eq✓(z) [f(z)] asEq0(✏) [f(T (✏;✓))], where q0(✏) is a fixed distributionand T (✏;✓) is a differentiable ✓-dependent transfor-mation. Since the expectation w.r.t. q0(✏) has no ✓dependence, gradients w.r.t. ✓ can be computed bypushing r✓ through the expectation. This reparame-terization can be done for a number of distributions,including for example the Normal distribution. Asnoted in [18], in cases where the reparameterizationtrick is applicable, we have that

v✓↵ =@T (✏;✓)

@✓↵

����✏=T �1(z;✓)

(6)

is a solution to the transport equation, Eqn. 3.

Martin Jankowiak, Theofanis Karaletsos

3 Adaptive Velocity Fields

Whenever we have a solution to Eqn. 3 we can form anunbiased gradient estimator via Eqn. 4. An intriguingpossibility arises if we can construct a parameterizedfamily of solutions v✓↵

� : as we take steps in ✓-space wecan simultaneously take steps in �-space, thus adap-tively choosing the form the ✓↵ gradient estimator takes.This can be understood as an instance of an adaptivecontrol variate [36]; consider the solution v✓↵

� as wellas a fixed reference solution v✓↵

0 :

r✓↵L = Eq✓(z)

hrzf · v✓↵

�

i

= Eq✓(z)

hrzf · v✓↵

0

i+ Eq✓(z)

hrzf ·

⇣v✓↵� � v✓↵

0

⌘i

(7)

The final expectation is identically equal to zero sothe integrand is a control variate. If we choose �appropriately we can reduce the variance of our gradientestimator. In order to specify a complete algorithm weneed to answer two questions: i) how should we adapt�?; and ii) what is an appropriate family of solutionsv✓↵� ? We now address each of these in turn.

3.1 Adapting velocity fields

Whenever we compute a gradient estimate via Eqn. 4we can use the same Ns sample(s) zi ⇠ q✓ to form anestimate of the variance of the gradient estimator:

Var

@LV

@✓

!⌘

X

↵

Var

@LV

@✓↵

!=

1Ns

X

↵

X

zi

h(rzf · v✓↵

� )2i�

X

↵

1Ns

X

zi

hrzf · v✓↵

�

i!2

In direct analogy to the adaptation procedure used in(for example) [37] we can adapt � by following noisygradient estimates of this variance:

r�Var

@LV

@✓

!=

1Ns

X

↵

X

zi

hr�(rzf · v✓↵

� )2i

(8)

In this way � can be adapted to reduce the gradientvariance during the course of optimization. See Algo-rithm 1 for a summary of the complete algorithm in thecase where single-sample gradient estimates are usedat each iteration.

3.2 Parameterizing velocity fields

Assuming we have obtained (at least) one solution v✓0 to

the transport equation Eqn. 3, parameterizing a familyof solutions v✓

� is in principle easy. We simply solvethe null equation, rz ·

�q✓v✓

�= 0, which is obtained

Algorithm 1: Stochastic optimization with AdaptiveVelocity Fields. Note that steps 2 and 4 occur implicitlyfor the reparameterization trick estimator. What setsAdaptive Velocity Fields apart is steps 3 and 5.Initialize ✓, �, and choose step sizes ✏✓, ✏�.for i = 1, 2, ..., Nsteps do

1. Sample zi ⇠ q✓2. Compute g(✓)✓↵i ⌘ rzf(zi) · v

✓↵� for all ↵

(cf. Eqn. 7)3. Compute g(�)i ⌘ r�Var

�@L@✓

�(cf. Eqn. 8)

4. Take a (noisy) gradient step ✓ ! ✓ + ✏✓g(✓)i5. Take a (noisy) gradient step � ! �+ ✏�g(�)i

end

from the transport equation by dropping the sourceterm @

@✓ q✓. Then for a given family of solutions to thenull equation, v✓

�, we get a solution v✓� to the transport

equation by simple addition, i.e. v✓� = v✓

0 + v✓�. At

first glance solving the null equation appears to beeasy, since any divergence-free vector field immediatelyyields a solution:

rz ·w� = 0 and v✓� =

w�

q✓) rz ·

�q✓v

✓�

�= 0

However, in order to construct general purpose gradientestimators we need to impose appropriate boundaryconditions on the velocity field; for more on this subtletywe refer the reader to the supplementary materials.

In order to demonstrate how to construct suitable nullsolutions we henceforth focus on a particular familyof distributions q✓(z), namely elliptical distributions.Elliptical distributions are probability distributionswhose density is of the form q(z) / g(zT⌃�1z) forsome scalar density g(·) and positive definite symmetricmatrix ⌃. They include, for example, the multivariateNormal distribution and the multivariate t-distribution.For concreteness, in the following we focus exclusivelyon the multivariate Normal distribution. We stress,however, that the resulting adaptive gradient estima-tors are applicable to any elliptical distribution. Werefer to the supplementary materials for a descriptionof the specific case of the multivariate t-distribution.

3.3 Null Solutions for the MultivariateNormal Distribution

Consider the multivariate Normal distribution in D di-mensions parameterized via a mean µ and Cholesky fac-tor L. We would like to compute derivatives w.r.t. Lab.While the reparameterization trick is applicable in thiscase, we can potentially get lower variance gradientsby using Adaptive Velocity Fields. As we show inthe supplementary materials, a simple solution to the

Pathwise Derivatives for Multivariate Distributions

corresponding null transport equation is given by

vLabA = LAabL�1

(z � µ) (9)

where � ⌘ {Aab} is an arbitrary collection of antisym-

metric D ⇥D matrices (one for each a, b). This is justa linear vector field and so it is easy to manipulate and(relatively) cheap to compute. Geometrically, for eachentry of L, this null solution corresponds to an infinites-imal rotation in whitened coordinates z = L�1

(z �µ).Note that vLab

A is added to the reference solution vLab0 ,

and this is not equivalent to rotating vLab0 .

Since a, b runs over D(D+1)2 pairs of indices and each

Aab has D(D�1)2 free parameters, this space of solutions

is quite large. Since we will be adapting Aab via noisygradient estimates—and because we would like to limitthe computational cost—we choose to parameterize asmaller subspace of solutions. Empirically we find thatthe following parameterization works well:

Aabjk =

MX

`=1

B`aC`b (�aj�bk � �ak�bj) (10)

Here each M ⇥D matrix B and C is arbitrary and thehyperparameter M allows us to trade off the flexibilityof our family of null solutions with computational cost(typically M ⌧ D).

The computational complexity of using Adaptive Ve-locity Field gradients with this class of parameterizedvelocity fields (including the A update equations) isO(D3

+MD2) per gradient step. This should be com-

pared to the O(D2) cost of the reparameterization trick

gradient and the O(D3) cost of the OMT gradient (see

Sec. 5 for a discussion of the latter).5 We explore theperformance of the resulting Adaptive Velocity Field(AVF) gradient estimator in Sec. 6.1.

4 Pathwise Derivatives for

Multivariate Mixture Distributions

In this section we use the transport equation Eqn. 3to construct pathwise gradient estimators for mixturedistributions. Note that this section is logically indepen-dent from the previous section, although—as we discussbriefly in the supplementary materials—Adaptive Ve-locity Fields can be readily combined with the gradientestimators discussed here.

5Note, however, that the computational complexity ofthe AVF gradient is somewhat misleading in that the O(D3)term arises from matrix multiplications, which tend tobe quite fast. By contrast the OMT gradient estimatorinvolves a singular value decomposition, which tends to besubstantially more expensive than a matrix multiplicationon modern hardware.

Consider a mixture of K multivariate distributions inD dimensions:

q✓(z) =KX

j=1

⇡jq✓j (z) (11)

In order to compute pathwise derivatives of z ⇠ q✓(z)we need to solve the transport equation w.r.t. the mix-ture weights ⇡j as well as the component parameters ✓j .For the derivatives w.r.t. ✓j we can just repurpose thevelocity fields of the individual components. That is, ifv✓j

single is a solution of the single-component transport

equation�i.e.

@q✓j@✓j

+r · (q✓jv✓j

single) = 0�, then

v✓j =⇡jq✓j

q✓v✓j

single (12)

is a solution of the multi-component transport equa-tion.6 See the supplementary materials for the (short)derivation. For example we can readily computeEqn. 12 if each component distribution is reparam-eterizable. Note that a typical gradient estimator for,say, a mixture of Normal distributions first samples thediscrete component k 2 [1, ...,K] and then samples zconditioned on k. This results in a pathwise derivativefor z that only depends on µk and �k. By contrast thegradient computed via Eqn. 12 will result in a gradientfor all K component parameters for every sample z.We explore this difference experimentally in Sec. 6.2.2.

4.1 Mixture weight derivatives

Unfortunately, solving the transport equation for themixture weights is in general much more difficult, sincethe desired velocity fields need to coordinate masstransport among all K components.7 We now describea formal construction for solving the transport equationfor the mixture weights. For each pair of componentdistributions j, k consider the transport equation

q✓j � q✓k +rz ·�q✓v

jk�= 0 (13)

Intuitively, the velocity field vjk moves mass from com-ponent k to component j. We can superimpose thesesolutions to form

v`j = ⇡j

X

k 6=j

⇡kvjk (14)

As we show in the supplementary materials, the ve-locity field v`j yields an estimator for the derivative

6We note that the form of Eqn. 12 implies that wecan easily introduce Adaptive Velocity Fields for each indi-vidual component (although the cost of doing so may beprohibitively high for a large number of components).

7For example, we expect this to be particularly difficultif some of the component distributions belong to differentfamilies of distributions, e.g. a mixture of Normal distribu-tions and Wishart distributions.

Martin Jankowiak, Theofanis Karaletsos

w.r.t. the softmax logit `j that corresponds to compo-nent j.8 Thus provided we can find solutions vjk toEqn. 13 for each pair of components j, k we can forma pathwise gradient estimator for the mixture weights.We now consider one particular situation where thisrecipe can be carried out. We leave a discussion of othersolutions—including a solution for a mixture of Normaldistributions with arbitrary diagonal covariances—tothe supplementary materials.

4.2 Mixture of Multivariate Normals withShared Diagonal Covariance

Here each component distribution is specified byq✓j (z) = N (z|µj ,�), i.e. each component distribu-tion has its own mean vector µj but all K componentsshare the same (diagonal) square root covariance �.9We find that a solution to Eqn. 13 is given by

vjk=

⇣�(zjkk �µjk

k )� �(zjkk +µkjk )

⌘�(||zjk

? ||2)

(2⇡)D/2�1q✓QD

i=1 �i

�µjk

(15)where �µjk indicates matrix multiplication bydiag(�), z ⌘ z � ��1, and where the quantities{µj , µjk, µjk

k , zjkk , zjk? } are defined in the supplemen-

tary materials. Here �(·) is the CDF of the unit Normaldistribution and �(·) is the corresponding probabilitydensity. Geometrically, the velocity field in Eqn. 15moves mass along lines parallel to µj � µk. Note thatthe gradient estimator resulting from Eqn. 15 has costO(DK2

).10 We explore the performance of the result-ing pathwise gradient estimators in Sec. 6.2.

5 Related Work

There is a large body of work on constructing gradientestimators with reduced variance, much of which can beunderstood in terms of control variates [36]: for exam-ple, [27] constructs neural baselines for score functiongradients; [40] discuss gradient estimators for stochas-tic computation graphs and their Rao-Blackwellization;and [43, 13] construct adaptive control variates for dis-crete random variables. Another example of this lineof work is [25], where the authors construct controlvariates that are applicable when q✓(z) is a diagonal

Normal distribution and that rely on Taylor expansions

8That is with respect to `j where ⇡j = e`j/P

k e`k .

9To keep the equations compact we consider a diagonalcovariance matrix. The resulting velocity field is readilytransformed to a rotated coordinate system in which the(single shared) covariance matrix is arbitrary.

10As detailed in the supplementary materials, the path-wise gradient estimators for the familes of mixture distribu-tions we consider all have cost O(DK) or O(DK2).

of the test function f(z).11 In contrast our AVF gradi-ent estimator for the multivariate Normal has abouthalf the computational cost, since it does not rely onsecond order derivatives. Another line of work con-structs partially reparameterized gradient estimatorsfor cases where the reparameterization trick is diffi-cult to apply [38, 28]. In [14], the author constructspathwise gradient estimators for mixtures of diagonalNormal distributions. Unfortunately, the resulting esti-mator is expensive, relying on a recursive computationthat scales with the dimension of the sample space. An-other approach to mixture distributions is taken in [7],where the authors use quadrature-like constructionsto form continuous distributions that are reparame-terizable. Ref. [35] considers a gradient estimator formixture distributions in which all K components aresummed out. Our work is closest to [18], which usesthe transport equation to derive an ‘OMT’ gradient es-timator for the multivariate Normal distribution whosevelocity field is optimal in the sense of optimal trans-port. This gradient estimator can exhibit low variancebut has O(D3

) computational complexity due to an(expensive) singular value decomposition. In additionthe OMT gradient estimator is not adapted to the testfunction at hand; rather it has one fixed form.

As this manuscript was being completed, we becameaware of reference [9], which has some overlap with thiswork. In particular, [9] describes an interesting recipefor computing pathwise gradients that can be appliedto mixture distributions. A nice feature of this recipe(also true of [14]), is that it can be applied to mixturesof Normal distributions with arbitrary covariances. Adisadvantage is that it can be expensive, exhibiting aO(D2

) computational cost even for a mixture of Normaldistributions with diagonal covariances. In contrast thegradient estimators constructed in Sec. 4 by solvingthe transport equation are linear in D. We provide acomparison to their approach in Sec. 6.2.1.

6 Experiments

We conduct a series of experiments using synthetic andreal world data to validate the performance of the gra-dient estimators introduced in Sec. 3 & 4. Throughoutwe use single sample estimators. See Sec. E in thesupplementary materials for details on experimentalsetups. Open source implementations of all our gradi-ent estimators are available at https://git.io/fhbqH.All variational inference experiments were implementedusing the Pyro probabilistic programming language [2].

11Also, in their approach variance reduction for scale pa-rameter gradients r� necessitates a multi-sample estimator(at least for high-dimensional models where computing thediagonal of the Hessian is expensive).

Pathwise Derivatives for Multivariate Distributions

0.4

0.6

0.8

1.0VDrLDnFe RDtLR

AdDptLve VeORFLty )LeOdV (0 1)

FRVLneTuDdrDtLFTuDrtLF

0.0 0.2 0.4 0.6 0.8 1.00DgnLtude Rf 2ff-DLDgRnDO Rf L

0.0

0.2

0.4

0.6

VDrLDnFe RDtLR

20T VeORFLty )LeOdVFRVLneTuDdrDtLFTuDrtLF

(a) We compare the OMT and AVF gradient estimators forthe multivariate Normal distribution to the RT estimatorfor three test functions. The horizontal axis controls themagnitude of the off-diagonal elements of the Cholesky factorL. The vertical axis depicts the ratio of the mean varianceof the given estimator to that of the RT estimator for theoff-diagonal elements of L.

101 102

DimenViRn10−4

10−3

10−2

10−1

100

VDriD

nce

RDtiR

ZerR0eDnGS0GS0DiDg1RrmDlVDiDg1RrmDlVShDreGCRvDriDnce

(b) We compare the variance of pathwise gradient estima-tors for various mixture distributions with K = 10 to thecorresponding score function estimator. The vertical axis de-picts the ratio of the logit gradient variance of the pathwiseestimator to that of the score function estimator for the testfunction f(z) = ||z||2.

5 10 15 20 25 30 35 401umber Rf 0ixture CRmpRnentV K

10−1

100

101

Variance 5atiR

σ 0.75 σ 1.0 σ 2.0 σ 4.0

(c) We compare the variance of our pathwise gradient estima-tor for mixtures of diagonal Normal distributions to the onedescribed in ref. [9]. Here D = 2, � controls the covarianceof each component distribution, and the vertical axis is theratio of variances (with our estimator in the denominator).See supplementary materials for details.

Figure 1: We compare the gradient variances for dif-ferent gradient estimators and synthetic test functions.See Sec. 6.1.1 and Sec. 6.2.1 for details.

100 200 300 400 500IWeraWLRn

−3

−2

−1

(LB2 ×10

−3

57207AV) (0 5)AV) (0 1)

10 20 30 40 50 60 70WaOO CORFk 7Lme (V)

−3

−2

−1

(LB2 ×10

−3

57207AV) (0 5)AV) (0 1)

Figure 2: We compare the performance of variouspathwise gradient estimators on the GP regression taskin Sec 6.1.2. The AVF gradient estimator with M = 5

is the clear winner both in terms of wall clock time andthe final ELBO achieved.

6.1 Adaptive Velocity Fields

6.1.1 Synthetic test functions

In Fig. 1a we use synthetic test functions to illustratethe variance reduction that is achieved with AVF gra-dients for the multivariate Normal distribution as com-pared to the reparameterization trick and OMT gradi-ents from [18]. The dimension is D = 50; the resultsare qualitatively similar for different dimensions. Notethat, as demonstrated here, whether AVF outperformsOMT is problem dependent.

6.1.2 Gaussian Process Regression

We investigate the performance of AVF gradients forthe multivariate Normal distribution in the context ofa Gaussian Process regression task reproduced from[18]. We model the Mauna Loa CO2 data from [20]considered in [32]. We fit the GP using a single-sampleMonte Carlo ELBO gradient estimator and all N = 468

data points (so D = 468). Both the OMT and AVFgradient estimators achieve higher ELBOs than the RTestimator (see Fig. 2); however, the OMT estimatordoes so at substantially increased computational cost(⇠1.9x), while the AVF estimator for M = 1 (M = 5)requires only ⇠6% (⇠11%) more time per iteration. Byiteration 250 the AVF gradient estimator with M = 5

has attained the same ELBO that the RT estimatorattains at iteration 500.

6.2 Mixture Distributions

6.2.1 Synthetic test function

In Fig. 1b we depict the variance reduction of pathwisegradient estimators for various mixture distributions

Martin Jankowiak, Theofanis Karaletsos

compared to the corresponding score function estima-tor. We find that the typical variance reduction hasmagnitude O(D). The results are qualitatively simi-lar for different numbers of components K. See thesupplementary materials for details on the differentfamilies of mixture distributions. In Fig. 1c we com-pare the gradient estimator for mixtures of Normaldistributions described in Sec. 4.2 to the estimatordescribed in [9]. We find that our estimator exhibitsbetter performance when the mixture components havemore overlap (� � 0), while their estimator exhibitsbetter performance in the opposite regime (� ! 0).Since their estimator has a cost that is quadratic inthe dimension D, while our estimator has a cost thatis linear in the dimension D, which estimator is to bepreferred will be problem specific.

6.2.2 Baseball Experiment

We consider a model for repeated binary trial datausing the data in [8] and the modeling setup in [1]with partial pooling. The model has two global la-tent variables and 18 local latent variables so that theposterior is 20-dimensional. While mean field SGVIgives reasonable results for this model, it is not ableto capture the detailed structure of the exact posterioras computed with the NUTS HMC implementationin Stan [16, 4]. To go beyond mean field we considervariational distributions that are mixtures of K diago-nal Normal distributions.12 Adding more componentsgives a better approximation to the exact posterior,see Fig. 3a. Furthermore, using pathwise derivativesin the ELBO gradient estimator leads to faster con-vergence, see Fig. 3b. Here the hybrid estimator usesscore functions gradients for the mixture logits butimplements Eqn. 12 for gradients w.r.t. the componentparameters. Since there is significant overlap amongthe mixture components, Eqn. 12 leads to substantialvariance reduction.

6.2.3 Continuous state space model

We consider a non-linear state space model with ob-servations xt and random variables zt at each timestep (both two dimensional). Since the likelihood is ofthe form p(xt|zt) = N (xt|µ(zt),�x12) with µ(zt) =(z21t, 2z2t), the posterior over z1:T is highly multi-modal.We use SGVI to approximate the posterior over z1:T .To better capture the multi-modality we use a varia-tional family that is a product of two-component mix-ture distributions of the form q(zt|zt�1,xt:T ), one ateach time step. As can be seen from Fig. 4, the varia-tional family makes use of the mixture distributions atits disposal to model the multi-modal structure of theposterior, something that a variational family with a

12Cf. the experiment in [26]

(a) Two-dimensional cross-sections of three approximate pos-teriors; each cross-section includes one global latent variable,log(), and one local latent variable, logit(✓0). The whitedensity contours correspond to the different variational ap-proximations, and the background depicts the approximateposterior computed with HMC. The quality of the vari-ational approximation improves as we add more mixturecomponents.

(b) ELBO training curves for three gradient estimators forK = 20. The lower variance of the pathwise and hybridgradient estimators speeds up learning. The dotted linesdenote mean ELBOs over 50 runs and the bands indicate1-� standard deviations. Note that the figure would be qual-itatively similar if the ELBO were plotted against wall clocktime, since the estimators have comparable computationalcost (within 10%).

Figure 3: Variational approximations and ELBO train-ing curves for the model in Sec. 6.2.2.

unimodal q(zt|·) struggles to do. In the next sectionwe explore the use of mixture distributions in a muchricher time series setup.

6.2.4 Deep Markov Model

We consider a variant of the Deep Markov Model(DMM) setup introduced in [23] on a polyphonic musicdataset. We fix the form of the model and vary thevariational family and gradient estimator used. Weconsider a variational family which includes a mixturedistribution at each time step. In particular each fac-tor q(zt|zt�1,xt:T ) in the variational distribution isa mixture of K diagonal Normals. We compare theperformance of the pathwise gradient estimator intro-duced in Sec. 4 to two variants of the Gumbel Softmaxgradient estimator [17, 24], see Table 1. For a descrip-tion of these two gradient estimators, which use theGumbel Softmax trick to deal with mixture weight gra-dients, see Sec. E in the supplementary materials. Notethat both Gumbel Softmax estimators are biased, whilethe pathwise estimator is unbiased. We find that theGumbel Softmax estimators—which induce O(2 � 8)

times higher gradient variance for the mixture logits

Pathwise Derivatives for Multivariate Distributions

(a) K = 1 (no mixture)

(b) K = 2 (mixture)

Figure 4: Approximate marginal posteriors for the toystate space model in Sec. 6.2.3 for a randomly chosentest sequence for t = 4, 5, 6. The mixture distributionsthat serve as components of the variational family onthe bottom help capture the multi-modal structure ofthe posterior.

as compared to the pathwise estimator—are unable toachieve competitive test ELBOs.13

Note that while the numerical differences in Table 1are not large (⇠ 0.05 nats), they are comparable tothe quantitative improvement that one can achievein this context when equipping the variational familywith normalizing flows [41, 33]; for this particular taskthere appears to be little headroom above what can beattained with a (unimodal) Normal variational family.

Gradient EstimatorVariational Family Pathwise GS-Soft GS-Hard

K=2 -6.84 -6.89 -6.88K=3 -6.82 -6.87 -6.85

Table 1: Test ELBOs for the DMM in Sec. 6.2.4. Higheris better. For comparison: we achieve a test ELBO of-6.90 for a variational family in which each q(zt|·) is adiagonal Normal.

6.2.5 VAE

We conduct a simple experiment with a VAE [22, 34]trained on MNIST. We fix the form of the model andvary the gradient estimator used. The variational fam-ily is of the form q(z|x), where q(z|·) is a mixtureof K 2 [3, 4, 5] diagonal Normals. We find that on

13We also found the Gumbel Softmax estimators to beless numerically stable, which prevented us from usingmore aggressive KL annealing schedules. This may havecontributed to the performance gap in Table 1.

this task our gradient estimator achieves similar per-formance to the two Gumbel Softmax estimators andthe score function estimator, see Table 2. This is nottoo surprising, since VAE posteriors tend to be quitesharp for any given x so that gradient variance is notthe bottleneck for learning.

Gradient EstimatorVar. Family Pathwise GS-Soft GS-Hard SF

K=3 -95.88 -95.95 -95.96 -95.98K=4 -95.89 -95.84 -95.88 -95.99K=5 -96.04 -95.94 -95.89 -95.98

Table 2: Test ELBOs for the VAE in Sec. 6.2.5. Higheris better.

7 Discussion

We have seen that the link between the transport equa-tion and derivatives of expectations offers a powerfultool for constructing pathwise gradient estimators. Inthe case of variational inference, we emphasize that—inaddition to providing gradients with reduced variance—pathwise gradient estimators can enable quick iterationover complex models and variational families, sincethere is no need to reason about ELBO gradients. Themodeler need only construct a Monte Carlo estimatorof the ELBO; automatic differentiation will handle therest. Our hope is that an expanding toolkit of pathwisegradient estimators can contribute to the developmentof innovative applications of variational inference.

One limitation of our approach is that solving thetransport equation for arbitrary families of distribu-tions presents a challenge. Of course, this limitationalso applies to the reparameterization trick. Whilewe acknowledge this limitation, we note that since:i) the Normal distribution is ubiquitous in machinelearning; and ii) mixtures of Normals are universalapproximators for continuous distributions, improvingthe gradient estimators that can be used in these twocases is already of considerable value.

Given the richness of the transport equation and thegreat variety of multivariate distributions, there areseveral exciting possibilities for future research. Itwould be of interest to identify more cases where wecan construct null solutions to the transport equation,e.g. generalizing Eqn. 9 to higher orders. It could alsobe fruitful to consider other ways of adapting veloc-ity fields. In some cases this could lead to improvedvariance reduction, especially in the case of mixturedistributions, where the optimal transport setup in[6] might provide a path to constructing alternativegradient estimators.

Martin Jankowiak, Theofanis Karaletsos

References

[1] Stan modeling language users guide andreference manual, version 2.17.0. http:// mc-stan. org/ users/ documentation/case-studies/ pool-binary-trials. html ,2017.

[2] E. Bingham, J. P. Chen, M. Jankowiak, F. Ober-meyer, N. Pradhan, T. Karaletsos, R. Singh, P. Sz-erlip, P. Horsfall, and N. D. Goodman. Pyro: Deepuniversal probabilistic programming. Journal of

Machine Learning Research, 2018.

[3] N. Boulanger-Lewandowski, Y. Bengio, and P. Vin-cent. Modeling temporal dependencies in high-dimensional sequences: Application to polyphonicmusic generation and transcription. arXiv preprint

arXiv:1206.6392, 2012.

[4] B. Carpenter, A. Gelman, M. D. Hoffman, D. Lee,B. Goodrich, M. Betancourt, M. Brubaker, J. Guo,P. Li, and A. Riddell. Stan: A probabilistic pro-gramming language. Journal of statistical software,76(1), 2017.

[5] G. Casella and C. P. Robert. Rao-blackwellisationof sampling schemes. Biometrika, 83(1):81–94,1996.

[6] Y. Chen, T. T. Georgiou, and A. Tannenbaum.Optimal transport for gaussian mixture models.arXiv preprint arXiv:1710.07876, 2017.

[7] J. Dillon and I. Langmore. Quadrature compound:An approximating family of distributions. arXiv

preprint arXiv:1801.03080, 2018.

[8] B. Efron and C. Morris. Data analysis using stein’sestimator and its generalizations. Journal of the

American Statistical Association, 70(350):311–319,1975.

[9] M. Figurnov, S. Mohamed, and A. Mnih. Im-plicit reparameterization gradients. arXiv preprint

arXiv:1805.08498, 2018.

[10] M. C. Fu. Gradient estimation. Handbooks in op-

erations research and management science, 13:575–616, 2006.

[11] P. Glasserman. Monte Carlo methods in finan-

cial engineering, volume 53. Springer Science &Business Media, 2013.

[12] P. W. Glynn. Likelihood ratio gradient estimationfor stochastic systems. Communications of the

ACM, 33(10):75–84, 1990.

[13] W. Grathwohl, D. Choi, Y. Wu, G. Roeder, andD. Duvenaud. Backpropagation through the void:Optimizing control variates for black-box gradi-ent estimation. arXiv preprint arXiv:1711.00123,2017.

[14] A. Graves. Stochastic backpropagation throughmixture density distributions. arXiv preprint

arXiv:1607.05690, 2016.

[15] M. D. Hoffman, D. M. Blei, C. Wang, and J. Pais-ley. Stochastic variational inference. The Journal

of Machine Learning Research, 14(1):1303–1347,2013.

[16] M. D. Hoffman and A. Gelman. The no-u-turnsampler: adaptively setting path lengths in hamil-tonian monte carlo. Journal of Machine Learning

Research, 15(1):1593–1623, 2014.

[17] E. Jang, S. Gu, and B. Poole. Categorical reparam-eterization with gumbel-softmax. arXiv preprint

arXiv:1611.01144, 2016.

[18] M. Jankowiak and F. Obermeyer. Pathwise deriva-tives beyond the reparameterization trick. In Pro-

ceedings of the 35th International Conference on

Machine Learning, pages 2235–2244, 2018.

[19] M. I. Jordan, Z. Ghahramani, T. S. Jaakkola,and L. K. Saul. An introduction to variationalmethods for graphical models. Machine learning,37(2):183–233, 1999.

[20] C. D. Keeling and T. P. Whorf. Atmospheric co2concentrations derived from flask air samples atsites in the sio network. Trends: a compendium

of data on Global Change, 2004.

[21] D. P. Kingma and J. Ba. Adam: A methodfor stochastic optimization. arXiv preprint

arXiv:1412.6980, 2014.

[22] D. P. Kingma and M. Welling. Auto-encodingvariational bayes. arXiv preprint arXiv:1312.6114,2013.

[23] R. G. Krishnan, U. Shalit, and D. Sontag. Struc-tured inference networks for nonlinear state spacemodels. In AAAI, pages 2101–2109, 2017.

[24] C. J. Maddison, A. Mnih, and Y. W. Teh. Theconcrete distribution: A continuous relaxationof discrete random variables. arXiv preprint

arXiv:1611.00712, 2016.

[25] A. Miller, N. Foti, A. D’Amour, and R. P. Adams.Reducing reparameterization gradient variance. InAdvances in Neural Information Processing Sys-

tems, pages 3711–3721, 2017.

Pathwise Derivatives for Multivariate Distributions

[26] A. C. Miller, N. Foti, and R. P. Adams. Variationalboosting: Iteratively refining posterior approxima-tions. arXiv preprint arXiv:1611.06585, 2016.

[27] A. Mnih and K. Gregor. Neural variational in-ference and learning in belief networks. arXiv

preprint arXiv:1402.0030, 2014.

[28] C. Naesseth, F. Ruiz, S. Linderman, and D. Blei.Reparameterization gradients through acceptance-rejection sampling algorithms. In Artificial Intel-

ligence and Statistics, pages 489–498, 2017.

[29] J. Paisley, D. Blei, and M. Jordan. Variationalbayesian inference with stochastic search. arXiv

preprint arXiv:1206.6430, 2012.

[30] R. Price. A useful theorem for nonlinear deviceshaving gaussian inputs. IRE Transactions on In-

formation Theory, 4(2):69–72, 1958.

[31] R. Ranganath, S. Gerrish, and D. Blei. Black boxvariational inference. In Artificial Intelligence and

Statistics, pages 814–822, 2014.

[32] C. E. Rasmussen. Gaussian processes in machinelearning. In Advanced lectures on machine learning,pages 63–71. Springer, 2004.

[33] D. J. Rezende and S. Mohamed. Variational in-ference with normalizing flows. arXiv preprint

arXiv:1505.05770, 2015.

[34] D. J. Rezende, S. Mohamed, and D. Wierstra.Stochastic backpropagation and approximate in-ference in deep generative models. arXiv preprint

arXiv:1401.4082, 2014.

[35] G. Roeder, Y. Wu, and D. K. Duvenaud. Stick-ing the landing: Simple, lower-variance gradientestimators for variational inference. In Advances

in Neural Information Processing Systems, pages6925–6934, 2017.

[36] S. M. Ross. Simulation. Academic Press, SanDiego, 2006.

[37] F. J. Ruiz, M. K. Titsias, and D. M. Blei. Overdis-persed black-box variational inference. arXiv

preprint arXiv:1603.01140, 2016.

[38] F. R. Ruiz, M. T. R. AUEB, and D. Blei. Thegeneralized reparameterization gradient. In Ad-

vances in Neural Information Processing Systems,pages 460–468, 2016.

[39] T. Salimans, D. A. Knowles, et al. Fixed-form vari-ational posterior approximation through stochasticlinear regression. Bayesian Analysis, 8(4):837–882,2013.

[40] J. Schulman, N. Heess, T. Weber, and P. Abbeel.Gradient estimation using stochastic computationgraphs. In Advances in Neural Information Pro-

cessing Systems, pages 3528–3536, 2015.

[41] E. Tabak and C. V. Turner. A family of non-parametric density estimation algorithms. Com-

munications on Pure and Applied Mathematics,66(2):145–164, 2013.

[42] M. Titsias and M. Lázaro-Gredilla. Doublystochastic variational bayes for non-conjugate in-ference. In International Conference on Machine

Learning, pages 1971–1979, 2014.

[43] G. Tucker, A. Mnih, C. J. Maddison, J. Lawson,and J. Sohl-Dickstein. Rebar: Low-variance, unbi-ased gradient estimates for discrete latent variablemodels. In Advances in Neural Information Pro-

cessing Systems, pages 2624–2633, 2017.

[44] R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforce-ment learning. Machine learning, 8(3-4):229–256,1992.

[45] D. Wingate and T. Weber. Automated variationalinference in probabilistic programming. arXiv

preprint arXiv:1301.1299, 2013.