V ARIATIONAL M ESSAGE P ASSING WITH S TRUCTURED I NFERENCE N ETWORKS
Anonymous authors
Paper under double-blind review
A BSTRACT
Powerful models require powerful algorithms to be effective on real-world problems.
We propose such an algorithm for a class of models that combine deep models with probabilistic graphical models to develop interpretable deep-learning models.
Our algorithm is a natural-gradient message-passing method whose messages automatically reduce to stochastic-gradients for the deep components of the model.
Using a special-structure inference network, our algorithm exploits the structural properties of the model to gain computational efﬁciency while retaining the simplicity and generality of deep-learning algorithms.
By combining the strength of two different types of inference procedures, our approach offers a framework that simultaneously enables structured, amortized, and naturalgradient inference for complex models.
I NTRODUCTION
To analyze real-world data, machine learning relies on models that can extract useful patterns.
For this purpose, deep models such as Deep Neural Networks (DNNs) are a popular choice and can learn ﬂexible representations.
Another popular choice are probabilistic graphical models (PGMs) which can ﬁnd interpretable structures in the data.
Recent work on combining these two types of models hopes to borrow their complimentary strengths and provide powerful models that are also easy to interpret (Johnson et al., 2016; Krishnan et al., 2015;
Archer et al., 2015; Fraccaro et al., 2016).
To apply such hybrid models to real-world problems, we need efﬁcient algorithms that can extract useful structure from the data.
However, the two ﬁelds of deep learning and PGMs traditionally use different types of algorithms.
For deep learning, stochastic-gradient (SG) methods are the most popular choice, e.g., those based on back-propagation.
These algorithms are not only widely applicable, but also employ amortized inference to enable fast inference at test time (Rezende et al., 2014; Kingma & Welling, 2013).
On the other hand, most popular algorithms for PGMs exploit the model’s graphical structure to gain computational efﬁciency, e.g., variational message passing (VMP) (Winn & Bishop, 2005), expectation propagation (Minka, 2001), Kalman ﬁltering, and more recently stochastic variational inference (SVI) (Hoffman et al., 2013).
The two ﬁelds of deep learning and PGMs employ fundamentally different inferential strategies and a natural question is, whether we can design algorithms that combine their respective strengths.
In this paper, we propose such an algorithm.
Our algorithm is a natural-gradient message-passing algorithm whose messages automatically reduce to stochastic-gradients for the deep components of the model.
In addition to such a ﬂexible message-passing scheme, our algorithm performs amortized inference by using a structured inference network that mimics the structure of the model.
Overall, this enables Structured, Amortized, and Natural-gradient (SAN) updates and therefore we call our algorithm the SAN algorithm.
Unlike many existing methods that continue using SG methods even when some of the model components are PGMs, our SAN algorithm exploits the probabilistic structure to gain computational efﬁciency.
Our algorithm is in a similar spirit to the work of Johnson et al.
but it enjoys much simpler implementation.
Experimental comparisons show that our method gives comparable performance to the method of Johnson et al.
(2016) while simplifying and generalizing it.
Latent mixture model
Latent state-space model y n x n z n θ NN ψ π
N (a) Model y n x n z n θ NN ψ π φ NN ¯ ψ ¯ π
N (b) Inference Network yyyyxxxxθ PGM θ NN (c) Model yyyyxxxxθ PGM θ NN φ PGM φ NN (
[Figure]
T HE M ODEL AND C HALLENGES WITH I TS I NFERENCE
We consider the modeling of data vectors y n by using local latent vectors x n .
The correlations among N data vectors y := { y 1 , y 2 , . . .
, y N } can then be captured by using a probabilistic graphical model (PGM) over all the latent vectors x := { x 1 , x 2 , . . .
, x N } .
This combination of probabilistic graphical model and neural network is referred to as structured auto-encoder (SVAE) by
Johnson et al.
and has also been used in other recent works (Krishnan et al., 2015; Archer et al., 2015).
Following previous approaches, we consider the following model in this paper:
[Equation]
where θ NN are the parameters of DNN, θ PGM parameters of PGM, and θ := { θ NN , θ PGM } .
Below, we give two examples of graphical models.
Mixture Models: To discover patterns in the data, we can model x n by using a ﬁnite mixture model with K mixture components.
To do so, we deﬁne additional discrete latent vector z n whose k ’th element is z nk ∈ { 0 , 1 } .
We can use then as a Gaussian mixture model as the PGM prior,
[Equation]
where µ k and Σ k are mean and covariance of the k ’th mixture and π k are the mixture proportions summing to 1, and θ PGM := { ψ k , π k } Kk =1 and ψ k := { µ k , Σ k } .
The parameter prior can be chosen as described in Bishop (2006) and the corresponding graphical model is shown in Figure 1a.
This is an example of a model that combines PGM with deep networks.
State-Space Models: When y n is a time series, we can model the latent x n using a linear dynamical system, for example, we can use the following model:
[Equation]
where A is the transition matrix, Q is the process noise, and µ 0 and Σ 0 are the mean and covariance of the initial distribution.
Here, θ PGM := { A , Q , µ 0 , Σ 0 } .
Again, we can use standard conjugatepriors for θ PGM (Beal, 2003).
The graphical model is shown in Figure 1c.
For these models, our goal is to approximate the posterior distribution p ( x , θ | y ) using an inference network that can be used to perform amortized inference (Ritchie et al., 2016).
Such inference networks have previously been used in variational auto-encoders (VAE) (Kingma & Welling,
2013) where they are shown to perform extremely well, however, in our case, designing such highperformance networks is a challenging task.
The main challenge lies in designing networks that can ﬂexibly capture the structure present in the
PGM component, i.e., p ( x | θ ) .
In VAEs, a DNN with parameters φ is used as the inference network q ( x | y , φ ) .
This choice is reasonable since all of the structure in the posterior is attributed to the DNN likelihood p ( y | x , θ ) .
The PGM part in this case is simply a standard normal distribution N ( x | 0 , I ) which does not have any special structure.
Therefore using a pure DNN based inference-network is sufﬁcient.
Our model is explicitly designed to exploit the PGM structure so that we can have interpretable deep models.
Therefore, it is essential to capture such structure in the inference network.
However, a purely DNN based inference network puts all its power in the likelihood term and ignores the structure.
This might give inaccurate predictions.
For example, consider the case where x is a time series (as shown in Fig. 1c).
Here, the prediction of a future observation, say, y 10 at time n = 10 given past data, say, { y 1 , y 2 , y 3 } would be inaccurate because the inference network does not capture the correlation between x 10 and x 3 .
Rather, the network relies on a DNN that is conditioned on the unavailable observation y 10 .
For such forecasting and imputation tasks, the inference network must capture the structure of the posterior accurately.
A solution to this problem is to exploit the PGM structure in the inference network by replacing all of its edges by neural networks (Krishnan et al., 2015; Fraccaro et al., 2016).
When the PGM itself is complex, this solution is reasonable, but when it is a simple model, such as linear dynamical system, this approach might be too aggressive.
Using DNNs in such cases would dramatically increase the number of parameters.
This, in turn, will lead to a loss in speed and a possible deterioration in performance.
Johnson et al.
(2016) propose an alternative double-loop procedure to make sure that the probabilistic structure in PGM can be captured in the posterior distribution.
Their algorithm however requires running a double loop algorithm where the inner loop needs to be run until convergence otherwise the algorihtm might diverge.
We give more details in the next section, but their approach is computationally intensive and might converge slowly.
To summarise, existing methods either simply ignore the graphical model structures or do not exploit them efﬁciently.
To solve this problem, we propose an inference network that preserves the structure of the posterior to facilitate fast and accurate amortized inference.
We incorporate this network within a messagepassing framework.
For the deep components of the model, these messages automatically reduce to stochastic-gradients, while for the rest of the model, they are used to perform natural-gradient updates, just like in SVI (Hoffman et al., 2013).
We give two examples of our approach and apply them to the experiments seen in Johnson et al.
Our results show that our method performs comparably, while simplifying and generalizing their approach.
C ONJUGATE -S TRUCTURE I NFERENCE N ETWORKS
As explained in the previous section, designing an inference network that preserves the prior’s probabilistic structure and enables ﬂexible amortized inference is a challenging task.
We solve this problem by proposing a structured inference network that takes the following form:
[Equation]
where p ( x | φ PGM ) is the same distribution as p ( x | θ PGM ) but is parametrized by a new parameter φ PGM , and q ( x | y , φ NN ) is an exponential-family distribution whose natural parameters are parametrized by a DNN in a similar fashion as in VAE.
The role of the PGM term is to enforce the PGM structure, while the role of DNN term is to enable ﬂexibility.
For amortized inference, we need to be able to generate samples from q ( x | y , φ ) which is an easy when we the normalizing constant Z ( φ ) is available in closed form.
To satisfy this requirement, we choose the network q ( x | y , φ NN ) to be conjugate to p ( x | φ PGM ) , i.e., they both must take the same functional form with respect to x .
For example, if PGM is a Gaussian, we can choose q ( x | y , φ NN ) to be a Gaussian as well.
This way q ( x | y , φ ) will also be Gaussian and Z ( φ ) can be obtained in closed-form.
In general, we choose the ﬁrst term to be a member of conjugate exponential family.
Sampling can then be performed using conjugate message-passing algorithm algorithm.
We call our inference network a conjugate-structure inference networks.
An important feature of our inference network is that it is independent of θ .
Due to this our method does not suffer from the “implicit gradient” problems encountered in the method of Johnson et al.
Their method approximates q ( x ) ∝ q ( x | y , φ ) p ( x , θ ) by optimizing a variational lower bound.
As a result, the parameters of q ( x ) depend on θ .
Due to this dependence, if the inner loop is not run until convergence, the outer loop might not have correct gradients which might result in the divergence of their method.
By simply using new parameters φ PGM , we make sure that the normalizing constant Z as well as the generated samples are all independent of θ PGM .
This simpliﬁes the inference because we do not have to worry about the inference network when taking gradients with respect to the variational parameters corresponding to the variational distribution of θ .
Another consequence of our proposed inference network is that when the original PGM is nonconjugate (e.g., a T-distribution), we can still use a conjugate term in the second term (e.g., a Gaussian) in (4).
We only have to make sure that the inference network captures useful structure in the posterior.
This modiﬁcation enables us to design easy-to-sample inference networks for complex models.
The conjugate structure in our inference network simpliﬁes the inference by replacing difﬁcult terms in the model by conjugate terms in the inference network.
The additional parameters φ PGM do not lead to a massive increase in parameters since the number of global parameters φ PGM is usually quite small.
Finally, using the above inference network, we can now specify the overall posterior approximation by assuming a mean-ﬁeld approximation among x and θ :
[Equation]
where q ( θ | λ ) is a ﬁxed-form minimal exponential-family distribution parametrized by λ .
Below, we give two examples of inference networks and show how to compute Z ( φ ) .
Mixture Models: As described earlier, we can easily construct a conjugate-structure inference network where the normalizing constant is available in closed form.
An example is shown below where the ﬁrst component is an inference network as in VAE, while the second and third components correspond to a GMM:
[Equation]
where m y n := m ( y n , φ NN ) and V y n := V ( y n , φ NN ) are mean and covariance parameterized by a
DNN with parameter φ NN , φ PGM := { ¯ ψ k , ¯ π k } Kk =1 and ¯ ψ k = { ¯ µ k , ¯ Σ k } are the parameters on the
PGM part in the inference network.
The increase in number of parameter is only O ( K ) which is usually much smaller than N .
The graphical model of this inference network is shown in Figure 1b, where we clearly see that the structured inference network mimics the structure of the model.
An attractive property of the conjugate-structure inference network is that its normalizing constant is available in closed-form.
In this case, we can obtain the following closed-form solution as described in Bishop (2006):
[Equation]
The above inference network is valid even when the mixture components are non-Gaussian, e.g., when they follow the Student’s t distribution.
For such cases, our inference network provides an efﬁcient way to perform by replacing the non-Gaussian distribution by a Gaussian one.
State-Space Models:
In this case, we can use the following conjugate-structure inference network:
[Equation]
where m y n and V y n are deﬁned similarly to the mixture-model case, and the PGM parameters are φ PGM := { ¯ A , ¯ Q , ¯ µ 0 , ¯ Σ 0 } .
The graphical model of this inference network is shown in Figure 1d.
The normalizing constant for this inference network can be computed in closed-form.
To see this, note that the ﬁrst term can be written as a Gaussian measurement N ( x n | m y n , V y n ) :=
N ( m y n | x n , V y n ) , i.e., m y n is a Gaussian measurement of x n .
Therefore, (8) is a linear dynamical system where the marginal likelihood (which is equal to Z ( φ ) , can be computed using the forward-backward algorithm in O ( N ) (a.k.a.
Kalman ﬁlter)
Bishop (2006).
Here again, when the dynamical system in (3) is non-linear, we can still use the inference network (8).
Therefore, the structured inference network is simpliﬁes the inference by replacing the nonlinear terms in the model by linear ones.
V ARIATIONAL M ESSAGE P ASSING
In the previous section, we proposed a conjugate-structure inference network.
Now, we describe a variational message-passing algorithm to perform efﬁcient inference with it.
An attractive property of our algorithm is that it enables natural-gradient updates for the PGM part, while its messages reduce to stochastic gradients when applied to deep neural networks.
Our algorithm is an extension of the method proposed by Khan & Lin (2017) and we extend to allow the use of inference network.
Our goal is to maximize the following lower bound:
[Equation]
A straightforward solution is to use stochastic-gradient methods, e.g., RMSprop or Adam.
These methods are suitable for estimating deterministic parameters φ , however they may be inefﬁcient to estimate λ .
This is because λ is the parameter of a distribution and SG methods are not suited fot this purpose (see a discussion in Hoffman et al.
A preferred way to optimize with respect to λ is to use natural-gradient methods such as stochastic variational inference (Hoffman et al., 2013).
Natural gradients exploit the geometry of the probabilistic model and show faster convergence.
For mean-ﬁeld variational inference, they are also easier to compute and updates can be performed locally at a node within the message-passing framework.
SG methods cannot exploit nice properties of natural-gradients and might be slow to converge.
An alternative is to use natural-gradients for both φ and λ , but this is computationally intensive because a direct computation of the Fisher matrix for deterministic parameters is difﬁcult.
We could also simply apply natural-gradient updates to λ while use SG methods for φ , but this approach is not principled and the step-size selection might be difﬁcult to tune due to the interaction between these different types of updates.
We propose a principled approach to solve this problem.
The beneﬁt of our approach is that we can automatically recover SG updates for φ as a special case, while enabling natural-gradient updates for λ .
The key idea is to associate φ with a Gaussian distribution q ( φ ) :
= N ( φ | µ φ , diag ( σ 2 φ )) where µ φ and σ φ are vectors of mean and standard deviation.
Given this distribution, we maximize the following lower bound:
[Equation]
This approach is known as variational optimization (Staines & Barber, 2012) and it works due to the following inequality: max L ( φ , λ ) ≥ E q ( φ )
[ L ( φ , λ )] , i.e., the maximum of a function is lower bounded by its expectation.
Therefore, maximizing the expectation also maximizes the function.
The variance eventually shrinks to zero when the algorithm converges, therefore the resulting algorithm converges to a local minima of L even though we optimize ¯ L .
The good news is that now all of the optimization parameters, namely λ , µ , and Σ are parameters of distributions.
Therefore, we can use natural gradient updates.
A ﬁnal step to obtain our message-passing algorithm is the choice of the natural-gradient method.
The approach of SVI only works for conjugate models, but the DNN likelihood in our network is non-conjugate.
We will therefore use a method proposed by Khan & Lin (2017) which is a generalization of SVI.
In this method, the messages for the conjugate part of the model reduces to SVI-like updates, but the messages for the rest of the model use stochastic-gradients.
We now give detailed updates for the parameter q ( θ PGM | λ PGM ) , q ( θ NN | λ NN ) , and q ( φ | µ φ , σ 2 φ ) .
Update for q ( θ PGM | λ PGM ) :
We assume that p ( θ PGM ) and p ( x | θ PGM ) are conjugate to each other, although the algorithm can easily be modiﬁed when the relationship is non-conjugate.
Conjugacy implies that the two distribution can be rewritten as follows:
[Equation]
where η is the natural parameter, T is the sufﬁcient statistics, (cid:104)· , ·(cid:105) is an inner product, h is the base measure, A ( η ) is the log-partition function, and η xθ and f x are functions that depend on x only.
The factors that involve θ PGM are p ( x | θ PGM ) and p ( θ PGM ) , therefore using Algorithm 2 of Khan & Lin (2017), we can write the natural-gradient updates as follows:
[Equation]
where β > 0 is a step-size.
This update is simple and does not require any gradient computations.
To compute the expectation we simply generate samples x ∗ from the inference network.
Given this sample, the computation of η xθ is equivalent to a variational inference in the model p ( x ∗ , θ PGM ) which can be performed using standard methods such as SVI and VMP.
Therefore we can reuse existing software to perform this update.
Update for q ( θ NN | λ NN ) :
To derive the updates, we will assume that the prior of all elements of θ NN is a Gaussian with mean µ 0 and variance σ 2 0 .
Since the number of parameters for DNNs is usually large, we will use a mean-ﬁeld approximation.
We will use a Gaussian approximation which makes sense since the prior is Gaussian too.
For simplicity, we show updates only for the i ’th element of θ NN .
We denote the mean and variance of this element by µ i and σ 2 i respectively.
The factors that involve θ NN in the joint distribution (1) are p ( θ NN ) and p ( y | x , θ NN ) .
To derive the updates using Algorithm 2 of Khan & Lin (2017), we ﬁrst note that the natural parameter of a Gaussian is equal to µ i /σ
2 i .
Using this and a few straightforward manipulations to express the gradients in terms of µ i and σ 2 i , we get the following updates:
[Equation]
These gradients can be approximated by using samples from the distribution and by using the reparameterization trick.
These updates are very similar to adaptive-gradient methods because the update for the mean is scaled by the variance which acts as the scaling vector in the adaptive methods.
There are a two small differences.
First, a regularization term µ 0 /σ 20 and 1 /σ 20 is added.
Second, gradient with respect to variance is computed.
Third gradient are computed at a randomly sampled location from q .
In our experiments, this does not make a big difference and our method converges at a very similar rate as the adaptive gradient methods.
Update for q ( φ | µ φ , σ 2 φ ) :
To derive these updates, we note that the neighbors of φ are the factors that use x , i.e., p ( y | x , θ NN ) , p ( x | θ PGM ) and q ( x | y , φ ) itself.
It is easier to derive the update by simply taking the derivative of the lower bound as a function of φ shown below:
[Equation]
Since the distribution is Gaussian we can use updates very similar to (14)-(15).
To compute the gradient we can use reparameterization trick since a sample from q ( x | y , φ ) is a function of µ φ and σ 2 φ .
The updates again resemble adaptive-gradient methods and can be easily obtained.
Summary of the method: The method of Khan & Lin (2017) is guaranteed to converge under stochastic SVI like updates where a local variable update is followed by the update of global variable.
Since we are using an inference network for the local variables, this is equivalent to generating samples from q ( x | y , φ ) .
Using the generated samples we can either perform the above three updates in parallel of use a sequential mode.
It is also allowed to randomly pick a global variable and update them.
E XPERIMENTS AND R ESULTS
The main goal of our experiment is to show that our SAN algorithm gives similar results to the method of Johnson et al.
The advantage of our algorithm over the method of Johnson et al. [Figure] (Johnson et. al.)
[Figure] (2016) is that our algorithm is simpler and more general than their method, and does not involve a double-loop procedure.
For this reason, we apply our algorithm to the two examples considered in
Johnson et al.
(2016), namely the latent Gaussian mixture model (GMM) and latent linear dynamical system (LDS).
In this section, we show results on GMM (see details of the model in in section 3).
An additional result for LDS is included in
Appendix A.
Our results show that the SAN algorithm can learn complex representations with interpretable structure, similar to the method of Johnson et al.
For the GMM case, we compare SAN to three algorithms: ﬁrst, the standard variational EM algorithm applied to GMM (referred to as ‘GMM’), second, the VAE approach of Kingma & Welling (2013), and third, the SVAE approach of Johnson et al.
We use two datasets.
The ﬁrst dataset is the synthetic two-dimensional Pinwheel dataset ( N =and D = 2 ), and the second dataset is the Auto dataset ( N = 392 and D = 6 , available in the UCI repository).
The latter contains information about cars.
The dataset also contains a ﬁve-class label which indicates the number of cylinders in a car.
We use these labels to validate our results.
For both datasets we use 70 % data for training.
For all methods, we tune their stepsizes, the number of mixture components, and the latent dimensionality on a validation set.
We train the GMM baseline on the full dataset, and, for VAE and SVAE, we use stochastic minibatches of size 64 .
Our neural-network likelihood consists of two layers with 50 hidden units and an output layer of dimensionality 6 for the Auto dataset and 2 for the Pinwheel dataset.
In Figure 2, we compare the performance during training.
In the left ﬁgure, we compare to SVAE and GMM, where we see that SAN converges faster than SVAE.
Both methods achieve similar performance upon convergence and perform better than GMM.
In the right ﬁgure, we compare to VAE and observe similar trends.
In this ﬁgure, the performance of GMM is represented as a constant because training converges after a few iterations already.
We found that the implementation provided by Johnson et al.
(2016) does not perform well on the Auto dataset which is why we have not included it in the comparison.
We also compared test log-likelihood and imputation error which show very similar trends.
In Figure 3, we show samples generated from the learned distribution p ( y ) for all methods.
On the top of this plot, we show the data with true labels.
These labels were not used during training.
The plots (a)-(c) show results on Pinwheel dataset, while plots (d)-(e) shows results on the Auto dataset.
For the Auto dataset, each label corresponds to the number of cylinders present in a car.
We observe [Figure] that SAN can learn meaningful partition of the data space, while VAE does not posses this capability.
We also observe that SAN and VAE can learn ﬂexible patterns while GMM does not do so.
These results show that with SAN we are able to combine complementary characteristics of VAE and
GMM.
For both datasets, our method learns ﬂexible representations containing meaningful structural information.
D ISCUSSION AND C ONCLUSION
We propose a message-passing algorithm for the models that use both deep networks and probabilistic graphical models.
We discussed the challenges involved when designing an algorithm to exploit the structural properties.
Our algorithm can automatically adjust its messages and adapt it to the type of the node.
For the DNN part, the messages reduce to computation of stochastic gradients, while on the PGM part, the messages are equal to natural gradients.
Experimental comparisons conﬁrm that our method gives similar results to existing methods, while using simpler updates.
Our experimental results in this paper are limited to small scale data.
We found that it is nontrivial to implement a message-passing framework that goes well with the deep learning framework.
Our current implementation is based on AutoGrad and TensorFlow, both of which are slow due to additional communication overhead due to messages passing.
We are going to pursue this direction in the future and investigate good platforms to integrate the capabilities of these two different ﬂavors of algorithms.
A R ESULTS FOR L ATENT L INEAR D YNAMICAL S YSTEM
In this experiment, we apply our SAN algorithm to the latent LDS discussed in Section 3.
For comparison, we compare our method, Structured Variational Auto-Encoder (SVAE)
(Johnson et al., 2016), and LDS on the Dot dataset used in Johnson et al.
Our results show that our method achieves comparable performance to SVAE.
For LDS, we perform batch learning for all model parameters using the EM algorithm.
For SVAE and SAN, we perform mini-batch updates for all model parameters.
We use the same neutral network architecture as in Johnson et al.
(2016), which contains two hidden layers with tanh activation function.
We repeat our experiments 10 times and measure model performance in terms of the following mean absolute error for τ -steps ahead prediction.
The error measures the absolute difference between the ground truth and the generative outputs by averaging across generated results.
[Equation]
(17) where N is the number of testing time series with T time steps, d is the dimensionality of observation y , and observation y ∗ t + τ,n denotes the ground-truth at time step t + τ .
[Figure]
From Figure 4, we can observe that our method performs as good as SVAE and outperforms LDS.
Our method is slightly robust than SVAE.
In Figure 5, there are generated images obtained from all methods.
From Figure 5, we also see that our method performs as good as SAVE and is able to recover the ground-truth observation.
[Figure]