Fast and Simple Natural-Gradient Variational Inferencewith Mixture of Exponential-family Approximations

Wu Lin 1 Mohammad Emtiyaz Khan 2 Mark Schmidt 1

Abstract
Natural-gradient methods enable fast and simple
algorithms for variational inference, but due tocomputational difficulties, their use is mostly limited to minimal exponential-family (EF) approximations. In this paper, we extend their application to estimate structured approximations suchas mixtures of EF distributions. Such approximations can fit complex, multimodal posterior
distributions and are generally more accurate thanunimodal EF approximations. By using a minimal
conditional-EF representation of such approximations, we derive simple natural-gradient updates.
Our empirical results demonstrate a faster convergence of our natural-gradient method compared toblack-box gradient-based methods. Our work expands the scope of natural gradients for Bayesianinference and makes them more widely applicablethan before.

1. Introduction
Variational Inference (VI) provides a cheap and quick approximation to the posterior distribution, and is now widelyused in many areas of machine learning (Kingma & Welling,2013; Furmston & Barber, 2010; Wainwright & Jordan,
2008; Hensman et al., 2013; Nguyen et al., 2017). In recent
years, many natural-gradient methods have been proposed
for VI (Sato, 2001; Honkela et al., 2007; 2011; Hensman
et al., 2012; Hoffman et al., 2013; Khan & Lin, 2017). These
works have shown that, for specific types of models and approximations, natural-gradient methods can result in simple
updates which converge faster than gradient-based methods.
For example, stochastic variational inference (SVI) (Hoffman et al., 2013) is a popular natural-gradient method for
conjugate exponential-family models. Unfortunately, the1

University of British Columbia, Vancouver, Canada. 2 RIKEN
Center for Advanced Intelligence Project, Tokyo, Japan. Correspondence to: Wu Lin <wlin2018@cs.ubc.ca>.
Proceedings of the 36 th International Conference on Machine
Learning, Long Beach, California, PMLR 97, 2019. Copyright2019 by the author(s).

Mixture of EF
Curved EF
MinimalConditionally
EF (MCEF)

Minimal EF

Examples of MCEFs
Finite Mixture of Gaussians
Birnbaum-Saunders Distribution
Multivariate Student’s T-distribution
Multivariate Normal Inverse GaussianMultivariate Skew-GaussianMultivariate Exponentially Modified Gaussian
Low-rank Covariance Gaussian

Figure 1. We derive simple natural-gradient updates for approximations with a minimal-conditional EF (MCEF) representation. Suchapproximations include all minimal (normalized) EF distributions,and some curved and mixture of EF distributions.

simplicity of natural-gradient updates is currently limited
to VI with minimal exponential-family (EF) approximations. For such approximations, we can efficiently computenatural-gradients in the natural-parameter space without
explicitly computing the Fisher information matrix (FIM)(Khan & Nielsen, 2018). Unfortunately, this property does
not extend to many other approximations such as mixturesof EF distributions. Such structured approximations are
more appropriate for complex and multi-modal posterior
distributions, giving a more accurate fit than minimal EF
distributions. However, computation of natural-gradients ischallenging for them.
In this paper, we propose a simple new natural-gradientmethod for VI with structured approximations. We define aclass of distributions which take a minimal conditional-EF(MCEF) form. This includes many members of mixture
and curved EF distributions (see Fig. 1). Using the MCEF
representation and an expectation parameterization associated with it, we derive simple natural-gradient updates.
We show examples on a variety of models where simple
natural-gradient updates can be used to estimate flexible
and accurate structured approximations. Our empirical results show faster convergence of our method compared togradient descent for VI. Our work extends the simplicity ofnatural-gradient methods making them more widely applicable than before, while maintaining their fast convergence.
1.1. Related Works
Existing work on natural-gradient VI to obtain EF approximations all assume a minimal representation to ensure invertiblility of the FIM. Such methods show fast convergence,Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

result in simple updates, and lead to a straightforward implementation for many types of models. They have been used
for conditionally-conjugate EF models (Sato, 2001; Hoffman et al., 2013) and non-conjugate models (Khan & Lin,2017), including deep neural networks (Khan et al., 2018;
Zhang et al., 2018; Mishkin et al., 2018) and Gaussian processes (Khan et al., 2016; Salimbeni et al., 2018). Our work
presents fast and simple updates for approximations that are
beyond the reach of these works.
For structured approximations, existing works employ
stochastic-gradient methods (Salimans & Knowles, 2013;
Hoffman & Blei, 2015; Ranganath et al., 2016; Titsias &Ruiz, 2018; Yin & Zhou, 2018). Such methods are widely
applicable, but not as widely used as their mean-field counterparts. This is because they are computationally expensive
and slow to converge. Our work attempts to improve these
aspects for a flexible class of approximations.

2. Natural-Gradient Variational Inference
We begin with a description of natural-gradient descent forvariational inference in probabilistic models. Given a probabilistic model p(D, z) to model data D using latent vectorz, the goal of Bayesian inference is to compute the posterior
distribution: p(z|D) = p(D, z)/p(D). This requires
comR
putation of the marginal likelihood p(D) = p(D, z)dz,
which is a high dimensional integral and difficult to compute.
VI simplifies this problem by approximating the posterior
distribution p(z|D) by another distribution whose normalizing constant is easier to compute. For example, a common
choice to approximate the posterior is to use a regular1
exponential-family (EF) approximation,q(z|λz ) := hz (z) exp [hφz (z), λz i − Az (λz )] ,(1)

where q denotes the approximating distribution, φz (z) are
the sufficient statistics, hz (z) is the base measure, andλz ∈ Ω is the natural parameter with Ω being the set of valid
natural-parameters (the set of λz where the log-partition
function Az (λz ) is finite) and h·, ·i denotes an inner product.2 Such parametrized approximations can be estimatedby maximizing the variational lower bound:L(λz ) := Eq [log p(D, z) − log q(z|λz )] ,(2)

which can be solved by gradient descent, as shown below:
GD :λz ← λz + α∇λz L(λz ),(3)where ∇ denotes the gradient and α > 0 is a scalar learning
rate. The GD algorithm is simple and convenient to implement by using modern automatic-differentiation methods
and the reparameterization trick (Ranganath et al., 2014;
1
2

An EF is called regular when Ω is an open set.
We assume Az (λz ) can be efficiently computed.

Titsias & Lázaro-Gredilla, 2014). Unfortunately, such firstorder methods can show suboptimal rates of convergenceand be slow in practice.
An alternative approach is to use natural-gradient descent
which exploits the information geometry of q to speed-up
convergence. Assuming that the Fisher information matrix(FIM) of q(z|λz ), denoted by Fz (λz ), is positive-definite
for all λz ∈ Ω, the natural-gradient descent (NGD) for VI
in the natural-parameter space is given as follows:
NGD :

−1

λz ← λz + β [Fz (λz )]∇λz L(λz ),(4)

The preconditioning of the gradients by the FIM leads to aproper scaling of the gradient in each dimension, and takes
into account dependencies between variables. This often
leads to faster convergence, particularly when the FIM is
well-conditioned.
A naive implementation of (4) would require computation and inversion of the FIM. However, for specific typesof models and approximations, NGVI could be simpler
to compute than GD. This is true for mean-field VI inconjugate-exponential family models (Hoffman et al., 2013)
as well as Bayesian neural networks (Khan et al., 2018;
Zhang et al., 2018). This computational efficiency is a
result of a simple NGD update to estimate EF approximations, as shown by Khan & Nielsen (2018). For EFs, we
can use the expectation-parameter, defined as the function
mz (λz ) := Eq [φz (z)] from Ω → M, to compute naturalgradients. This is possible because of the following relation,
∇λz L = ∇λz mTz ∇mz L = [Fz (λz )] ∇mz L (5)where the first equality is obtained by applying the chain ruleand second equality is obtained by noting that ∇λz mTz =∇2λz A(λz ) = Fz (λz ). When the FIM is invertible, we geta simple update for NGD,λz ← λz + β ∇mz L(λz ),(6)where the gradient is computed with respect to mz . Whenthe above gradient is easier to compute than the gradientwith respect to λz , NGD admits a simpler form than GD.
This is the case for many existing works on NGD for VI.
To rewrite NGD as in (6), we need the FIM to be invertible.
For EFs, a sufficient condition for invertibility is to use aminimal representation which is defined3 below using thedefinition given in Wainwright & Jordan (2008).
Definition 1 (Minimal EF) A regular EF representation
is said to be minimal when there does not exist a nonzero
vectors λ such that hφz (z), λi is equal to a constant.
3
A more complete definition is given in Definition 1.3 of Johansen (1979).

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

This essentially means that they are no linear dependencies
in the parameterization of the distribution. When such anonzero vector exists, we can add/subtract it from λ without changing the distribution. Minimality ensures that thisnever happens and the parametrization λz is unique andidentifiable up to multiplication with a nonsingular affine
transformation.
Under a minimal representation, the log-partition function
Az (λz ) is strictly convex, implying that the FIM is positivedefinite.4 For other types of representation, like curvedEFs, the FIM may not be positive-definite. Minimality
ensures that the FIM is positive-definite and that NGD iswell defined 5 . The NGD update can be then carried out
using the expectation parameter which are a one-to-one
function, as stated below (Wainwright & Jordan, 2008).
Theorem 1 The representation (1) is minimal if and only if
the mapping mz (·) : Ω → M is one-to-one.
Unfortunately, minimal EF approximations are not always
appropriate. Such approximations, especially unimodal
ones, usually yield poor approximations of multimodal
posterior distributions. Structured approximations, suchas mixtures of EF distributions, are more suitable for approximating multimodal posteriors. Unfortunately, for suchapproximations, there is no straightforward way to definea minimal EF representation which can be exploited to derive a simple NGD update. For example, a mixture of EF
distributions expressed asZq(z) := q(z|w)q(w)dw,(7)
with q(z|w) as the component and q(w) as the mixing distribution, may not even have an EF form even when boththe terms above are in the EF. A famous example is thefinite mixture of Gaussians. The conditions under which theFIM of (7) is invertible are also difficult to characterize ingeneral. Due to these reasons, it is difficult to simplify the
NGD update for such structured distributions.
In this paper, we propose a new way to derive simple naturalgradients for structured approximations (7) using a minimal
conditional representation which we define next.

3. Minimal Conditional-EF Representation
In this section we define a minimal conditional-EF (MCEF)
representation of the joint distribution q(z, w) for structured
approximations that take the form (7). Using this, we derive
4

A formal proof can be found in Johansen (1979).
In some cases, when FIM is not invertible, it is still possible
to perform NGD by, for example, ignoring the zero eigenvalues
or using damping, but the simplification shown in (6) may not be
possible.
5

conditions under which the FIM of the joint is invertible,
and show that it leads to a simple NGD similar to (6).
We begin with a definition of the conditional EF distribution.
Definition 2 (Conditional EF) We call the joint distribution q(z, w) defined in (7) a conditional EF when its components take the following form:6q(z|w) := hz (z, w) exp [hφz (z, w), λz i − Az (λz , w)] ,q(w) := hw (w) exp [hφw (w), λw i − Aw (λw )] ,(8)

with λw , φw (w) and Aw (λw ) being the natural parameters,
sufficient statistics, and log-partition function of q(w), andλz , φz (z, w) and Az (λz , w) are the same for q(z|w). We
denote the set of natural parameters for q(z|w) and q(w)by Ωz and Ωw respectively, and assume them to be open.7
Note that the sufficient statistics φz (z, w) and log-partition
Az (λz , w) both depend on w, but, conditioned on w, thedistribution is parametrized by the natural parameter λz .
For a well-defined conditional EF distribution, q(w) is aregular EF distribution and conditioned on w, q(z|w) isalso a regular EF distribution. This is a type of conditional
exponential-family distribution (Xing et al., 2002; Lianget al., 2009; Lindsey, 1996; Feigin, 1981) with a special
conditional structure.
We are interested in an NGD update that can exploit theFIM of the joint distribution. We denote the set of naturalparameters by λ := {λz , λw } and the set of valid λz and
λw by Ωz and Ωw respectively. We define the following
FIM in the natural-parameter space Ωz × Ωw as follows:Fwz (λ) := −Eq(z,w) ∇2λ log q(z, w) .
(9)This is the FIM of the joint distribution q(z, w) which isdifferent from the one for the marginal distribution q(z).
Similar to the minimal EF case, our goal now is to findrepresentations where the above FIM is invertible and canbe exploited to compute NGD using expectations of thesufficient statistics. The expectation parameters of a CEF
can be defined as shown below:mz (λz , λw ) := Eq(z|w)q(w) [φz (z, w)] ,(10)mw (λw ) := Eq(w) [φw (w)] .

(11)

We denote the ranges of mz and mw by Mz and Mw
respectively, and the whole set of expectation parameters
6
We assume that Az (λz , w) and Aw (λw ) can be efficiently
computed.
7
It is possible for Ωz to be non-open since it is an intersectionof all of the valid natural-parameter of q(z|w) for each w ∼ q(w).
The set is open when the set of valid natural-parameters for q(z|w)
conditioned on w does not depend on w, or when the cardinalityof the support of w is finite. For all the examples given in thispaper, the set is an open set.

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximationsby m(λ) := {mz , mw }. Since φz (z, w) and Az (λz , w)both depend on w, we may or may not be able to perform
NGD using these expectation parameters. However, we next
show NGD is always possible by restricting to a minimal
representation of CEFs.
Below we define a minimal representation for CEF which
ensures that the Fwz (λ) is positive definite, and NGD canbe performed using m.
Definition 3 (Minimal Conditional-EF (MCEF)) A conditional EF defined in Definition 2 is said to have a minimal
representation when mw (·) : Ωw → Mw and mz (·, λw ) :Ωz → Mz are both one-to-one, ∀λw ∈ Ωw .
In the next section, we will show that all the examples
shown in Figure 1 (See “Examples of MCEFs”) have an
MCEF representation. Similar to minimal EFs, an MCEF
representation implies that Fwz (λ) is positive-definite andinvertible, as stated in the following theorem (see a proof in
Appendix A).
Theorem 2 For an MCEF representation given in Definition 3, the FIM Fwz (λ) given in (9) is positive-definite andinvertible for all λ ∈ Ω.
Since the FIM is well-defined, it is reasonable to perform
natural-gradient steps in the Riemannian manifold definedby Fwz (λ). The steps can be taken using an update thattakes a simple form, similar to the one shown in (6). As
shown in Lemma 5 in Appendix A, similar to (5) for minimal EF, we have the following relationship for MCEFs:∇λ L = ∇λ mT ∇m L = [Fwz (λ)] ∇m L
(12)Since the FIM is invertible, we can compute the naturalgradient by using gradients with respect to m. The following theorem shows the simplicity of NGD.
Theorem 3 For minimal conditional-EF approximations,the following updates are equivalent:−1

λ ← λ + β [Fwz (λ)]λ ← λ + β ∇m L(λ).

∇λ L(λ)(13)(14)

The above theorem generalizes the previous result for minimal EF approximations to minimal conditional-EF approximations. It also implies that using the expectation parameterization m enables us to exploit the geometry of the jointq(z, w) to improve convergence.

4. Examples
In this section, we give examples of approximations witha minimal conditional-EF form. We discuss finite mixtures of EFs, scale mixtures of Gaussians, and multivariate

Skew-Gaussians. We give updates that can be implemented
efficiently by using existing NGD implementations, e.g.,
the variational online Newton (VON) method of Khan et al.
(2018). We also propose many new versions. To derivethese updates, we use the extended Bonnet’s and Price’s
theorems (Lin et al., 2019) for Gaussian mixtures. For
exponential-family mixtures, we use the implicit reparameterization trick (Salimans & Knowles, 2013; Figurnov et al.,2018). Lin et al. (2019) discuss a weaker version of thereparameterization trick for exponential-family mixtures.
4.1. Finite Mixture of Exponential Family Distribution
Finite mixtures of EFs are a powerful approximation wherecomponents in EF form are mixed using a discrete distributions such as a multinomial distribution, as shown below:q(z) =

KX
c=1

πc q(z|λc ), such thatKX

πc = 1,(15)c=1where q(z|λc ) are EF distributions with natural parameters
λc . The support of z of all the components
PK−1 is assumed tobe the same, and πK is fixed to 1 − c=1 πc .
This distribution cannot be written in an exponential formin general, and therefore none of the existing NGD methods
can be used to derive a simple expression for NGVI. Directly
applying NGD using the FIM of q(z) would be too expensive since the number of parameters are in O(D2 K), andthe FIM could be extremely large when computed naively.
Fortunately, the joint distribution does take an MCEF form
when all q(z|λc ) are minimal EFs. A formal statement isgiven in Appendix B. Using this conditional EF form, we
can derive a much simpler NGD update which reduces tosimple parallel updates on mixture components.
We will now demonstrate the simplicity of our update fora finite mixture of Gaussians (MOG) where componentsq(z|λc ) := N (z|µc , Σc ) are Gaussian. The sufficient statistics, natural parameters, and expectation parameters of thecorresponding CEF are given in Table 1.
We consider the following model where the likelihoodp(Dn |z) is defined using neural-network weights z withQN
a prior p(z): p(D, z) = n=1 p(Dn |z)p(z). We approximate the posterior with a MOG by optimizing the following variational lower bound:PL(λ) := Eq(z) [−h(z)] whereh(z) := log [q(z)/p(z)] − n log p(Dn |z).
We now summarize the NGD update derived in Appendix B.
We first generate samples (z, w) from q(z, w). The mean µc
and covariance Σc are similarly updated by the variational
online Newton (VON) algorithm (Khan et al., 2018):
 2

−1
Σ−1
c ← Σc + βδc ∇z h(z) ,µc ← µc − βδc Σc [∇z h(z)] ,log (πc /πK ) ← log (πc /πK ) − β(δc − δK )h(z)(16)

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations
Table 1. We give expressions for various components of q(z, w) as specified in (8) for a minimal conditional-EF representation. We give
three examples. More examples can be found at the Appendix. The first two rows give expressions for the sufficient statistics, and thesubsequent rows show natural and expectation parameters. For the first column, Ic (w) denotes the indicatorfunction which is 1 whenpw = c and 0 otherwise. For the second column, ψ is the digamma function. For the third column, c := 2/π. Note that both the naturaland expectation parameters may lie in a constrained set. Due to space limitations, we have not explicitly given the description of these sets.

φw (w)φz (z, w)
λw
mw (λw )λz
mz (λ)

Mixture of Gaussians
K−1{Ic (w)}c=1
K
Ic (w)z, Ic (w)zzT c=1
K−1{log (πc /πK )}c=1{πc }K−1 −1c=1KΣ µ , − 12 Σ−1
c
c=1
 c c
K
πc µc , πc µc µTc + Σc c=1

T-Distribution−1/w − log wz/w, zzT /w
a−1 − log a + ψ(a)
 −1Σ µ, − 12 Σ−1
µ, µµT + Σ

PKwhere δc := N (z|µc , Σc )/ k=1 πk N (z|µk , Σk ). Thisupdate can be implemented efficiently using approximations
discussed in Khan et al. (2018); Zhang et al. (2018); Mishkin
et al. (2018). For example, implementation of the Adam
optimizer can be utilized. This is simpler and more efficientthan the update which requires computation of the FIM.
A similar example to MOG is the fatigue life distribution(Birnbaum & Saunders, 1969) discussed in Appendix C.
4.2. Scale Mixture of Gaussians
A multivariate scale-mixture of Gaussian (SMG) distribution (Andrews & Mallows, 1974; Eltoft et al., 2006) takes
the following form where the covariance matrix is “scaled”
by a vector w sampled from q(w) such that wi > 0:ZdYq(z) = N (z|µ, LWLT )q(wi )dw,i=1

where W := diag(w) is a diagonal matrix containing w
as the diagonal and L is a matrix with determinant 1 (e.g.,a Cholesky factor) that determines the covariance matrix.
SMG includes many well-known distributions like Student’st, Laplace, logistic, doubly-exponential, normal-gamma,
normal-inverse Gaussian, normal-Jeffreys, and their nonparametric extensions (Caron & Doucet, 2008). For example, the multivariate Student’s t-distribution is obtained byusing a scalar w from an inverse-gamma distribution:Z
T (z|µ, Σ, a) := N (z|µ, wΣ) IG(w|a, a)dw, (17)where Σ := LLT and both the shape and scale parametersof the inverse-gamma distribution are equal to a. We assume
a > 1, since for a ≤ 1 the variance of q(z) does not exist.
SMG is another class of approximations where existing
methods cannot be applied to obtain a simple naturalgradient update. For example, for Student’s t the jointdistribution q(z, w) takes an EF form, but, as we show in

Skew-Gaussianw, w2
z, |w|z, zzT
constants 0, − 12constants {0, 1}
 −1
Σ µ, Σ−1 α, − 12 Σ−1µ + cα, α + cµ,µµT + ααT + Σ + c µαT + αµT

Lemma 8 in Appendix D, the number of free parameters isequal to 3, while the number of natural parameters is equal
to 4. Such representations are not minimal EFs but curved
where the FIM is not positive-definite because the numberof free parameters is less than the number of natural parameters. Therefore, the update (6) does not apply. In contrast,
our update can be applied since the joint distribution is aminimal CEF. This can be verified from Table 1 where bothmw (λw ) and mz (λ) are one-to-one functions. A formal
proof is given in Lemma 10 in Appendix D.
We now demonstrate the simplicity of our update to obtaina Student’s t-approximation on a Bayesian neural network.
We assume the following model with a likelihood p(Dn |z)
specified using a neural networkwith weights z and a StuQN
dent’s t-prior: p(D, z) = n=1 p(Dn |z)T (z|0, I, a0 ). The
Student’s t-prior is better than a Gaussian one if we expect
the weights to follow a heavy-tailed distribution. We approximate the posterior by the approximation (17) usingparameters Σ, µ, and a. We use the variational lower bounddefined in the joint distribution p(D, z, w) space::=−fn (z)

N z
hX
}|
{
p(z, w) i
L(λ) = Eq(z,w|λ)log p(Dn |z) + log(18)
q(z, w)
n=1

Below, we summarize the NGD updates derived in Appendix D. We first sample (z, w) from q and randomly
sampled an example n. The update is then a small modification of the VON update (Khan et al., 2018):
Σ−1 ← (1 − β)Σ−1 + β u∇2z fn (z) + I/N ,(19)µ ← µ − βΣ [∇z fn (z) + µ/N ] ,a ← (1 − β)a + β a0 − δTr ∇2z fn (z)Σ ,

(20)
(21)

2

where β > 0, δ ← N w /(2(1 − w)), and u is a premultiplierh defined below (d is length ofi z): u = (a − 1 +Td/2)−1 a + (z − µ) Σ−1 (z − µ) /2 . Similarly to theprevious section, this update can be implemented efficiently
using the method of Khan et al. (2018).

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

4.3. Gaussian Mean Mixture

Then we divide the vector λ := {λ1 , λ2 , . . . , λK } into K
blocks with λj ∈ Ωj being the j-th block of parameters. In
EF, the function f is just linear in λz at (1). We can generalize the notion of linearity to multiple blocks of parameters
by considering f to be a multilinear function.

We consider the following Gaussian mixture.
Z
k
k
X
Y
q(z|µ, α, Σ) = N (z|µ +u(wi )αi , Σ)q(wi )dw,

Definition 4 (Minimal Multilinear-EF) We call f a multilinear function when, for each block j, there exist functions
φj and rj such that f is linear with respect to λj , i.e.,

Another example is given at Appendix H. Extensions using
an Adam-like optimizer for this kind of mixtures are given
in Appendix J.

i=1

i=1

f (z, λ) = hλj , φj (z, λ−j )i + rj (z, λ−j ) ,

(26)

d

where α is a d-by-k matrix and z ∈ R . An example is
the rank-1 covariance Gaussian with k = 1, u(w) = w,
andR D as a diagonal covariance matrix: q(z|µ, Tα, D) =N (z|µ + wα, D)N (w|0, 1)dw = N (z|µ, αα + D).
The multivariate skew Gaussian (Azzalini & Valle, 1996;
Genton, 2004), defined below, is another example and allows for non-zero skewness (asymmetric approximations):Zq(z|µ, α, Σ) = N (z|µ + |w|α, Σ)N (w|0, 1)dw.
This distribution is not a minimal EF and the FIM of themarginal q(z) can be singular (Azzalini, 2013). However,
the joint distribution is a minimal conditional-EF distribution as we show in Lemma 12 and 13 in Appendix E. The
sufficient statistics, natural parameters, expectation parameters of the conditional EF form are given in Table 1.
Similar to other examples, we get a simple and efficient NGD update. We summarize the updates for amodel with a neural-network likelihood p(Dn |z) using
weights z and a Gaussian prior N (z|0, I). Denoting
fn (z) =P − log p(Dn |z), the lower bound is L(λ) :=Eq(z) [− n fn (z) + log p(z) − log q(z)]. The updates derived in Appendix E is a variant of the VON update:Σ−1 ← (1 − β)Σ−1 + β(I + N gnS )µ ← µ − βΣ c̄(gnµ − cgnα ) + µα ← α − βΣ c̄(gnα − cgnµ ) + α

(22)
(23)
(24)

N/(1 − 2/π) and gnS , gnµ , gnα are gradientswhere c̄ :=obtained by gradient and Hessian of fn (z) at a sample ofq(z). The gradients are defined in (53)-(55) in Appendix E.
Another example of the mixture is the exponentially modified Gaussian distribution (Grushka, 1972; Carr & Madan,2009) given in Appendix F. Extensions using an Adam
optimizer for the class of mixtures are given in Appendix J.

5. Generalization to Multilinear EF
We now extend the approach to an approximation with multilinear EFs which contain blocks of natural parameters. We
start by specifying a distribution over z by a function f (·):
q(z|λ) = hz (z) exp [f (z, λ) − Az (λ)] .

(25)

where λ−j is the parameter vector containing all λ except λj . A distribution q(z|λ) defined as in (25), but witha multilinear f , is called a multilinear EF. Additionally,
when Ωj is openand the following
expectation parametersmj (λ) := Eq(z) φj (z, λ−j ) are one-to-one, we call the
distribution a minimal multilinear EF distribution.
Clearly, minimal EFs are minimal multilinear EFs. The
following theorem give a result about a block NGD update
performed on individual blocks of parameters λj .
Theorem 4 For approximations with multilinearly-minimal
EF representation, the following updates are equivalent:−1

λj ← λj + β [Fj (λ)]∇λj L(λ)λj ← λj + β ∇mj L(λ)

(27)
(28)

The proof of this theorem is similar to Theorem 3. We now
give an example and demonstrate the simplicity of the NGD
update. Let’s consider the Matrix-Variate Gaussian (MVG)
distribution defined as follows:
MN (Z|W, U, V) := N (vec(Z)|vec(W), U ⊗ V).
This distribution has been used for Bayesian neural networks(Louizos & Welling, 2016; Sun et al., 2017). An approximate NGD update is also derived by Zhang et al. (2018)
where the FIM is approximated by a block-diagonal matrix
and K-FAC approximation. Our update have a similar blockdiagonal approximation, but the update for the each blockis an exact NGD unlike Zhang et al. (2018) where the steps
are approximated by K-FAC.
In Appendix I, we show that the MVG distribution
MN (Z|W, U, V) can be written in the minimal multilinear form. The NGD update, derived in Appendix I, issummarized below to optimize the lower bound as L(λ) =Eq [−h(Z)] where h(Z) := − log p(D, Z) + log q(Z). To
simplify our implementation, we use the Gauss-Newton approximation (Graves, 2011) although it is not necessary to
do so. The resulting block NGD update is shown below,
W ← W − β1 UGV,
U

−1

V

−1

←U

−1

←V

−1

(29)
>

+ β2 GVG ,

(30)

>

+ β2 G UG,(31)

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

10
5

10
5

-6

-7

log P

5

-7

10
5

-7

-6

log P

5

10
5

-7

-6

logit u

4
3.5

6

3

5

2.5

-6

4

2

logit u
exact

3

1.5

2

1

-7

logit u
8

4.5

7

10

-6

5

-6

5
-7

logit u7
108

logit u
6

log P

log P

10

9

5
-6

10 -3

10

10

logit u
5

log P

-7

logit u
4

log P

3

log P

2

log P

log P

1

10
1

0.5

5
-7

-6

-7

logit u-60logit u

0

0

5

10

15

20

Figure 2. Quantitative results on three toy examples showing the flexibility obtained by using structured approximations considered in thispaper. The leftmost figure shows the MOG approximation (with K = 20) to fit an MOG model with 10 components in a 20 dimensional
problem. The first 9 dimensions are shown in the figure where we see that MOG approximation fits the marginals well. The middle figure
shows MOG approximation fit to a beta-binomial model for a 2-D problem. The number indicates the number of mixture components
used. By increasing the number of components in our approximation, we get better results. The last figure shows a Skew Gaussian (in red)and a Gaussian (in blue) fit on a 2-D logistic regression posterior. We see that skew-Gaussian captures the skewness in the distribution in
the right direction, and gives better approximation than a single Gaussian.

6. Experimental Results
The code is available at:https://github.com/yorkerlin/VB-MixEF.
6.1. Qualitative Results on Synthetic Examples

breast_cancer_scale

100

10−1

10−2

101

102

103

Iterations

sonar

100

BBVI-1BBVI-3
BBVI-5BBVI-10NGVI-1NGVI-3
NGVI-5NGVI-10KL(q|p)KL(q|p)where we sample Z from the MVG distribution and evaluate the gradient G := ∇Z h(Z). These updates extend
the VON update obtained in Khan et al. (2018) to MVG
approximations. The gradient G is pre-conditioned, which
is very similar to other preconditioned algorithms, such as
K-FAC (Martens & Grosse, 2015; Zhang et al., 2018) and
Shampoo (Gupta et al., 2018). The update can be extended
to Tensor-Variate Gaussian (Ohlson et al., 2013).

104

BBVI-1BBVI-3
BBVI-5BBVI-10NGVI-1NGVI-3
NGVI-5NGVI-10

10−1

10−2

101

102

103

104

Iterations

Figure 3. Bayesian logistic regression approximated by MoG: This
figure demonstrates a fast convergence of NGVI over BBVI. We
use a mixture of Gaussians with full covariance matrix as the approximating distribution. The number next to the method nameindicates the number of mixture components used. The plot shows
the KL obtained using 106 MC samples, where p is the true posterior distribution. For both algorithms, we used full-batches byusing 20 MC samples to compute stochastic approximations.

First, we show qualitative results on three toy examples andvisualize the results obtained by structured approximations.
The first toy example is the Gaussian mixture example
from Wang et al. (2018). In this example, the true distribution is a finite mixture of Gaussians (MOG) p(z) =PC 1di=1 C N (z|ui , I), z ∈ R , where each element of ui is
uniformly drawn from the interval [−s, s]. We approximate
the posterior distribution by an MOG approximation described in Section 4.1. We consider a case with K = 20,
1
C = 10, s = 20, and d = 20. We initialize πc = K
andΣc = 100I. Each element of µc is randomly initialized by
Gaussian noise with mean 0 and variance 100. We use 10
Monte Carlo samples to compute the gradients. The leftmost plot in Figure 2 shows the first 9 marginal distributionsof the true distribution and its approximations, where weclearly see that MOG closely matches the marginals. All 20
marginal distributions are in Figure 6 in Appendix B.5.

In the second toy example we approximate the beta-binomial
model for overdispersion considered in Salimans & Knowles(2013); Salimans et al. (2015) by using MOG (N = 20,
d = 2). The model is to used to estimate the rates of deathfrom stomach cancer for cities in Missouri. The exact posterior of the model is non-standard and extremely skewed. In
the middle plot in Figure 2, we see that our MOG approximation approximates the true posterior better and better as
we increase the number of mixture components.
In the last toy example, we visualize the skew-Gaussian
approximations for the two-dimensional Bayesian Logistic
regression example taken from Murphy (2013) (N = 60,
d = 2). In the rightmost plot in Figure 2, we can see thatthe skewness of the true posterior is captured well by the
skew-Gaussian distribution. The Gaussian approximation
results in a worse approximation than the skew-Gaussian.

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

6 × 104

BBVI
NGVI

106

boston

103

M=290

wine

102

BBVI(Gauss)
NGVI(Skew-Gauss)
BBVI(Skew-Gauss)
102

M=32

Test RMSE

7 × 10

M=232

Negative ELBO

Negative ELBO

8 × 104

4

covtype_scale

107

BBVI
NGVI

Test RMSE

covtype_scale

105

9 × 104

BBVI(Gauss)
NGVI(Skew-Gauss)
BBVI(Skew-Gauss)

101

M=32

100

101

102

103

104

Iterations

105

101
101

102

103

104

Iterations

Figure 4. Bayesian logistic regression with Student’s t (left) andskew-Gaussian approximations (right). We use 10 MC samples
for training, and M denotes the mini-batch size.

6.2. Results on Real Data
Next, we show results on real-world datasets. We consider
two models in our experiments. We start with Bayesian
Logistic regression (BLR) and present results for MOG
approximations on two small UCI datasets. The ‘BreastCancer’ dataset has N = 683, d = 10 with 341 chosen
for training, and regularization parameter is set to 1.0. The‘Sonar’ dataset has N = 208, d = 6 with 100 chosen for
training, and regularization parameter is set to 0.204. Wevary the number of mixture components to K = 1, 3, 5, 10.
In Figure 3, we plot the KL divergence between the true
posterior and the MOG approximation, and compare ourmethod (referred to as ‘NGVI’) proposed in Section 4.1 tothe black-box gradient method (referred to as ‘BBVI’) withthe re-parametrization trick (Salimans & Knowles, 2013;
Kingma & Welling, 2013; Figurnov et al., 2018). For bothmethods, we use a full batch for each update. We observethat NGVI always converges faster than BBVI. We also seethat MoG is a better approximation than the single Gaussian,and the quality of the posterior approximation improves as
the number of mixture increases.
Next, we show results on a larger dataset using two otherkinds of variational approximation: Skew-Gaussian andStudent’s t. We use the UCI dataset “covtype-binary-scale”
with d = 54, N = 581, 012 with 464, 809 chosen for training and regularization parameter 0.002. We use the algorithm discussed in Section 4.2 and 4.3. For black-box methods, we use the Adam optimizer and refer to it as BBVI.
In the skew-Gaussian case, we use a Gaussian prior, and,for the Student’s t-distribution, we use a Student’s t-prior as
shown in (18). Figure 4 demonstrates the fast convergenceof our method compared to BBVI.
Finally, we discuss results on Bayesian neural network(BNN) with a standard normal prior on weights. We use one
hidden layer, 50 hidden units, and ReLU activation functions. We approximate the posterior by a skew-Gaussian
distribution using NGVI. We also compare to two othermethods where we used BBVI to fit a skew-Gaussian approximation as well as a Gaussian approximation. For scalability reasons, we use of a diagonal covariance for all

0

1000

2000

Iterations

3000

4000

0

5000

10000

Iterations

Figure 5. BNN using skew-Gaussian approximation: This figure
shows a fast convergence of NGVI over BBVI to approximate theposterior of BNN. For all methods, the prior is a Gaussian. We
use 10 MC samples for training. M in the figures denotes the sizeof a mini-batch. BBVI (Gauss) uses a Gaussian approximation
while BBVI (Skew-Gauss) uses a skew-Gaussian approximation.
Our NGVI method with skew-Gaussian approximation converges
faster than the other two methods.

methods. We use 10 Monte Carlo (MC) samples and minibatch size of 32. For NGVI, we use the gradient magnitudeapproximation as explained in Appendix J. Figure 5 shows
the performance of all methods in terms of the test RMSE.
We can see that our method converges faster than BBVI,
although the performance of skew-Gaussian methods seemto be similar to a Gaussian.

7. Discussion
In this paper, we present fast and simple NGD for VI withstructured approximations. The approximations we haveconsidered are currently beyond the reach of existing methods, and our approach extends these existing approaches
to perform NGD updates with a simple update which canalso be implemented efficiently in some cases. Our current
proposal is limited to a certain class of approximations, andfurther work is needed to generalize our results to manyother types of structured approximations. The minimality
condition we proposed uses one-to-one mappings of the
expectation parameterization. We believe that this condition
can be relaxed which will enable simple NGD update formany types of approximations.
Our main focus has been on the derivation of simple updates.
We have presented examples where the updates can also be
implemented efficiently. There are however implementation
bottlenecks in existing software frameworks to implement
some of the reparameterization tricks used in our algorithms.
It is important to find ways to enable efficient updates bymodifying the existing software frameworks. Another issue
is that the NGD update needs to make sure that the parameters stay inside Ω, and this issue deserves further exploration.
Another important direction is to apply structured approximations to large problems, especially those involving deep
networks. We hope to perform such extensive experimentsin the future to establish the benefits obtained by our NGD
update for Bayesian deep learning.

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

Acknowledgements
We would like to thank Voot Tangkaratt for useful discussions during this work. WL is supported by a UBC International Doctoral Fellowship.

References
Andrews, D. F. and Mallows, C. L. Scale mixtures of normal
distributions. Journal of the Royal Statistical Society.
Series B (Methodological), pp. 99–102, 1974.
Arellano-Valle, R. B., Contreras-Reyes, J. E., and Genton,
M. G. Shannon Entropy and Mutual Information for
Multivariate Skew-Elliptical Distributions. Scandinavian
Journal of Statistics, 40(1):42–62, 2013.
Azzalini, A. The skew-normal distribution and related multivariate families. Scandinavian Journal of Statistics, 32(2):159–188, 2005.
Azzalini, A. The skew-normal and related families, volume 3. Cambridge University Press, 2013.
Azzalini, A. and Valle, A. D. The multivariate skew-normal
distribution. Biometrika, 83(4):715–726, 1996.
Balakrishnan, N. and Kundu, D. Birnbaum-Saunders distribution: A review of models, analysis, and applications.
Applied Stochastic Models in Business and Industry, 35(1):4–49, 2019.
Barndorff-Nielsen, O. E. Normal inverse Gaussian distributions and stochastic volatility modelling. Scandinavian
Journal of statistics, 24(1):1–13, 1997.
Batir, N. Some new inequalities for gamma and polygammafunctions. J. Inequal. Pure Appl. Math, 6(4):1–9, 2005.
Birnbaum, Z. W. and Saunders, S. C. A new family oflife distributions. Journal of applied probability, 6(2):319–327, 1969.
Bonnet, G. Transformations des signaux aléatoires a travers
les systemes non linéaires sans mémoire. In Annales desTélécommunications, volume 19, pp. 203–220. Springer,
1964.
Caron, F. and Doucet, A. Sparse Bayesian nonparametric regression. In Proceedings of the 25th international
conference on Machine learning, pp. 88–95. ACM, 2008.
Carr, P. and Madan, D. Saddlepoint methods for option
pricing. The Journal of Computational Finance, 13(1):49, 2009.
Choi, J., Du, Y., and Song, Q. Inverse Gaussian quadratureand finite normal-mixture approximation of generalized
hyperbolic distribution. arXiv preprint arXiv:1810.01116,
2018.

Contreras-Reyes, J. E. and Arellano-Valle, R. B. Kullback–Leibler divergence measure for multivariate skew-normal
distributions. Entropy, 14(9):1606–1626, 2012.
Culham, R. Lecture Notes: Advance differential equations
and special functions. www.mhtlab.uwaterloo.
ca/courses/me755/web_chap4.pdf, 2004. Accessed: 2019/03/25.
Desmond, A. On the relationship between two fatigue-lifemodels. IEEE Transactions on Reliability, 35(2):167–
169, 1986.
Eltoft, T., Kim, T., and Lee, T.-W. Multivariate scale mixtureof Gaussians modeling. In International Conference on
Independent Component Analysis and Signal Separation,
pp. 799–806. Springer, 2006.
Feigin, P. D. Conditional Exponential Families and a Representation Theorem for Asympotic Inference. The Annalsof Statistics, pp. 597–603, 1981.
Figurnov, M., Mohamed, S., and Mnih, A. Implicit Reparameterization Gradients. 2018.
Furmston, T. and Barber, D. Variational methods for reinforcement learning. In Proceedings of the ThirteenthInternational Conference on Artificial Intelligence andStatistics, pp. 241–248, 2010.
Genton, M. G. Skew-elliptical distributions and their applications: a journey beyond normality. CRC Press, 2004.
Graves, A. Practical variational inference for neural networks. In Advances in neural information processing
systems, pp. 2348–2356, 2011.
Grushka, E. Characterization of exponentially modified
Gaussian peaks in chromatography. Analytical Chemistry,
44(11):1733–1738, 1972.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned Stochastic Tensor Optimization. In Proceedings ofthe 35th International Conference on Machine Learning,pp. 1842–1850, 2018.
Hensman, J., Rattray, M., and Lawrence, N. D. Fast variational inference in the conjugate exponential family. In
Advances in neural information processing systems, pp.
2888–2896, 2012.
Hensman, J., Fusi, N., and Lawrence, N. D. Gaussian
processes for big data. arXiv preprint arXiv:1309.6835,
2013.
Hoffman, M. D. and Blei, D. M. Structured stochastic variational inference. In Artificial Intelligence and Statistics,
2015.

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J.
Stochastic variational inference. The Journal of Machine
Learning Research, 14(1):1303–1347, 2013.
Honkela, A., Tornio, M., Raiko, T., and Karhunen, J. Natural
conjugate gradient in variational inference. In International Conference on Neural Information Processing, pp.
305–314. Springer, 2007.
Honkela, A., Raiko, T., Kuusela, M., Tornio, M., andKarhunen, J. Approximate Riemannian Conjugate Gradient Learning for Fixed-form Variational Bayes. Journalof Machine Learning Research (JMLR)), 11:3235–3268,
2011.
Johansen, S. Introduction to the Theory of Regular Exponential Famelies. 1979.
Jørgensen, B., Seshadri, V., and Whitmore, G. On the
mixture of the inverse Gaussian distribution with its complementary reciprocal. Scandinavian Journal of Statistics,
pp. 77–89, 1991.
Khan, M. and Lin, W. Conjugate-computation variational inference: Converting variational inference in nonconjugate models to inferences in conjugate models. In
Artificial Intelligence and Statistics, pp. 878–887, 2017.
Khan, M. E. and Nielsen, D. Fast yet Simple NaturalGradient Descent for Variational Inference in Complex
Models. arXiv preprint arXiv:1807.04489, 2018.
Khan, M. E., Babanezhad, R., Lin, W., Schmidt, M., andSugiyama, M. Faster stochastic variational inference using Proximal-Gradient methods with general divergence
functions. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, pp. 319–
328. AUAI Press, 2016.
Khan, M. E., Nielsen, D., Tangkaratt, V., Lin, W., Gal, Y.,
and Srivastava, A. Fast and Scalable Bayesian Deep
Learning by Weight-Perturbation in Adam. In Proceedings of the 35th International Conference on Machine
Learning, pp. 2611–2620, 2018.
Kingma, D. P. and Welling, M. Auto-encoding variational
bayes. arXiv preprint arXiv:1312.6114, 2013.
Kotz, S. and Nadarajah, S. Multivariate t-distributions andtheir applications. Cambridge University Press, 2004.
Kumar, S. and Tsvetkov, Y. Von Mises-Fisher Loss forTraining Sequence to Sequence Models with Continuous
Outputs. arXiv preprint arXiv:1812.04616, 2018.
Liang, P., Jordan, M. I., and Klein, D. Learning from measurements in exponential families. In Proceedings of the
26th annual international conference on machine learning, pp. 641–648. ACM, 2009.

Lin, W., Khan, M. E., and Schmidt, M. Stein’s Lemma for
the Reparameterization Trick with Exponential-family
Mixtures. https://github.com/yorkerlin/
VB-MixEF/blob/master/report.pdf, 2019.
Accessed: 2019/06.
Lindsey, J. K. Parametric statistical inference. Oxford
University Press, 1996.
Louizos, C. and Welling, M. Structured and efficient variational deep learning with matrix gaussian posteriors.
In International Conference on Machine Learning, pp.
1708–1716, 2016.
Martens, J. and Grosse, R. Optimizing neural networks with
Kronecker-factored approximate curvature. In International Conference on Machine Learning, pp. 2408–2417,
2015.
Mishkin, A., Kunstner, F., Nielsen, D., Schmidt, M., andKhan, M. E. SLANG: Fast Structured Covariance Approximations for Bayesian Deep Learning with Natural
Gradient. In Advances in Neural Information Processing
Systems, pp. 6246–6256, 2018.
Murphy, K. P. Machine learning : a probabilistic perspective. MIT Press, Cambridge, Mass. [u.a.], 2013. ISBN
9780262018029 0262018020.
Nguyen, C. V., Li, Y., Bui, T. D., and Turner,
R. E. Variational Continual Learning. arXiv preprint
arXiv:1710.10628, 2017.
Oh, C., Adamczewski, K., and Park, M. Radial and Directional Posteriors for Bayesian Neural Networks. arXiv
preprint arXiv:1902.02603, 2019.
Ohlson, M., Ahmad, M. R., and Von Rosen, D. The multilinear normal distribution: Introduction and some basic
properties. Journal of Multivariate Analysis, 113:37–47,
2013.
Opper, M. and Archambeau, C. The variational Gaussian
approximation revisited. Neural computation, 21(3):786–
792, 2009.
Price, R. A useful theorem for nonlinear devices having
Gaussian inputs. IRE Transactions on Information Theory, 4(2):69–72, 1958.
Ranganath, R., Gerrish, S., and Blei, D. Black box variational inference. In Artificial Intelligence and Statistics,
pp. 814–822, 2014.
Ranganath, R., Tran, D., and Blei, D. Hierarchical variational models. In International Conference on Machine
Learning, pp. 324–333, 2016.

Fast and Simple Natural-Gradient Variational Inference with Mixture of Exponential-family Approximations

Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic
backpropagation and approximate inference in deep generative models. arXiv preprint arXiv:1401.4082, 2014.
Ruiz-Antolı́n, D. and Segura, J. A new type of sharp bounds
for ratios of modified Bessel functions. Journal of Mathematical Analysis and Applications, 443(2):1232–1246,
2016.
Salimans, T. and Knowles, D. Fixed-form variational posterior approximation through stochastic linear regression.
Bayesian Analysis, 8(4):837–882, 2013.
Salimans, T., Kingma, D., and Welling, M. Markov chainmonte carlo and variational inference: Bridging the gap.
In International Conference on Machine Learning, pp.
1218–1226, 2015.
Salimbeni, H., Eleftheriadis, S., and Hensman, J. Natural
Gradients in Practice: Non-Conjugate Variational Inference in Gaussian Process Models. International Conference on Artificial Intelligence and Statistics (AISTATS),2018.
Sato, M.-A. Online model selection based on the variational
Bayes. Neural computation, 13(7):1649–1681, 2001.
Staines, J. and Barber, D. Variational optimization. arXiv
preprint arXiv:1212.4507, 2012.
Sun, S., Chen, C., and Carin, L. Learning Structured Weight
Uncertainty in Bayesian Neural Networks. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, AISTATS, pp. 1283–1292, 2017.
Titsias, M. and Lázaro-Gredilla, M. Doubly stochastic variational Bayes for non-conjugate inference. In International
Conference on Machine Learning, pp. 1971–1979, 2014.
Titsias, M. K. and Ruiz, F. J. Unbiased Implicit Variational
Inference. arXiv preprint arXiv:1808.02078, 2018.
Wainwright, M. J. and Jordan, M. I. Graphical models, exponential families, and variational inference. Foundations
and Trends in Machine Learning, 1–2:1–305, 2008.
Wang, D., Liu, H., and Liu, Q. Variational Inference with
Tail-adaptive f-Divergence. In Advances in Neural Information Processing Systems, pp. 5742–5752, 2018.
Xing, E. P., Jordan, M. I., and Russell, S. A generalized
mean field algorithm for variational inference in exponential families. In Proceedings of the Nineteenth conference
on Uncertainty in Artificial Intelligence, pp. 583–591.
Morgan Kaufmann Publishers Inc., 2002.
Yang, Z.-H. and Chu, Y.-M. On approximating the modified
Bessel function of the second kind. Journal of Inequalities and Applications, 2017(1):41, 2017.

Yin, M. and Zhou, M. Semi-Implicit Variational Inference.
arXiv preprint arXiv:1805.11183, 2018.
Zhang, G., Sun, S., Duvenaud, D., and Grosse, R. Noisy
Natural Gradient as Variational Inference. In Proceedings of the 35th International Conference on Machine
Learning, 2018.