Unbiased learning of deep generative models with structured discrete representations
Authors: Harry Bendekgey
Gabriel Hope
Erik B. Sudderth
Venue: NeurIPS 2023 poster (2023)
Keywords: Generative Models
Graphical Models
Variational Inference
Amortized Inference

Abstract: By composing graphical models with deep learning architectures, we learn generative models with the strengths of both frameworks. The structured variational autoencoder (SVAE) inherits structure and interpretability from graphical models, and flexible likelihoods for high-dimensional data from deep learning, but poses substantial optimization challenges.  We propose novel algorithms for learning SVAEs, and are the first to demonstrate the SVAE's ability to handle multimodal uncertainty when data is missing by incorporating discrete latent variables.  Our memory-efficient implicit differentiation scheme makes the SVAE tractable to learn via gradient descent, while demonstrating robustness to incomplete optimization. To more rapidly learn accurate graphical model parameters, we derive a method for computing natural gradients without manual derivations, which avoids biases found in prior work.  These optimization innovations enable the first comparisons of the SVAE to state-of-the-art time series models, where the SVAE performs competitively while learning interpretable and structured discrete data representations.

Unbiased Learning of Deep Generative Models with Structured Discrete Representations
Harry Bendekgey, Gabriel Hope, Erik B. Sudderth
(hbendekg, hopej, sudderth)@uci Suuceuuuiueueeuue edu Department of Computer Science, University of California, Irvine
Abstract
By composing graphical models with deep learning architectures, we learn gen- erative models with the strengths of both frameworks. The structured variational autoencoder (SVAE) inherits structure and interpretability from graphical models, and flexible likelihoods for high-dimensional data from deep learning, but poses substantial optimization challenges. We propose novel algorithms for learning SVAFs, and are the first to demonstrate the SVAE's ability to handle multimodal uncertainty when data is missing by incorporating discrete latent variables. Our memory-efficient implicit differentiation scheme makes the SVAE tractable to learn via gradient descent, while demonstrating robustness to incomplete optimization. To more rapidly learn accurate graphical model parameters, we derive method for computing natural gradients without manual derivations, which avoids biases found in prior work. These optimization innovations enable the first comparisons of the SVAE to tate-of-theeart time series models, where the SVAE performs com- petitively while learning interpretable and structured discrete data representations.
Introduction
Advances in deep learning have dramatically increased the expressivity of machine learning models at great cost their interpretability. This trade-off can be in deep generative models that produce remarkably accurate synthetic data, but often fail to illuminate the data's underlying factors of variation, and cannot easily incorporate domain knowledge. The structured variational autoencoder (SVAE, Johnson et al. (29]) aims to elegantly address these issues by combining probabilistic graphical models [62] with the VAE [33], gaining both flexibility and interpretability. But since its 2016 introduction, SVAES have seen few applications because their expressivity leads to optimization challenges. This work proposes three key fixes that enable efficient training of general SVAES.
SVAE inference requires iterative optimization [62, 20] of variational parameters for latent vari- ables associated with every observation. Johnson al. [29] backpropagate gradients through this multi-stage optimization, incurring prohibitive memory cost. We resolve this issue via an implicit differentiation scheme that shows empirical robustness even when inference has not fully converged.
Prior work [29] also identifies natural gradients [2, 23] as an important accelerator of optimization convergence, but apply natural gradients in manner that requires dropping parts of the SVAE loss, yielding biased learning updates. We instead derive unbiased natural gradient updates that are easily and efficiently implemented for any SVAE model via automatic differentiation.
Basic VAEs require carefully tuned continuous relaxations [27, 44] for discrete latent variables, but SVAEs utilize them seamlessly. We incorporate adaptive variational inference algorithms [24, 25] to robustly avoid local optima when learning SVAEs with discrete structure, enabling data clustering. SVAE inference casily accommodates missing data, leading to accurate and multimodal imputations. We further improve training speed by generalizing prior work on parallel Kalman smoothers [56].
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Poing
Taking phoee
predictiun
Stattes
predictive
effen
The SVAF-SIDS segments each sequence of human motion, which we display as sequence of discrete colors. Discrete variables interpretable: Below each segmentation, we show five segmentations other subjects performing the same action, noting similarity across semantically similar series. Discrete variables are compact representations: Drawing multiple samples from the generative model conditioned on ground-truth segmentations yields the stick figures in grey, which track closely with the observed data.
Figure I:
We begin in Sec. and 3 by linking variational inference in graphical models and VAEs. Our optimization innovations (implicit differentiation in Sec. 4, unbiased natural gradients in Sec. 5, variational inference advances in Sec. 6) then enable SVAE models to efficiently trained to their full potential. Although SVAEs may incorporate any latent graphical structure, we focus temporal data. In Sec. 8, we are the first to compare SVAE performance to and wransformer-baee architectures time series benchmarks [21], and the first demonstrate that SVAEs provide a principled method for multimodal interpolation missing data.
recurrent neural network-
2 Background: Graphical Models and Variational Inference
We learn generative models that produce complex data æ via lower-dimensional latent variables z. The distribution p(=10) is defined by graphical model (as in Fig. 2) with parameters 0, and is processed by (deep) neural network with weights to compute the data likelihood p., (z|2).
Exact evaluation simulation of the posterior p. (z, 0|z) is intractable due the neural network likelihood. Variational inference (VI [62]) defines family approximate posteriors, and finds the distribution that best matches the true posterior by optimizing the evidence lower bound (ELBO):
p(0)p(z (x12))
Cla(O; ()q(;))()))))) 7] [((0=))((=:w)
log q(0; (z)
S log (x).
(I)
Here, q(0; n)g(2; w) = (2, 0|z)) are known as variational factors. We parameterize these distribu- tions via arbitrary exponential families with natural parameters 7, w. This implies that
((z;w) = exp (ws, ((2)) log (ww),
Z(w) =
dz.
(2)
An exponential family is log-linear in its sufficient statistics t(2), where the normalizing constant Z (ww) ensures it is a proper distribution (see App. C.1 for properties of exponential families). For models where P., (xr|z) has restricted conjugate form (rather than deep neural network), we can maximize Eq. (1) by alternating optimization of 1), w; these coordinate ascent updates have closed form [62]. Stochastic V/ [23] improves scalability (for models with exponential-famie likelihoods) by sampling batches of data r, fitting locally-opting (22:)) ) to the latent variables in that batch, and updating q(:: (9:n) 17) by the resulting (natural) gradient.
Amortized VI. Because it is costly to optimize Eq. (1) with respect to w for each batch of data, VAEs employ amortized VI [32, 47, 53] to approximate the parameters of the optimal q(2; w) via a neural network encoding of r. The inference network weights ø for this approximate posterior (=/=x) are jointly trained with the generative model. potentially substantial amortization gap exists [14, 36]: the inference network does not globally optimize the ELBO of Eq. (1) for all
Structured VAEs. For a fixed q()), the true optimizer (across all probability distributions) of Eq. (1) is given by ((2) x po(z; p)(2;1))p-(r (xl=)) where po(2; 17) exp(Eyteen [logg/==]))) p(=(0)}} is expected prior on (see App. C.2 for derivations). This simple product of expected prior and likelihood cannot be normalized because of the neural network parameterization of p. (x)2)). Rather than approximating

---

Algorithm 1
Structured Mean Field Variational Inference
Require: graphical model potentials v). network poten- tials À = encode(;; d) defining (s(22))
(11.....5 (MM) e- init, expactaddstataan A) while converged do for m + 1 to M do
Wm (M(M-m)
BP( A) return concat( (Xx)
Line Sysie (1u1s)
Switching Lamer Dynerkol Systems (SU0))
syliiemo
Figure 2: Left: Generative (above) and inference (below) graphical models for SVAE-LDS and EVAE-SLDS For the SLDS, we show the prior and posterior as factor graphs [38]. (=(=])) combines potentials from the inference network with the true prior, Structured variational inference separates continuous from discrete latent variables for tractability, and mean field messages propagated across residual edges between disentangled factors. Right: General form of iterative mean-field (MF) and peliaa-propagatio (BP) updates for SVAE inference in model where (()) factorizes into groups of latent variables. For the SVAE-SLDS, M
this whole posterior in amortized VI, the SVAE [29] only approximates the likelihood function, and explicitly multiplies it by the known expected prior to obtain an approximate posterior.
In more detail, the SVAE approximates the likelihood E(z) Pr (zl2) with function parameterized
by neural network (=(x)) The optimal posterior given this approximation is then equal
p00)p(=10)
(0(2|<;1)) arg max ((q n)g(2), o] arg max (g(0:))(=)))
log
(3)
g(0; ())
((2)
(())
The posterior that optimizes the surrogate loss L): o) of Eq. (3) is go (2(22;))) x Po(2; 1) (x]x)) If
(=/2r) is chosen to be conjugate to p(=)0), this multiplication and normalization is easy for many exponential families. Note that (=(xx) does not need to be normalized, as any multiplicative factors would disappear when the posterior normalized. The overall ELBO is then L(g(0: 1)) (2 (z; m). 7)-
If the encoder's approximate likelihood l is (up normalization) close to the true likelihood f. the surrogate loss 2 will closely approximate the true loss L, and there will be little amortization gap. SVAE inference has the advantage that (2(2;;;))) depends on the learned posterior q(0: n)) of graphical model parameters, so as learning improves the graphical generative model, the coupled inference model remains closely aligned with the generative process. The ladder VAE [57] and related hierarchical VAŻS [59, 12] also incorporate generative parameters in amortized variational inference, but impose a restricted generative hierarchy that is less flexible and interpretable than the SVAE.
Example Standard Normal. For a basic VAE, is fixed and z ~ (O, 1). The SVAE inference network outputs a Gaussian (= = su(;; Φ), T = ((r; x)) x (22;; diag(r-1)) with mean and inverse-variance (precision) т. The product of this Gaussian with the standard normal prior has simple form: (2(x)=) = (2; + μt, diag((7 + 1)) This reparameterization of the standard VAE posterior imposes the (useful) constraint that posterior variances must be smaller than prior variances. Example 2: Linear Dynamical System (LDS). A LDS model for temporal data assumes the latent state variables evolve linearly with Gaussian noise: 2, ~ ((A2 1 + bt, Qt), 21 A (11, E,). In this case, we expand the exponential family distribution ((2:w) as sum across time steps:
T
T
(<(:w) = 9((1)
II ((()1<11) = exp (W1, t(21)) + ((ve, (we, (()-11 log [(w) (4) t=2
1=2
where = concat(ux) and the prior p(z10) belongs to this family. The prior induces temporal dependence between the zt, but we assume the likelihood factorizes as (x(z)) = II. p, (zrlzt).
For models like the LDS, approximating the likelihood function with conjugate, time- inneeindependent Gaussian distributions is a much simpler task than approximating the comporally-coupled posterior. In addition, as the generative parameters A,, b,, Qe of p(z(0) are learned through optimization of q()), the inference routine shares those parameters, improving accuracy. These advantages were not present in the standard normal VAF, where the prior on lacks structure and is not learned. Inference, normalization, and sampling of (= xx; 7) in the LDS model is feasible via Kalman smoothing algorithm [4] that efficiently (with cost linear in T') aggregates information across time steps.
3

---

3 Structured Variational Inference
For complex graphical models, the distributions p(=|0) and g(2) typically factorize across subsets of the latent variables 2, as illustrated in Fig. 2. We thus generalize Eq. (4) by partitioning into local variable groups, and representing the dependencies between them via a set of factors F:
((2; w) exp
Σ (zy)) JEF
log 3(w)).
(5)
For certain factor graphs we can efficiently compute marginals and draw samples via the belief propagation (BP) algorithm [48, 38]. However, exact inference is intractable for many important graphical models, making it impossible compute marginals or normalize the qs (=|2;;; defined in Sec. SVAE training addresses this challenge via structured variational inference [20, 63, 62], which optimizes the surrogate loss across a restricted family of tractable distributions. We connect structured VI to SVAEs in this section, and provide detailed proofs in App. C.2.
3.1 Background: Block Coordinate Ascent for Mean Field Variational Inference
Let {am JMM be partition of the variables in the graphical model, chosen so that inference within each 2m is tractable. We infer factorized (approximate) marginals (&m 77) for each mean field cluster by maximizing L[q((;;) n)) IIm (<mm), øj. The optimal qu g0(zm|z; 17) inherit the structure of the joint optimizer 9o (2(2|a;; 77). replacing any factors which cross cluster boundaries with factorized approximations (see Fig. 2). The optimal parameters for these disentangled factors are a linear function of the expected statistics of clusters connected m via residual edges. These expectations in turn depend on their clusters' parameters, defining a stationary condition for the optimal wt
Wm MF 1)),
Mm BP 77, d,
(6)
Here, BP is a belief propagation algorithm which computes expected statistics am for cluster m. and the linear mean field function MF updates parameters of cluster given the expectations of other clusters μ-m along residual edges. We solve this optimization problem via the block updating
coordinate ascent in Alg. 1, which is guaranteed to converge to a local optimum of L(g(0) IIm {(zm)]-
3.2 Reparameterization and Discrete Latent Variables
While optimizing (2m x 7) at inference time requires some computational overhead, it allows us to bypass the typical obstacles to training VAEs with discrete latent variables. To learn the parameters of the inference network, conventional VAE training backpropagates through samples of latent variables via a smooth reparameterization [32], which is impossible for discrete variables. Many alternatives either produce biased gradients [6] or extremely sigh-variance gradient estimates [50, 28). Continuous relaxations of discrete variables [44, 27, 7] produce biased approximations of the true discrete ELBO, and are sensitive annealing schedules for temperature hyperparameters.
SVAE training only requires reparameterized samples of those latent variables which are direct inputs to the generative network Py (2] By restricting these inputs continuous variables, and using other discrete latent variables to capture their dependencies, discrete variables are marginalized via structured VI without any need for biased relaxations. With a slight abuse of notation, we will denote continuous variables in z by 2m, and discrete variables by kim
Example 3: Gaussian Mixture. Consider a generalized VAE where the state is sampled from mixture model: l ~ Cat(T), 2 ~ ([ Σx). The likelihood (x(2) directly conditions on only the continuous latent variable Variational inference produces disentangled factors go (2 rx; n)gs (Re)e; m)) and we evaluate likelihoods by decoding samples from (x 77) without sampling go (k|z))
Example 4: Switching Linear Dynamical System (SLDS). Consider set of discrete states which evolve according to a Markov chain k ~ Cat(To), b; ~ Cat(Tk,_1). .) and a continuous state evolving according to switching linear dynamics: 20 N( Σn), 2 ~ ((AA,, Æt-1 + bR, Qk.). The transition matrix, offset, and noise at step depends on k. Exact inference SLDS is intractable [39], but structured VI [20] lears partially factorized posterior qs (2|2; n) (k|z;; (x;;)) 7) that exactly captures dependencies within the continuous and discrete Markov chains.
BP for SLDS uses variational extensions [5, 4] of the Kalman smoother to compute means and variances of continuous states, and forward-back backaadd message-passing to compute marginals of discrete states (see App. D). Let kūz if the SLDS is in discrete state at time t, kij 0 otherwise.

---

Time gradient (ms)
Method
B 1
B = 32
B = 64
B 128
Implicit + Parallel
603
922
1290
2060
Unrolled + Parallel
659
1080
n/a
nía
Implicit + Sequentil
2560
3160
3290
3530
Unrolled Sequential
2600
3290
3980
n/a
Table 1: Time of ELBO backpropagation in an SVAE-SIDS with K 50 discrete states, dimension D 16, and T 250 time steps. For varying batch sizes B, we compare capped implicit gradients to unrolled gradients for L = 10 block updates of two inference algorithms: standard sequential BP, and our parallel extension. For large batch sizes, unrolled gradients crashed because it attempled to allocate more than 48GB of GPU memory.
[0y] be the expected (natural) parameters of the LDS for discrete state j. Structured VI updates the natural parameters of discrete states Wk,, and continuous states Wat follows:
and f = (m))
Σ
Bglkes]0,.
Wkis = (0), Bg (22))1,2))
Stable and Memory-Efficien Learning via Implicit Gradients
When (=(2z)) computed via closed-form inference, gradients of the SVAE ELBO may be obtained via automatic differentiation. This requires backpropagating through the encoder and decoder networks, as well as through reparameterized sampling z ~ (=(2)2; 1)) from the variational posterior.
For more complex models where structured VI approximations are required, gradients of the loss become difficult to compute because we must backpropagate through Alg. For the SLDS this unrolled gradient computation must backpropagate through repeated application of the Kalman smoother and discrete BP, which often has prohibitive memory cost (see Table 1).
We instead apply the implicit function theorem (IFT [34]) to compute implicit gradients dw Ou without storing intermediate states. We focus gradients with respect to v) for compactness, but gradients with respect to & are computed similarly. Let ( (1) be the sequence of w values produced during the "forward" pass of block coordinate ascent, where (4) are the optimized structured VI parameters. The IFT expresses gradients via the solution of of linear equations: (w(L) On Here we apply the BP and MF updates in parallel for all variable blocks m, rather than sequentially as in Eq. (6). At VI fixed point, these parallel updates leave parameters unchanged and g(w) 0.
Onj as
(gg(w; ( ow
Ag(w; 7 Φ,
g(w)
MBP n, o, .): (8)
On
For an SLDS with latent dimension and discrete states, w has O(K + D²) parameters at each 0g is thus a matrix with O + K)) rows/columns and O D2 K) time step. Over T time steps, 8 non-zero elements. For even moderate-sized models, this is infeasible to explicitly construct solve.
We numerically solve Eq. (8) via Richardson iteration [54, 64] that repeatedly evaluates matrix- vector products AJ- to solve A Such numerical methods have been previously used for other tasks, like hyperparameter optimization [43] and meta-learning [49], but not for the training of SVAEs. The resulting algorithm resembles unrolled gradient estimation, but we repeatedly backpropagate through updates at the endpoint of optimization instead of along the optimization trajectory.
(w(l)
(g(w(t); xx) )) Ow
3 "g((w()) 0)
Richardson:
(9)
Om
On
(
0g(w((; n S0 ))))
0g(w<0 n 2,
Unrolled:
= Σ
(10)
On
Ow
On
1=0 i=1
Lorraine al. [43] tune the number of Richardson steps balance speed and accuracy. However, there is another reason to limit the number of iterations: when the forward pass is not iterated until convergence, (4) is not a stationary point of g(w) and therefore Eq. (9) is not guaranteed to converge as / oc. For batch learning, waiting for all VI routines to converge (local) optimum might be prohibitively slow, we might halt VI before (L) converges numerical precision.
Seeking robustness even when the forward pass has not converged, we propose a capped implicit gradient estimator that runs one Richardson iteration for every step of forward optimization, that
5

---

-XIE
Laxs
Unitowed Na1
Marioral de tedee Post Soga
Marirraf wpd.aatitnnarra......... Rast Stage
Step
Figure 3: We compare implicit gradient estimators' stability (left, middle), and gradient conditioning methods' loss trajectory (right), on human motion capture data (Sec. 8). Stability: Gradient estimate rMSE relative to the No-Solve estimator (smaller is better) for various numbers of VI block updates L, and SVAE-SLDS models taken from the start of training (left) and after 20 epochs (middle). Solid lines show median rMSE ratio across batch of 128 data points, and dashed lines show 90th percentiles. Conditioning: Convergence of SVAE-LDS negative-El ELO loss versus number of optimization steps (log-scalo) for conventional nionnnatural) gradients, biased natural gradients [29], and unbiased natural gradients computed via automatic differentiation (Sec. 5).
L. In this regime, implicit gradient computation has a one-to-one correspondence unrolled gradient computation, while requiring a small fraction of the memory. This can be thought of as form of gradient regularization: if we take very few steps in the forward pass, we should have low confidence in the optimality of our end-point and compute fewer terms of the Neumann series (9).
Experiments. In Fig. 3 (left, middle) we compute the accuracy of different approximate gradient estimators for training a SVAE-SLDS as in Fig. 2. To our knowledge, we are the first to investigate the quality of implicit gradients evaluated away from an optimum, and we compare our capped implicir proposal to other gradient estimators. Ground truth gradients are computed as the implicit gradient at the optimum, and we compare the root-mean- root-mean-squuree--ee squared-error (rMSE) of various gradient estimators to that of the naïve No-Solve solution, which replaces the inverted matrix in Eq. (8) with Identity.
We consider two models: with randomly initialized parameters, and one that has been trained for 20 epochs. The newly-initialized model requires more forward steps for the block updating routine converge. We compare the memory-intensive unrolled estimator (Unrolled) three versions of the implicit gradient estimator. First, an uncapped version (Implicit) always performs J = 50 Richardson iterations regardless of the number of forward iterations, thus incurring high computation time. Note that evaluating implicit gradient far from an optimum can produce high error; in the newly-initialized model, many of these iterations diverge to infinity when fewer than 20 forward steps are taken. Second, we consider capped implicit estimator (Implicit+ Cap) which sets J = L to match the number of forward steps. Finally, we consider a capped implicit estimator which also includes a threshold (Implicit+Cap- Cap+T'rrshh:: if the forward pass has not converged in the specified number of steps, the thresholded estimator simply returns the No-Solve solution. This gradient is stable in all regimes while retaining desirable asymptotic properties [64]. Our remaining experiments therefore use this method for computing gradients for SVAE training.
Prior work. Johnson et al. [29] consider implicit differentiation, but only very narrowly. They derive implicit gradients by hand in cases (like the LDS) where exact inference tractable, so the linear solve in Eq. (8) cancels with other terms, and gradients may be evaluated via standard automatic differentiation. For models requiring structured VI (like the SLDS). [29] instead computes unrolled gradients for inference network weights Φ suffering high memory overhead. They compute neither unrolled nor implicit gradients with respect to generative model parameters 77: in practice they set the gradient of the inner optimization to O, yielding biased training signal. Our innovations instead enable memory nemory-efficient and unbiased gradient estimates for all parameters, for all graphical models.
5 Rapid Learning via Unbiased Natural Gradients
SVAE training must optimize the parameters of probability distributions. Gradient descent implicitly
uses Euclidean distance as its notion of distance between parameter vectors, which is often a poor indicator of the divergence between two distributions. The natural gradient [2] resolves this issue by rescaling the gradient by the Fisher information matrix F,, of q(:: 17). given by:
Fy = Ey(e.n (V)g(0:))) (0 (V,g(0:)))
(1I)

---

Johnson et al. [29] demonstrate the advantages of natural gradients for the SVAE, drawing parallels to the natural gradients of stochastic VI (SVI [23]). SVI extends the variational EM algorithm to mini-batch learning: similar to the SVAE, it fits q(=) in an inner optimization loop and learns g(0: ()) in an outer loop by natural gradient descent. The key difference between SVT and the SVAE that SVI's inner optimization is done with respect to the true loss function L, whereas the SVAE uses surrogate L. SVI can only do this inner optimization by restricting all distributions to be conjugate exponential family members, giving up the flexibility provided by neural networks in the SVAE.
Let M, be the expected sufficient statistics of q(0; 1)) Exponential family theory tells us that du By
Fi
[30, 45], allowing Johnson et al. [29] to derive the natural gradients of the SVAE loss:
correction term
updale
2C F 0L 0u On
dc
dc
OC
Ow
3z On Fr 9 az.²
7/0 + E (alemm) [M=))
(12)
3p
dw On
This gradient differs from the SVI gradient by the final term: because SVI's inner loop optimizes w with respect to the true loss f., 0c = 0 for conjugate models. Johnson al. [29] train their SVAE by dropping the correction term and optimizing via the SVI update equation, yielding biased gradients.
There are two challenges to computing unbiased gradients in the SVAE. First, in the structured mean field case aw involves computing an implicit or unrolled gradient, as addressed by our numerical By methods in Sec. 4. Second, including the correction term in the gradient costs us a desirable property of the SVI natural gradient: for step size less than 1, any constraints on the distribution's 's natural parameters are guaranteed to be preserved, such as positivity or positive-definitensss.
We resolve this issue by reparameterizing 1) into an unconstrained space, and computing natural gradients with respect to those new parameters. Letting ij be an unconstrained reparameterization of 17, such as 17 = Softplus [ii} = log(1 + e²)) for positive precision parameter, we have:
ac Dn
OC on -T = -(0)) (a
T
(13)
R 0p
Onj
On
See App. F for proof. This differs from the non-natural gradient in two ways. First, the Jacobian of the " + y map is dropped, as before. Unlike Johnson et al. [29], we do not hand-derive the solution; we employ straight-throug gradieni estimator [6] to replace this Jacobian with the identity. Then, the Jacobian of the r) + 1j map is replaced by its inverse- transss-. This new gradient can be computed without any matrix arithmetic by noting that the inverse of Jacobian is the Jacobian of the inverse function. Thus Eq. (13) can be computed by replacing the reverse-mode backpropagation through the
v) - 1) map with forward-mode differentiation through the inverse - fj map.
In Fig. 3 (right) we show the performance benefits of our novel unbiased natural gradients with stochastic gradient descent, compared to regular gradients with an Adam optimizer [31], and stochastic gradient descent via biased natural gradients [29] that drop the correction term. Results are shown for an SVAE-LDS model whose pre-trained encoder and decoder are fixed.
6 Adapting Graphical Model Innovations
Efficient implementations of BP inference, parameter initializations that avoid poor local optima, and principled handling of missing data are well-established advantages of the graphical model framework. We incorporate all of these to make SVAE training more efficient and robust.
Parallel inference. The BP algorithm processes temporal data sequentially, making it poorly suited for large-scale learning of SVAEs on modern GPUs Särkkä & Farnnd-Fernández [56] developed method to parallelize the usually-sequentie
Kalman smoother algorithm across time for jointly Gaus- sian data. Their algorithm is not directly applicable to our VI setting where we take expectations over q(0) instead of having fixed parameters 0. but we derive analogous parallelization of variational BP in App. D.2. We demonstrate large speeds gains from this adaptation in Table 1.
Initialization. Poor initialization of discrete clusters can cause SLDS training to collapse to single discrete state. This problem becomes worse when the graphical model is trained on the output of a neural network encoder, which when untrained produces outputs which do not capture meaningful statistics of the high- nigh-dimensional data. We therefore propose a three-stage training routine: a basic VAE is trained to initialize py (x]22,, and then the output of the corresponding inference network is used for variational learning of graphical model parameters [25]. Once the deep network and graphical model are sensibly initialized, we refine them via joint optimization while avoiding collapse. For details of this initialization scheme, see App. A.5.
7

---

265 Conerizzn
116V Sanastoon
Des Straaye Thespp Beveraion
DC's Coneent Gereritenn
Figure 4: Unconstrained generation of 513-dim. speech spectrogram data T 500 time-steps (horizontal; models are trained on data with 7 50). An example sequence of real speech data is shown. For models which use discrete latent variables, the sequence of discrete states is shown as a changing colorbar beneath the generation, with solid colorbar meaning constant discrete state for the entire sequence.
Missing data. The structure provided by the SVAE graphical model allows us to solve marginal and conditional inference queries not seen at training time. In particular, we explore the ability of a trained SVAE to impute data that is missing for an extended interval of time. By simply setting (z |mas v)) to be uniform at a particular timestep 4, our posterior estimate of z is only guided by the prior, which aggregates information across time to produce smooth estimate of the posterior on zt.
While discriminative methods may be explicitly trained to impute time series, we use imputation performance measure of generative model quality, so do not compare to these approaches. Unlike discriminative methods, SVAE imputation does not require training data with aligned missing-ness.
7 Related Work
Dynamical VAEs. Girin et al. [21] provide a comprehensive survey of dynamical VAEs (DVAEs) for time series data, which use recurrent neural networks to model temporal dependencies. The Stochastic Recurrent Neural Network (SRNN [17]), which has similar structure to the Variational Recurrent Neural Network (VRNN [13]), the highest-performing model in their survey; models data via one-step-ahead prediction, producing probabilities p(zz|2, 21-1). This model therefore reconstructs x using more information than is encoded in by skipping over the latent state and directly connecting ground truth to reconstruction, reducing the problem of sequence generation to series of very-local one-step predictions. On the other hand, the Deep Kalman Smoother (DKS [35]) extension of the Deep Kalman Filter [37] is the best- cest-performin model which generates observations independently across time, given only information stored in the latent encoding z.
RNNs lack principled options for handling missing data. Heuristics such as constructing a dummy neural network input of all-zeros for unobserved time steps, interpolating with exponentially decayed observations
[11], effectively require training to learn these imputation heuristics. RNNs must thus be trained on missing-ness ness that similar test missing-ness, unlike the SVAE.
Transformers [61] have achieved state-of-the-ar generative performance on sequential language modeling. However, Zeng al. [66] argue that their cermutation-invariann results in weak per- formance for time-series data where each observation carries low semantic meaning. Unlike text, many time series models are characterized by their temporal dynamics rather than collection of tokens. Lin et al. [42] propose a dynamical VAE with encoder q( ( decoder p(#2+1 [ii,t, (1:++1) and latent dynamics p(ze+1 | Z1:e) parameterized by transformers.
Structured VAEs. We, in Johnson et al. [29], only consider SVAEs where the inference network output factorizes across (temporal) latent variables. Orthogonal to our contributions, Yu et [65] investigate the advantages of taking the SVAE beyond this restriction, and building models where the recognition network outputs proxy-likelihon functions on groups of latent variables.
In recent independent work, Zhao & Linderman [67] also revisit the capabilities of the SVAE. However, their work differs from ours in a few key respects. First, because their experiments are restricted to the LDS graphical model (which requires no mean field factorization nor block updating), they do not need implicit differentiation, and do not explore the capacity of the SVAE include discrete latent variables. Second, because they optimize point-estimates of parameters instead of variational factors g(),, they do not make use of natural gradients. In this point-estimate formulation,
8

---

(SVAE-SLDS)
(SIN-SLDS)
(SRNN)
Figure 5: Interpolations of human motion capture data. Red figures (covering 150 time steps) are generated by each model to interpolate between the black figures four times. We see that our SVAE-SLDS provides varied and plausible imputations with corresponding segmentations (colors shared with Fig. 1). During training the SIN-SLDS [41] collapses to only a single discrete state, and thus cannot produce diverse imputations. The SRNN [17] produces varied sequences, but autoregressive generation is sometimes unstable and unrealistie, and its inability to account for future observations prevents smooth interpolation with the observed sequence end.
they apply the parallel Kalman smoother [56] off-the-shelf, whereas we derive novel extension for our variational setting. Finally, their experimental results are confined toy and synthetic data sets.
The most directly comparable model to the SVAE is the Stochastic Inference Network (SIN [41]), which employs graphical model prior p(=(0) but estimates q(2) through traditional amortized inference; a parameterized function that shares no parameters with the graphical model produces variational posteriors. The authors consider discrete latent variable models like the SLDS, but due to the intractability of discrete reparameterization, their inference routine fits the continuous latent variables with a vanilla LDS. Thus training pushes the SIN to reconstruct, and therefore model, the data without the use of switching states. (Experiments in [41] consider only the LDS, not the SLDS.)
The Graphical Generative Adversarial Network [40] integrates graphical models with GANs for structured data. Experiments in [40] solely used image data; we compared to their implementation, but it performed poorly on our time-series data, generating unrecognizable samples with huge FID.
Vector Quantization. Vector- cucioo--uatizeed VAEs (VQ- VAQ- [60]) are another family of deep gen- erative models that make use of discrete latent representations. Rather than directly outputting an estimated posterior, the encoder of a VQ-VAE outputs point in an embedding space which is quantized to one of learnable quantization points. The encoder trained using a straight-through estimator [6]. While discrete SVAE-SIDS states switch among multiple continuous modes, the VQ- VAE representation is purely discrete, limiting information encoded by a latent variable to log, bits. In order to generate plausible and diverse samples, VAEVAEs require large values of K, structured collections discrete variables, and/or post-hoc training of autoregressive priors [51, 52, 16, 15].
Experiments
We compare models via their test likelihoods, the quality of generated data, and the quality of interpolations. We consider joint positions from human motion capture data (MOCAP [10. 26]) and audio spectrograms from recordings of people reading Wall Street Journal headlines (WSJO [19]); Table 2. MOCAP has 84-dimensional data and training sequences of length T = 250. WSJO has 513-diimessoona data and training sequences of length = 50. See App. A.2 for further details.
Generation quality is judged via modified Frechét inception distance (FID (18)) metric. We re- place the Inception V3 network with appropriate classifiers for motion capture and speech data (see App. A.3). SVAE-SLDS-Bia runs the SVAE presented by Johnson et al. [29], with unrolled gradi- ents, dropped correction term, sequential Kalman smoother, and no pre-training scheme. We match encoder-decoder architectures for all SVAE and SIN models using small networks (about 100,000 total parameters for motion data). The DKS, SRNN, and Transformer DVAE have approximately 300,000, 500,000, and 1.4 million parameters each; App. A.5 for details.
To demonstrate the SVAE's capacity to handle discrete latent variables in principled manner, we compare to two DVAE baselines which incorporate discrete variables via biased gradients: the straight-through estimator [6] and the concrete (Gumbel-softmal distribution [27, 44]. To our knowledge, one has successfully integrated either method into dynamical VAEs for temporal data. Thus to make comparison possible, we have devised a new model which adds discrete latent variables to the generative process of the DKS. Specifications for this model is provided in App. A.4. We evaluated this model with both gradient estimators, and reported the results in Table 2.

---

Interpolaticn FiDs (4)
Method
1(2) 2 (T)
Sample FID (L)
0.0-0.8
0.2-1.0
0.2-0.8
Human Motion Capture (h3.6m)
SVAE-SLDS
2.39
12.3 + 0.2
7.9 + 0.2
7.5 0.2
2.8 2 0.02
SVAE-SLDS-BEE [29]
2.36
34.6 0.7
28.8 0.2
25.8 0.3
6.71
SVAE-LDS
34.0
19.3 0.2
21.9
7.90
SIN-SLDS [41]
0.12
8.97
SIN-IDS [41]
2.33
65.2
18.3 0.2
15.5 0.5...2.
6.24 0.09
Transformer [42]
2.82
421 +
234
5
113:5
SRNN [17]
2.94
62.7
0.7
43.5 0.7
84.2
0.3
DKS [35]
2.31
136 1: 1
46.7 1 1.7
33.3 d: 1.1
9.0 0.3
OKneCCneeete
1.70
144 3
88
89 2
34.0 1.4
DKS+Straight-Thoon
22.0
17.15 + 0.14
13.8
Audio Spectrogram (WSJO)
SVAE
0.15
8.14
4.88
SVAE-SLDS-Bia
0.2
15.2
0.12
SVAE-LDS
1.56
19.1 0.3
17.9 = 0.2
16.6 4 0.3
7.2 1 0.3
SIN-SLDS
1.53
20.0
17.2 0.3
1..
9.5
SIN-LDS
1.54
17.8 0.2
17.21 = 0.11
13.2
0.4 0.4
10.1
Transformer
1.88
1.2
12.0 0.3
8.2
5.4
SRNN
1.94
23.6 0.3
0.5
1.3
1..7
DKS
1.55
12.9
10.8
02
10.8
7.7
DKS-Concrete
12.8
SKKS-Straighh-Thoough
1.48
15.51
0.133
10.07
9.02
6.29
Table 2: Comparison of model performance
ligelikd (higher is better), FIDs of unconditionally generated samples (lower better), and FIDs of interpolations on augmented human motion capture and audio spectrogram data. Each interpolation column corresponds to a masking regime where the shown range of percentiles of the data is masked, e.g. 0.0-0.8 means the first 80% of time steps are masked.
Interpretability. In Fig. we show the SVAE-SLDS's learned discrete encoding of several joint tracking sequences. While "sitting" sequences dominated by single dynamic mode, walking is governed by cyclic rotation of states. "Posing" and "Taking photo" contain many discrete modalities which are shared with other actions. The discrete sequences provide easily-readable insight into sequences, while also compactly encoding the high-dimensional data.
Generation. In Fig. 4 show example generated sequences from each model of audio data. Like true speech, the SVAE-SLD. moves between discrete structures over time representing individual sounds. In contrast, the SIN-SIDS [41] and ConcCene baselines collapse single discrete modality, blending together continuous dynamics. While the SKKa+SSraghhttThhree collapse, uses discrete states coarsely to inform the high-frequency dynamics of speech.
model does not
Interpolation. Amortized VI cannot easily infer ((z)) at time steps where the observations a, are missing. Thus, given observations
subset of times Tobss we can encode to obtain g( and infer the missing latent variables by drawing from the generative prior pl ((Zmissing| Robs, 0). Because baseline models parameterize p(=(0) ()) by one-directional neural networks, they can only condition Amissing Zobs at previous time steps, leading to discontinuity at the end of the missing window. For further specifications and for details of the SVAE approach described in Sec. 6, see App. A.6. An alternative approach of in-filling missing data with zeros causes models to reconstruct the zeros; see App. Fig. 7.
In Fig. 5, we see the SVAE-SLDS uses discrete states to sample variable interpolations, while the SRNN's one-step-ahead prediction scheme cannot incorporate future information imputation, producing discontinuities. We also note that despite achieving the highest test likelihoods, the SRNN produces some degenerate sequences when we iterate next-step prediction, and has inferior FID (see Table 2). The SIN-SLDS collapses to single discrete state in training, resulting in uniform imputations that lack diversity. Example imputations for all models are provided in App. Fig.
Transformers for time series. The cormutation-invaiinn performance on these two tasks. lack of sequential modeling can lead to discontinuities in the data sequence which are naturally present in speech. For MOCAP, joint locations are continuous across time, making
of transformers is visible in its varied
samples unrealistic (see Fig. and Table 2).
Discussion
The SVAE is uniquely situated at the intersection flexible, high-dimensional modeling and inter- pretable data clustering, enabling models which both generate data and help us understand it. Our optimization innovations leverage automatic differentiation for broad applicability, and provide the foundation for learning SVAEs with rich, non-temporal graphical structure in other domains.
10

---

Acknowledgements
This research supported in part by NSF RI Award No. IIS-1816365, ONR Award No. N00014-23-1- 2712, and the HPI Research Center in Machine Learning and Data Science at UC Irvine.
References
[1] Akiba, T., Sano, S., Yanase, T., Ohta, T., and Koyama, M. Optuna: A next-generation hyper- parameter optimization framework. Proceedings of the 25th ACM SIGKDD International Conference Knowledge Discovery and Data Mining, 2019.
[2] Amari, S.-I. Natural gradient works efficiently in learning. Neural computation, 10(2):251-276. 1998.
[3] Ba, L., Kiros, J. R., and Hinton, E. Layer normalization. arXiv preprint arXiv:1607.064433 2016.
[4] Barber, and Chiappa, Unified inference for variational Bayesian linear Gaussian state-space models. In Advances in Neural Information Processing Systems, 2007.
[5] Beal, M. J. Variational algorithms for approximate Bayesian inference. University College London, 2003.
[6] Bengio, Y. Léonard, N., and Courville, A. Estimating propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3332 2013.
[7] Berliner, A., Rotman, G., Adi, Y., Reichart, R., and Hazan, T. Learning discrete structured variational auto-encoder using natural evolution strategies. In International Conference on Learning Representations, 2022.
[8] Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: gigggut.......... com/ comm/oogle//e
[9] Cao, H., Cooper, D. G., Keutmann, M. K., Gur, R. C., Nenkova, A., and Verma, R. CREMA-D: Crowd-sourced emotional multimodal actors dataset. IEEE transactions on affective computing, 2014.
[10] Catalin lonescu, Fuxin Li, C. S. Latent structured models for human pose estimation. In International Conference on Computer Vision, 2011.
[11] Che, Z., Purushotham, S., Cho, K., Sontag, D., and Liu, Y. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1 8(1):6085, 2018.
[12] Child, R. Very deep VAEs generalize autoregressive models and can outperform them on images. In International Conference on Learning Representations, 2021.
[13] Chung, J., Kastner, K., Dinh, L., Goel, K., Courville, A. C., and Bengio, Y. recurrent latent variable model for sequential data. In Advances in Neural Information Processing Systems, 2015.
[14] Cremer, C., Li, X., and Duvenaud, D. Inference suboptimality in variational autoencoders. In International Conference on Machine Learning, 2018
[15] Dhariwal, P., Jun, H., Payne, C., Kim, J. W., Radford, A., and Sutskever, Jukebox: A generative model for music. arXiv preprint arXiv:2005.0034/ 2020.
[16] Dicleman, S., van den Oord, A., and Simonyan, K. The challenge of realistic music generation: modelling audio at scale. In Advances in Neural Information Processing Systems, 2018.
[17] Fraccaro, M., Sønderby, S. K., Paquet, U., and Winther, Sequential neural models with stochastic layers. In Advances in Neural Information Processing Systems, 2016.

---

[18] Fréchet, M. Sur la distance de deux lois de probabilité. In Annales de I'ISUP, volume 6, pp. 183-198, 1957.
[19] Garofolo, J., Graff, D., Paul, D., and Pallett, D. Csr-i (wsj0) sennheiser 1dc93s6b. Web Download. Philadelphia: Linguistic Data Consortium, 1993.
[20] Ghahramani, Z. and Hinton, G. E. Variational learning for switching space-spaee models. Neural Computation, 12:831 1-164, 2000.
[21] Girin, L., Leglaive, S., Bie, X., Diard, J., Hueber, T., and Mlaedd--innedi X. Dynamical variational autoencoders: A comprehensive review. Foundations and Trends in Machine Learning, 15(1-2):1-175, 2021.
[22] Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint 1606.0445 2016.
[23] Hoffman, M. D., Blei, M.. Wang, C., and Paisley, J. Stochastic variational inference. Journal of Machine Learning Research, 2013.
[24] Hughes, M. C. and Sudderth, E. Memoized online variational inference for dirichlet process mixture models. In Advances Neural Information Processing Systems. 2013.
[25] Hughes, M. C., Stephenson, T., and Sudderth, E. Scalable adaptation of state complexity for nonparametric hidden markov models. In Advances in Neural Information Processing Systems, 2015.
[26] Ionescu, C., Papava, D., Olaru, V., and Sminchisescu, C. Human3.6mm Large scale datasets and predictive methods for 3d human sensing in natural environments. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2014.
[27] Jang, E., Gu, S., and Poole, B. Categorical reparameterization with Gumbel-softmax. In International Conference on Learning Representations, 2017.
[28] Ji, G., Sujono, D., and Sudderth, E. B. Marginalized stochastic natural gradients for black-box variational inference. In International Conference on Machine Learning, 2021.
[29] Johnson, M. J., Duvenaud, D. K., Wiltschko, A., Adams, R. P., and Datta, R. Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems, 2016.
[30] Khan, M. E. and Nielsen, D. Fast yet simple natural- natural-gradient descent for variational inference in complex models. In 2018 International Symposium on Information Theory and Its Applications (ISITA), 2018.
[31] Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
[32] Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014.
[33] Kingma, D. and Welling, An introduction to variational autoencoders. Foundations and Trends in Machine Learning, 334):307-392, 2019.
[34] Krantz, S. G. and Parks, H. R. The implicit function theorem: history, theory, and applications. Springer Science & Business Media, 2002.
[35] Krishnan, R., Shalit, U., and Sontag, D. Structured inference networks for nonlinear state space models. In Proceedings of the AAAI Conference on Artificial Intelligence, 2017.
[36] Krishnan, R., Liang. D., and Hoffman, M. On the challenges of learning with inference networks on sparse, digh--imennional data. In AISTATS, pp. 143-151, 2018.
[37] Krishnan, R. G., Shalit, U., and Sontag, D. Deep kalman filters. arXiv preprint 1rXiv.151..5512, 2015.
12

---

[38] Kschischang, F. R., Frey, B. J., and Loeliger. H.-A Factor graphs and the sum sum-product algorithm. IEEE Trans. Info. Theory, 47(2):498 519, February 2001.
[39]
Lerner, U. and Parr, R. Inference in hybrid networks: Theoretical limits and practical algorithms. In Conference on Uncertainty in Artificial Intelligence, 2001.
[40] Li, C., Welling, M., Zhu, J., and Zhang, B. Graphical generative adversarial networks. In Advances in Neural Information Processing Systems, 2018.
[41] Lin, W., Hubacher, N., and Khan, M. E. Variational message passing with structured inference networks. In International Conference on Learning Representations, 2018.
[42] Lin, X., Bie, X., Leglaive, S., Girin, L., and Alameda-Pineda. X. Speech modeling with a hierarchical transformer dynamical vac. In IEEE International Conference Acoustics, Speech and Signal Processing, 2023.
[43] Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics, 2020.
[44] Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations, 2017.
[45] Malagò, L., Matteucci, M., and Pistone, G. Towards the geometry of estimation of distribution algorithms based on the exponential family. In Proceedings of the 11th workshop proceedings on Foundations of genetic algorithms, pp. 230-242, 2011.
[46] Marcel, S. and Rodriguez, Y. Torcchissio the machine-vision package of torch. In Proceedings of the 18th ACM international conference on Multimedia, 2010.
[47] Mnih, A. and Gregor, K. Neural variational inference and learning in belief networks. In International Conference Machine Learning, 2014.
[48] Pearl, J. Probabilistic Reasoning in Intelligent Systems. Morgan Kaufman, San Mateo, 1988.
[49] Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems, 2019.
[50] Ranganath, R., Gerrish, S., and Blei, D. Black Box Variational Inference. In Proceedings of the Seventeenth International Conference on Artificial Intelligence and Statistics, 2014.
[51] Rasul, K., Park, Y.-J., Ramström, M. N., and Kim, K.-M. Vq-ar: Vector quantized autoregressive probabilistic time series forecasting. arXiv preprint arXiv:2205. 15894, 2022.
[52] Razavi, A., Van den Oord, A., and Vinyals, O. Generating diverse high-fidelity images with vq-vae-2. In Advances in Neural Information Processing Systems, 2019.
[53] Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, 2014.
[54] Richardson, I.. F. Ix. the approximate arithmetical solution by finite differences of physical problems involving differential equations, with an application to the stresses masonry dam. Philosophical Transactions of the Royal Society of London. Series A, Containing Pupers of Mathematical or Physical Character, 210(459-470::307-337 1911.
[55] Rybkin, O., Daniilidis, K., and Levine, S. Simple and effective vae training with calibrated decoders. In International Conference on Machine Learning. 2021.
[56] Särkkä, S. and Marcfa-Fernánder A. F. Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control, 2020.
[57] Sønderby, C. K., Raiko, T., Maalge, L., Sønderby, S. K., and Winther, O. Ladder variational autoencoders. In Advances in Neural Information Processing Systems, 2016.
[58] Szegedy, C., Vanhoucke, V., loffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. In Proceedings of the IEEE conference on computer vision and pattern recognition, 2016.
13

---

[59] Vahdat, A. and Kautz, J. NVAE: A deep hierarchical variational autoencoder. In Advances in Neural Information Processing Systems, 2020.
[60] Van Den Oord, A., Vinyals, O., et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
[61] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.., Gomez, A. N., Kaiser, E.., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
[62] Wainwright, M. J., Jordan, M. I., al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1, 2008.
[63] Xing, E. P., Jordan. M. I., and Russell, S. A generalized mean field algorithm for variational inference in exponential families. Conference on Uncertainty in Artificial Intelligence, 2003.
[64] Young, D. M. Iterative Solution of Large Linear Systems. Academic Press, New York, 1971.
[65] Yu, C.,, Soulat, H., Burgess, N., and Sahani, M. Amortised inference in structured generative models with explaining away. In Advances in Neural Information Processing Systems, 2022.
[66] Zeng, A., Chen, M., Zhang, I.., and Xu, Q- Are transformers effective for time series forecasting? In Proceedings of the AAAI conference on artificial intelligence, 2023.
[67] Zhao. Y. and Linderman, S. W. Revisiting structured variational autoencoders. In International Conference on Machine Learning, 2023.
14

---

A Experimental Protocol
A.1
Implementation
Code can be found at httpp///gggttub...o////beeeee/// implemented with the JAX library [8].
All methods were
A.2 MOCAP Dataset
Our motion capture experiments conducted using a variant of the H3.6M dataset [26]. The original data consists of high-resolution video data of 11 actors performing 17 different scenarios, captured with 4 calibrated cameras. Following the procedure of [21] our tests use the extracted skeletal tracks, which contain 3D-coordinates for 32 different joints (for 96 observation dimensions), sampled at 25hz. Joint positions are recorded in meters relative the central pelvis joint. Our only changes from the pre-processing of [21] are (i) we remove dimensions with 0 variance across data sequences, resulting in 84 observation dimensions; (ii) we extract 250-step-long sequences instead of 50-step-long ones, corresponding to an increase from 2 seconds of recording to 10 seconds; (iii) we add Gaussian noise with variance of Imm to the data to reduce overfitting: and (iv) we employ a different train-valid-test split. In total, our training dataset consisted of 53,443 sequences, our validation set contained 2,752 sequences, and our test set contained 25,893 sequences.
While [21] split across subjects, we found this lead to substantially different training, validation, and testing distributions. We instead split across "sub-acts", or the 2 times each subject did each action. Thus the training set contained 7 actors doing 17 actions once, and the validation and test combined contained the second occurrence of each act. The validation set contained 3 subjects' second repetition of each action, and the test set contained the other 4. Two corrupted sequences were found in the validation set. 7-phoning-1 and 7-discussion-l, which were removed.
We model the likelihood p(z[2) )) of joint coordinates using independent Gaussian distributions. Following the work of [55] fit single global variance parameter each feature dimension (as part of training), rather than fixing the likelihood variances or outputting them from the decoder network. We found this leads to more stable training for most methods.
FID Scores The Inception V3 [58] network typically used to compute Frechet Inception Distance [18] scores was trained on natural images and therefore inappropriate for evaluating the generative performance of the H3.6M motion capture data. To resolve this, we compute the same score using a different underlying network. We trained convolutional network to predict which of the 15 different motions was being performed in each sequence (note that in H3.6M, motions are distinct from actions/scenerios, which may include more subtle variations). We used a residual network architecture, where each residual block consists of 1-D convolutional layers with batch normalization and a skip connection. We chose I-D convolutions (across time) as they are more appropriate for our temporal motion capture data. We used the Optuna library [1] to search across model depths, activations and layer sizes to find a network within our design space that performed optimally on a validation dataset (25% of the original training data). Our final network architecture is summarized in table 3. As with the original FID score, we use the output of final average pooling layer to compute the score.
A.3
WSJO Dataset
The WSJO dataset [19] contains recordings of people reading news from the Wall Street Journal. We followed all pre-processing of [21], which resulted in sequences of length 50 with 3i3-dimmnsi observations at each time step. We modeled the spectrogram data via Gamma distribution with fixed concentration parameter 2, equivalent to modeling the complex Fourier transform of the data with mean O complex normal distribution. We remove subject 442 from the test data set, as several of the recordings are substantially different from all others in the train and validation sets: instead of being isolated audios, some of these recordings include background noises and beeps which cannot be fit at training time. The encoder takes as input the log of the data and the decoder outputs the log rate of the gamma distribution, as the data spans many orders of magnitude. In total, our training dataset contained 93,215 sequences, our validation set contained 2,752 sequences, and our test set contained 4,709 sequences.
15

---

(SVAE-SLDS)
(SIN-SLDS)
(SRNN)
(SVAE-LDS)
SVAE-SLDS-biag
(Trnsforrmrr)
(DKS+Concrete)
(DKS+Straigh-TTrool
(DKS)
Figure 6: Interpolations of human motion capture data. Red figures (covering time steps) are generated by each model to interpolate between the black figures four times (see App. A.6 for methodology). Segmentation colorbars are shown below models which use discrete latent variables; all non-SVAE-SLA models collapse to single discrete state throughout training. Amortized VI models draw missing data conditioned previous data, introducing discontinuities the end the missing window.
Mocap Classifier
Residual Block
Downsampling block
Activation: Relu
Conv. window:
Cony.
Conv.
Residual block 128
Batch Norm
Batch Norm
Downsampling block 128
Activation
Activation
Residual block 128
Conv.
Conv. (stride
Residual block 128
Batch Norm
Batch Norm
Global average pool
Activation
Activation
Dense
Add input
Add downsampled input
Table 3: Summary of classifier network architecture used compute FID scores for Mocap data.
FID Scores As with the H3.6M dataset, it is necessary to define different base network to evaluate generative FID scores for the WSJO dataset. As the WSJ0 dataset does not have an appropriate set of labels to use to train a classifier, we chose to instead train classifier on the CREMA-D dataset [9], which consists of similar speech clips. We preprocessed each audio sequence in CREMA-D identically to the WSJO dataset and trained classifier to predict one of 6 different emotions that was expressed. For architecture, we used e ciiiiiiettt....mmmmmmmmmmm v2_s model from the torchvisi package [46] on the data in log space.
16

---

Fall Sequence
SVAR-SL DS
SWAE-SL DS-8Bis
DKS Izero
SHAN Ozeno
Continuous Latent Variable
Ifram
SANN (from
Discrete Labent Varia
Time
Figure 7: Segmentation and Imputation. An SLDS SVAE trained synthetic data with feature dimension 100, sequence length
and distinet dynamics. show SVAE's latent representation pictured sequence via the means of the 2-dimensional continuous latent variable and the probability distributions over 8 discrete states (which has coalesced its posterior probability into a single state at nearly every time step). In the following columns and SRNN, we show one imputation using the infill-with- weith-zeros method, and imputations drawn from the learned latent variable prior, conditioned on the preceding latent variables for which a posterior is properly defined.
show model imputations when time steps 20 to 70 not provided. For the DKS
A.4 Discrete Deep Kalman Smoother Baselines
To our knowledge, time series VAFs that utilize Concrete [44] or straight-through gradient estimators have not been previously evaluated in literature. In order perform fair comparison these approaches to gradient estimation with discrete latent variables, we have devised a new model that extends the Deep Kalman Smoother (DKS) [37] incorporate both discrete and continuous latent variables. We can define the generative model as follows:
Pa (kig 2z k
Cat()e,, (21-1,
(14)
po. (ze 21-1, ke) = -N(po, (i),, 06.)
(15)
(x 2x) = (x),
(16)
1, k2) and pay (2,) are neural networks, while dg, and dg are learnable variance parameters. Ma, (2x) uses the same decoder architecture as our implementation of the standard
Here λo, (0 Me, (2
DKS model. λo (z ki- 1), ne. (2+11, kt) use the following architecture in our experiments:
X(-)
((-)
Layer Norm
Layer Norm
Dense 64
Dense 64
Gelu
Gelu
Layer Norm
Layer Norm Dense 64
Dense 64
Gelu
Gelu
Layer Norm
Layer Norm
Dense
Dense
Softmax
The inference model can be written as:
Qdv (ke Iu. 2 kr-1) Cat(Aa, (To(Ett) 21,
(17)
go. (zx Ter, 11 k.) = Nuo, (4-
(18)
Here 9o (IET) is same encoder network used by our DKS implementation (see Table 4), while Ao (-) and Mor (.) are networks with the same architecture as Ao, (-) and μoo (-) shown above. We note that the output of (x(T)) is shared by both parts of the inference model.
17

---

In our experiments we match the latent dimension and number of discrete states to the corresponding Straight-Thoouhh
SVAE SLDS model. For "DKS Straight Through" we use trraiigh-throuuee estimates the gradient of samples of k, In this case we can directly compute the KI--divergence between po, (ki z 1i) and dox (k2 TET, 1t-1, in closed form.
For "DKS+ + Concrete" we replace the Categorical distribution in equations 14 and 17 with Concrete distribution. As the KL-divergence between concrete distributions does not have closed-form expression, we estimate the KL-divergence between the two distributions by sampling. During training. we anneal the temperature parameter for each distribution from to 0.1.
A.5 Architectures and Training Specifications
All experiments use latent space with dimension D = 16. We train using the Adam [31] optimizer for neural network parameters, and stochastic natural gradient descent for graphical model parameters, with batch size of B = 128 and learning rate 10 3 (the transformer DVAE uses learning rate 10 10-4, which improved performance). We train all methods for 200 epochs MOCAP and 100 epochs on WSJO (including VAE-pre-triining for 10 epochs for SVAE/SIN methods), which we found was sufficient for convergence.
For the SVAE-SLDS, we initialize the graphical model parameters after these 10 epochs of VAE pre-training of the networks. We found that the discrete latent variables could not meaningfully cluster the continuous latent variable at the beginning of training, when the encoder outputs do not yet contain high mutual information with the observations. After VAE pre-training, sampled 100 sequences and encoded them. We then normalized the resulting approximate likelihood functions and sampled to obtain plausible sequences in the latent space. We trained an auto-regressive HMM model on the resulting sequences using memoized variational inference, which uses adaptive proposals to avoid local optima implemented in bnpy [24, 25], to initialize the parameters r) of (()).
Table 4 summarizes the network architectures used, numbers of parameters, and compute needed to run the main experiments of this paper. Note that "Rev. LSTM" denotes a reverse-order LSTM and "TF-LSTM" refers to the "leacher-forcing" LSTT that is shared between the encoder and decoder [21]. For the transformer, attention modules in the decoder include causal masks, and all attentions are single-headed. Skip connections are not listen in the table. For a full specification of the architecture, see [42] (Table 4 outlines the architecture choices at each layer of that paper's 's framework).
For each model, list the number of parameters for the MOCAP dataset. Even though the same architecture is used for both models, WSJ0 has a higher data dimension (513 vs 84) resulting in 109,824 more parameters in the encoder (via the first layer) and the decoder (in the last layer). The total amount of computation across both datasets amounts to 6 GPU days for A10G GPUs on an EC2 instance.
A.6 Interpolation Procedure
Let Fmissing; Tabs be a partition of data into missing and observed sequences. For the following explanations, we assume that the observed data is the beginning and end of the sequence Tobs
00.0-0.2 U F.0 (first 20% and last 20% as in our results) leaving Fmissing = -0.2-0.8-
To interpolate missing data DKS/rrnnsformee/SRNNN case) or Eo (onss) (in the SVAE/SIN case), where Zabe is the subset of 2 corresponding to the observed time steps. For models which are separable across time this is trivial, and we handle the LSTMs in the SRNN and DKS encoders by encoding every connected block of observations separately so that we do not need to hand missing or infilled data into the encoders.
begin by encoding Lobs to obtain g(Zobs) (in the
For the DKS and its extensions, we then sample the remaining 2 from the generative prior given the last encoded latent variable that precedes the missing chunk. In the regime above, we therefore sample from 2020.2-0.8 (0.0-0.2 as we cannot easily condition on later timesteps of a forward deep recurrent model. Finally, we decode Emissing U zabs to obtain the reconstruction and interpolation.
For the SRNN and the transformer, due to their autoregressive nature we sample a value of z at single timestep (i.c. the first timestep in 20.2-0.8) and decode to obtain new We then repeat this
18

---

Method
Encoder arch.
Enc. Params Decocder arch.
Dec. Latent params
Dense 256
Dense 128
Gelu [22]
Gelu
16.4 GPU hours
Layer Norm [3]
Layer Norm
30,668
(MOCAP)
SVAE-SIDS Dense 128
57,504
Dense 256
57,640
(50 states)
12.5 GPU hours
Gelu
Gelu
(WSI0)
Layer Norm
Layer Norm
Dense 16
Dense
Dense 256
Dense 128
Gelu
Gelu
9 GPU hours
Layer Norm
Layer Norm
(MOCAP)
SVAE-LDS
Dense 128
57,504
Dense 256
57,640
716
4.6 GPU hours
Gelu
Gelu
(WSJO)
Layer Norm
Layer Norm
Dense 16
Dense
Dense 256
Dense 128
Gelu
Gelu
7.3 GPU (MOCAP)
Layer Norm
Layer Norm
34,018 (50 states)
SIN-SLDS
Dense 128
57,504
Dense 256
57,640
3.6 GPU hours (WSJO)
Gelu
Gelu
Layer Norm
Layer Norm
Dense 16
Dense
Dense 256
Dense 128
Gelu
Gelu
7.3 GPU hours (MOCAP)
Layer Norm
Layer Norm
SIN-LDS
Dense 128
57,504
Dense 256
57,,400
1516
3.6 GPU hours
Gelu
Gelu
(WSJ0)
Layer Norm
Layer Norm Dense
Dense 16
Dense 256 Gelu
Dense 256
Gelu
Layer Norm Positional Enc.
Layer Norm
Positional Enc.
Self Attention
2 GPU hours (MOCAP)
Self Attention
Layer Norm
Transformer
360,480
532,136
518,688
Layer Norm
Cross Attention
1.6 GPU hours (WSJO)
Dense 256
Layer Norm
Gelu Skip
Dense 256
Layer Norm
Gelu
Dense
Layer Norm
Dense
Dense 256 Gelu
Shared LSTM 128
3.7 GPU hours (MOCAP)
Layer Norm Dense 128
Dense 256 Gelu
SRNN
252,288
190,888
87,680
2.9 GPU hours
Shared LSTM 128
Layer Norm Dense
(WSJO)
Reverse-LSTT 128
Dense
Dense 128
Dense 256 Gelu Layer Norm Reverse-LSTM 128
Gelu
3.3 GPU hours (MOCAP)
Layer Norm
DKS
38.144
244,160
Dense 256
57,640
(+ Discrete)
(43,724)
2.4 GPU hours (WSJO)
Gelu
Dense 192
Layer Norm
Dense
Table 4: Summary of network architectures used for human motion capture and speech experiments. "Latent parameters" those which specify p(zl.). For the SRNN model, one L.STM layer shared hetween inference and generation; we denote this layer "shared ISST" count its parameters among those for decoder. -VAE-SL-S-BiE uses the same architecture as the SVAE-SLDS but takes 36.5 GPU hours MOCAP and 28.5 GPU hours on WSJO train.
19

---

process so as to update the hidden state h, while holding 2 in the observed and region constant.
For the SVAE and SIN methods, we do inference in qa (z(2))) as in training time, except setting the inference network potentials at missing timesteps to 0 (thus setting (=(zx) as uniform over the reals). For models with discrete random variables the block updating routine prone to local optima. so we initialize according to the following heuristic: we encode 20.0-0.2 and 20.8-1.0 and perform block updating to obtain ((k..000.2) ) and (k0.8-1.0). We then sample the endpoints ko.2, ko.s from these variational factors. We perform HMM inference sample from p( 2..0..-..sl /.... kp.s), using that sample to initialize the block updating.
The protocols above can be easily extended to the 0.0 0.8, 0.2 - 1.0, and the full generation cases, noting that the models which can only condition generation forwards in time (DKS, SRNN, Transformer) must sample the initial from the generative model in the 0.0 - 0.8 case instead of conditioning on an observation. The discrete random variable models initialize block updating for fully-generated sequences via draws from the prior p(k)).
Note that in Fig. we show the discrete states are discrete representation by initializing block updating with the ground truth discrete states as well as the first and last time steps (see grey background figures). We found the "anchoring" observations were necessary because the discrete states only encoder dynamics, as opposed to location in the z-space.
A.7 Additional Results
In Fig. 6 we show extended version of Fig. 5, showcasing the performance of all models interpolating human motion capture data.
In Fig. we compare model's interpolation ability several baselines synthetic data. The data constructed by evaluated a Laplacian distribution on 100-dimensional grid, where observations evolve over time by either increasing decreasing the mean of the distribution, or increasing or decreasing the variance of the distribution. Note that this data was used to test the runtimes of various methods, see Table 1.
Fig. 7 demonstrates the effect of the imputation method on the baseline models: for RNN-based models, if we simply infill the missing observations with O, the encoder-decoder structure simply reconstructs the Os. Drawing the missing region from the model's prior distribution yields generation closer to the data distribution but still infeasible under the true data dynamics. The prior SVAE implementation collapses to a single state, yielding interpolations that do not match any dynamics seen in the training data.
20

---

B
Notation
In this section we provide a glossary of notation used in the main paper and the appendix.
B.1
Notation in main paper
Notation
Meaning
BP())
belief propagation function, which takes (natural) parameters of the graphical model and returns expected sufficient statistics.
weights of the generative network (decoder)
7
natural parameters of q(O;
1)
unconstrained parameterization of rj
Fnj
Fisher information matrix of exponential family distribution with natural parameters
0 (lw; 1); 0,
equation which block updating attempts solve for the root of discrete local latent variable (subset of 2)
k:
(2(22)
approximate likelihood function outputted by recognition network (encoder)
C
the loss function, the Evidence Lower Bound Objective (ELBO),
C
Surrogate ELBO, obtained by replacing true likelihood function with conjugate surrogate.
MF())
mean field inference update, where messages are passed along residual edges (those edges removed by the structured mean field approximation).
/
expected sufficient statistics of g(2), i.e. E [+2)
Pm
expected sufficient statisties of particular mean field component g(=m) )- (t is defined analogously to w-m-
Hn
expected sufficient statistics of q(0), i.e. Eg [(0)
W
natural parameter of g(2) in any of its forms: ((;;)) (local
VI), go [2; n) (SVAE VI) or II. 9 m [.; r)) (Structured Mean
Field SVAE VT)
Wm
natural parameter of each mean field cluster IIm Iol (m I;; n) in Structured Mean Field SVAE
Um W-m is defined
analogously to _4-m-
W W-m
natural parameter all mean field clusters but m, W-m
UngmmWn
w (x,, (x,1).0)
function which runs block updating to return the optimal natural
parameters of ((z) with respect to surrogate loss funciton L Po(z;1) x exp (Ey(()))) [oz(((=()]]]
Po(2;
weights of recognition network (encoder)
q(0 17,, (2)
variational factors approximating the posteriors of global (0) and
local (z) latent variables.
(= [=)2),, 4s (z
variational factor for produced by amortized VI and SVAE VI respectively.
g m;
variational factor for subset 2m C
t(.)
sufficient statistics of an exponential family distribution. global latent variables defining graphical model factors vertex set of the graphical model
0
V
observation
2
local opeerobbervttion) latent variable
21

---

B.2
Notation in this appendix
Notation
Meaning
Sec.
Exponential families and mean field objectives
KL
Kullback-Leibe divergence, defined in Eq. (25)
F
set factors in the graphical model representation of g(=), intro-
duced in Eq. (32).
{x(x)
natural parameter of (z(2))) defined in Eq. (44).
rt, R
for Gaussian temporal models such as the LDS and SLDS,
\x(x) = {rr, R } as defined in Eq. (50).
arass
Matrix LDS transition parameters for each of the switching
states in the SLDS, defined in Eq. (58).
Sec. D
Belief propagation in time series
ho: Jo
natural parameters of =1 (0), defined in Eq. (64).
We
natural parameters of the LDS (or the continuous half of the SLDS) transition at time step t, defined in Eq. (65). Also used in the HMM the mean field message from the LDS at step t.
11,80 11.. J 11 2 12,5
components of wr, defined in Eq. (65).
M m.,,on (2))
belief propagation messages from variable 2; to factor with pa- rameter W/, introduced in Eq. (69)
Misin (2x)
belief propagation messages from factor with parameter wi to variable 24. introduced in Eq. (70)
Feje, Feje
parameters of ma4 (Zx), i.c. the filtered messages for the
Gaussian inference step & which has accrued evidence from the first t "observations"
S441,,
parameters of M (Z++1) i.e. the predicted messages for the Gaussian inference at step + 1 which has accrued evidence from the first "observations"
Jepr: Far
parameters of q( (1 i.e. the smoothed messages for the Gaussian inference at step f which has accrued evidence from all "observa- tions"
ker(.)
returns an un normozzal distribution with no log partition func- tion, defined in Eq. (111)
[(( 21), 2 ①b
sequence of functions (and their parameters) used by the filtering
step of the parallel Kalman filter, defined in Eq. (115)
(1 1), 7t:
sequence of functions (and their parameters) used by the filtering
step of the parallel Kalman filter, defined in Eq. (112)
e1(2,, 21 21-1), 1,, Fx
sequence functions (and their parameters) used by the smooth- ing step of the parallel Kalman smoother, defined in Eq. (148)
associative operator for parallel Kalman filtering, defined in
Eq. (119) and (120)
@
associative operator for parallel Kalman smoothing, defined in
Eq. (151)
22

---

C Exponential families and mean field objective
To rigorously define the SVAE inference routine, in Sec. C.I we (re- )introduce exponential family and factor graph notation for defining distributions. Then, in Sec. C.2 we derive the optimizer of [[...] and the mean field message passing update MF(.). Sec. D then contains the necessary equations for belief propagation in the LDS and SLDS models, including proofs and derivations of our novel inference algorithm which parallelizes belief propagation across time steps.
C.1 Exponential families of distributions
Given an observation E 2, an exponential family is family of distributions characterized by statistic function ((2)= Z R with probabilities given by
p(2: = exp (17, t(z)) - log (n))
(19)
Where E IR" is parameter vector known as the natural parameters, and the log partition funiitinn log (1) given by
log (11) = log / expl(n; {(z)))dz
(20)
and ensures that the distribution properly normalizes (to simplify notation, here we focus on continu- ous but this definition can be easily extended to discrete observations by replacing the integral with a sum and and restricting the output of ((z) to some discrete set). A single assignment of values to characterizes a member of this exponential family, and searching over a family means searching over valid values of 17 E H where H is the set of parameters for which the integral in Eq. (20) is finite. We begin with a couple simple properties of exponential families:
Proposition C.1. The gradient of a distribution 's log partition function with respect to its parameters its expected statistics:
V log ((1)) = Ep(cm) [())]
(21)
Proof.
J. expl (n, (z)))dde / (xx)))))),,,()))))
1
V, log 7(n)
(22)
f (x))xp((n log ( Ep(em [()]
(23)
12) in the same exponential family with parameters Tn and up, the KL divergence between the two distributions is given by
Proposition C.2. Given two distributions p(z; ), p((
p()
KL(p(z;m) |l p(2;ma))) E:pps:m) log
(24)
p(z:
=(m 72, Ep(zmmp [{(2]]) log ((m) + log (112)
(25)
Proof. The log probabilities are linear in the statistics, so the expectation can be pulled in.
Definition C.3. Given an exponential family distribution p(0; 17)
p(0:n) = exp (n, ((0)) - log (n))
(26)
we say that p(z(0) is conjugate to p(0; 1)) if one distribution's 's parameters is the other distribution's 's statistics, namely
p(z(0)= = (1())
(t(=), c) = expl ((()) - log ({{))))
(27)
Where i(0) is the first dim(!(=)) entries of ((0) and e a constant.
Note that t(f) and t(2) are not the same function with different inputs, but rather are unique to each distribution. As it is generally clear which function is needed, we share this notation, just as p(0), p(2) understood to refer to different densities.
Further note that the sufficient statistics of p(0) either become parameters of p(=]0) or its normalizers. This can be notated two ways: either by appending Is to the end of ((z) to align with the normalizers. or simply fold them into the log partition function. An example is provided below.
23

---

Example: NIW distribution. statistics given by:
A normal distribution p(2; 2) has natural parameters and sufficient
Σ
-
((x)
(28)
We bend notation here by concatenating vectors matrices. Assume all matrices are appropriately flattened to make the inner product a scalar, for more info see Sec. G. The normal distribution' 's log partition function is given by:
1 Meanwhile, the Normal Inverse Wishart (NIW) distribution has sufficient statistics:
log (1) 5" E +
log =
log(2z)
(29)
Σ
Σ
Σ
1(14 x) =
(30)
log Σ
Note that the first two are parameters of p(z), and the last two appear in the log partition function. A desirable property of conjugacy is that the posterior has simple form, namely:
p(0|=) (xp((1(0), (1)=),1))
(31)
This follows easily from the fact that p(0)=) & p(0)p(=10) which are both log-linear in t(0).
Factor Graphs. The exponential family framework can be applied to complicated distributions over many variables, in which case can be helpful to rephrase the distribution as a collection of factors subsets of the variables. We represent the variable such that the vertices partition the components of z. Let the factor set F. be a set of subsets V such that 2 z, F. Then write distribution
a graphical model with vertex set V
p m) = exp (n, {z)) - log ( = exp { Σ (n (2z)))
log [(n)
(32)
JEF
These two ways of writing the distribution are trivially equal by defining as the union of all ny and combining t(z)) appropriately. We demonstrate where this notation can be useful in two examples.
C.2 Mean field Optima
Proposition C.4. Given fixed distribution q(a) and a loss function of the form
f(a,b)
[[g(a)((b)] = (g(()) log q(a)g(b)
(33)
for non-negative, continuous function f(a, b) with finite integral (such un-normalized density), the partially optimal q(b) across all possible distributions given by
arg max (((())((b) x exp(Bgra) (a, b)]}
(34)
s(b)
Proof. This is a simple extension of proof in Johnson et al. [29]. Let p(b) be the distribution which satisfies pi(b) oc exp (T) [(a b)]} We can drop terms in the loss which are constant with respect (bb) and write:
log f(a
arg max (8) ((b)
= arg max Eg [ [log fla 6)] - log g(b)]
(35)
q(b)
q(b)
arg max E (b) (b)) [log exp By(ca) [log (a b)] log {(b)]
(36)
((b)
(xp [log (a, b)]}
arg max E (b) ((b)
log
(37)
(())
log ((b)
arg max E (b) q(b)
+ const
(38)
((b)
arg max (L(((b) |I ((b)) 7(b) a(b)
(39)
24

---

This is applicable to ELBO maximization problems, as well as the surrogate maximization problem arising in the SVAE:
Example: Exact Inference in the SVAE objective, we see:
Applying proposition C.4 to the SVAE's surrogate
log ((0)p(=1)){
arg max (ig(0: Be(Om)a(s)
(40)
((0:1))(=)
q(x)
o exp( [log p(0) p(z]0) (z(2)))
(41)
xx ( /lo (()) ((,,) plog)(000 log (s(22))
(42)
oc exp [(= 0)] log
(43)
Recall that have chosen (x(2(x) be conjugate likelihood function to p(210). Let (x(x) be the outputs of the recognition network such that
(œ(=|=) exp((Ap((t),
(44)
Then
arg max (=())) x
(x(x))
(45)
<(2)
x (x (t(0), (z)
(46)
exp (()m)) (.)))
(47)
Demonstrating that the optimal (2) is in the same exponential family as p(z with parameters w = Eg(0:n) [0) + (z). These parameters are the sum expected parameters from the prior and the parameters of the fake observations.
Example: Linear Dynamical System. Consider the linear dynamical system, where p(=1) NHo, Σo) and ((=(( = -( be, Ex) (in practice, we use a hom-hhomogenuu model in which A, b, E do not depend the time step t). We can write this distribution
p(= p(21) II p(2))=11)) exp exp (1(0)),
Σ (1(0)), 21)) - log ((1(1)) (48)
Where ((z2),((0;) are given by the NIW example in Sec. C.1, and for 2,...,
-A bb
-t-1
1 1AT AA
2t 2--12 7 Zt- 2
(()-1;
((0,)
A Q,
(49)
-1,
21 2t
2t
Q,
Note that both t(0,) and t(0++1) introduce parameters corresponding the statistics 24. In the context of the SVAE, we add observations to each time step.
[s(2|2) exp Σ
([<
T1
2 22
1}}
(50)
which is conjugate the likelihood p(=10) that the field optimum stays in the same exponential family.
Structured mean field. The main paper introduces the concept of dividing the latent variables into separate structured mean field components 2m C 2, m E M. We outline the math in this by noting that conjugacy allows us to break down statistics on groups of vertices into statistics on individual vertices.
t((0),((2-1122))) 2x)) (t(0, 2-),) = (((=-)))
(51)
25

---

Le. this expression is linear in the statistics of each argument. This proves important in deriving the optimal factor g(zm) given the other factors q(= By applying prop C.4 the mean field surrogate objective with q(a) = q(0) II-m ((zm) and ((b) (2m)). we see:
arg max Cla(0; m) II q (g )i o) x exppLeatreeets....lllll
+ log (z/1))
(52)
((am)
The expected-pric part of this expression can be rewritten (where z.. = 2④ n 2m and z..-m
z/ 0 n:))):
E (t0:)))(=...) [oogpppp000) 0) x (o((:))((=_.)) m)
Σ
(t(0)),
(53)
JEF
(x(0:))((=-m))
Σ (t(0g, (100),21.m)) t t(2g,m))
(54)
SEF
Σ E (0 ... .,)(t((,, 25. t m t(z,,m) JEF
(55)
Note that because the joint sufficient statistics of zf,-m, O, are linear in each input, we can distribute the expectations through the statistic function and write
(o((m)))(-_m)) (1,0,, =m]= = MF (Fy
(56)
Where (F(.) is a mean field function which takes the statistics of fy, 2..-m and produces the joint statistics t(O,, mm)) This function is linear in each of its inputs and generally cheap to compute.
Example: Switching Linear Dynamical System. The SLDS consists of the following gener- ative model: p(k²) Cat(To), p(ks) Cat( p(21) = NNoo Eg), p(k2) Cat(Tu) and p(=+==-1) = (A + bb,, Ex.)-
In an exponential family form, the model consists of the following factors: (i) an initial-state factor on 21 matching the NIW example in Eqs. (10-12), (ii) an initial categorical factor on k2 (see Sec. G for the form), (iii transition factor for p(ke which has parameters Ty and sufficient statistics (kk = (k<+1 = })} for j.j = I,..., K, and (iv) transition factors corresponding to p(ze|=e-11 kz), which we will examine further.
Recall from Eq. (49) the expression for the parameters of p(2) 11) in the LDS. Then we can write: p(z0| k2) = exp( (t(x,,)) t(2)1122)) - log Z (((O,,))
(57) We need to rewrite this to be an exponential family distribution of k as well (instead of only appearing as an index). Note that the sufficient statistics of a categorical variable is simply its one-ho hot encoding. If we define t(A {(Ouras) to be a matrix whose kth row is given by t(Ox). ). and log Z (0 to be a vector with components log Z (0x), we write:
(58) Thus the conditional distribution can be described by two factors, one which takes the outer product of the one-hot encoding ((kg) with the continuous observations (thus embedding them in the correct dimension for inner product with t()),, and one which simply takes the correct log normalizer based on what state we're re in (again, by multiplying a one-hot encoding with a vector of log probabilities).
| k2) exp (ttGuams), t(ke) -(2--1, log ((t(auam))) t(kx))
Given observations, exact inference in the SLDS is intractable. Thus we divide the graphical model via structured mean field inference into q(=)g(()) This separates out the joint factors (1(0mams) ((ke) 1( = ). However, as noted in Eq. (51) we can reframe this factor on each mean field cluster:
exp((t((Oran)) (kk)) (21-11 1 expl (1(x 22), ((Otran)T ((kr)))
(59)
= expl (t(ks), (OPrans) 2)))
(60)
Thus, applying Eq. (55), the optimal (k) inherits the (expected) initial-state factors from p(k2) and transition factors from pkk, 1k) as well as the normalizer factors (- log ((((Ora))) (kk)) from
Eq. (58), but replaces the joint factors of Eq. (58) which include 21 with
(By(e:n) [00mmas)] [Bgc))(((== 2t)], {(kx))
(61)
Which (when combined with the normalizer factor) is simply a vector of expected log probabilities that the transition at time t is caused by each state. Similarly, the continuous chain keeps its initial state factor corresponding to p(21) but all of its transition factors are replaced by
(E,():)) [(0raan)) Eg((k) (2)-11))))
(62)
26

---

Because t(kg) is a one-hot encoding of k4, (o())))((k( is simply probability vector of what value ky takes. Thus the new parameter for {( is simply a convex combination of the rows of t(Orans), yielding time-inhomogencous LDS with each transition computed as an expectation over discrete states, Our mean field optimization of the SLDS becomes:
Run a belief propagation algorithm to get (x(e)(((2))
For each f, compute (x(() it ((0rrs))) Bx(+)) (2)-11, as the new parameters g(k) corre- sponding to (kk))
Run a belief propagation algorithm to get (x)())(((h))
For each compute Eyy [(mmms)))] (B(())(((kx)] as the new parameters in ((2) correspond- ing to 2,)
and repeat until convergence.
Belief propagation in time series
D.1 Sequential LDS Inference
Consider again the generative model
N[uo, En), ~ ~r be, Q:)
(63)
Recall from Eq. (47) that the ((z;w) will have some parameters from the expected log prior and some parameters from the recognition potential. We separate out the parameters into individual factors, starting with a prior factor on the initial state:
ho J
w0
t(a1)
(64)
1 Jo are the expected statistics of an NIW prior. Further, we place transition factor on Where ho. 2 each adjacent pair of states, given by Eq. (49). To simplify notation, write
-,i,
ELA 6r]
zezl
JI1.
E[A Qr Ad]
Uz Z+41)
W--
J12,t
E[A Q
(65)
Z 2+11244
2 J22,t
3 [E,
2++1
h22
E E
b.]
For the SLDS, each of these parameters is outputted by the mean field message passing function and a convex combination of the corresponding parameters for each cluster, weighted by Ey(k.) (k:)]] For the LDS, these values are time-homogeneous.
Note that there is also a factor outputted by the recognition network output each time step:
A
T' Re
((zx)
2t
(66)
Because the parameters are expectations, there may be no values of A., Q,, d, which satisfy
- b
hi
2 Q
1 A
1 A
2 JuLe
AT Q-
J
(67)
Q
J22,1
Q b, And as result, it might not be equivalent to any single LDS with fixed parameters. This prevents us from using general Kalman smoother algorithms, families and belief propagation.
h2
must return to the definition of exponential
The goal of inference is to compute the marginals (=x), draw samples, and compute log Z. Using the belief propagation algorithm, we know q( proportional to the product of incoming messages from factors which involve 2₁. In belief propagation. messages pass from factors to variables and back. Thus write
a(2)) Musies (2x) mwet1.as (2x) mxi, (2))
(68)
27

---

Where (2) is the message from the factor with parameter w to variable 2. We similarly define mz (2) be the message from variable to factor. Note that all m., (2,) are provably Gaussian, and
can therefore each message function can be minimally described by the Gaussian natural parameters.
Messages from variables factors are the product of incoming messages, e.g.
ma (z+) = Miw.a (#x) -, (2+)
(69)
While messages from factors variables are computed by integrating over incoming messages times the factor potential, e.g.
S 1a (=c-1) exp( 2+) 22))
Mu (2x)
(70)
We attempt match our notation to the typical Kalman smoother algorithm, by defining filtered messages Jej. haje as the natural parameters of ma (2x), i.e. the distribution which has accu- mulated the first Wia and Aia which correspond the first & ""oseevvtions". We similarly define predicted messages J4t1if / as the natural parameters of Mu (2+1)), which corresponds distribution over 2 given the first È ""bservvtions". Finally we describe the smoothed messages Jej hup to be the natural parameters of g(z,) (note that for simplicity the natural parameters for each distribution are actually J. h.j, and the negative one half is assumed to be factored out).
The message passing algorithm is outlined below:
Input: wi, Az- Initialize:
Fip Jo + Rı
(71)
f1j ho + +ri
(72)
log Z
(73)
Filter:
for 1. 1, ,,...,7 T 1
Predict:
P. Feje + J11,t
(74)
F4t1 = J22,1 J22 P." J12,t
(75)
fet1/e = -h2,1 + J22 P. (fel - h1.1)
(76)
log Z = log Z + 2
(Fip - h1,))² P. (fule h1.e) - log PPl
(77)
Measurement:
F411111 = Fit1i + Re+1
(78)
1441441 = S0t1/e + +r411
(79)
Finalize:
log 7 = =og 7 + PT Prir log Fryr)
(80)
Smoother: for = 1 1....,1
C, = F441/ Pit2 + J22
(81)
Fip F.J, + J11, J12, - GT! J22
(82)
Sur = Sup h1, + J12, (Sttp - set1je + h221)
(83)
Sample: (the A below are taking natural parameters input)
2T ~ N - (fri - Frye)
(84)
1 for = 1, 1....,1 21 ~ No + J12 - h1s (Frle + J12,))
(85)
Return: marginal params fo T. FijT: log Z; and samples æ1.7.
28

---

From the log marginals, we can recover the expected statistics:
Σ; I/T
(86)
H4 Exfer
(87)
[[z]] M4i
(88)
F. 2[ Σ; + phlta
(89)
E [1 = C, + papp+1
(90)
Derivation of filter: The measurement step trivially follows from Eq. (69). The predict step, on the other hand, attempts to compute Eq. (70), which can be written out:
J11.c J12.t
2
J exp {[-
ft Pele
Zx (zzz)
J22,1
Ze4
2141
)
(91)
-1,
Zt
112,1
2141
{[
J222
20+12(+1 2 2
>>
/ exp {{
Jaj
J12,1 h1,4
24
1)
Jdee
exp
(Fepe + J11,e)
2
(92)
J22, {[[
24412+11 Zz41
| + log Z(feje hie + J12,14411 Faje + Jrna)
exp
(93)
h2.
What is that log partition function term the end?
2 (Jep - hi, + J12,22411 )a (Faja + Jn.a))
b1,, + J12 log Fije + Jniel (94)
Simplifying P. = Feje + U11. we recover normalization term
// Fill P, - h11e) 2 log Prl
(95)
Linear term
2
I JT. P hnd))
(96)
And quadratic term
- 1 Te, Pi J12,1
(97)
Thus the entire message is:
(Laat h2g + Fe. (Sele h1.1)
- ( P. h1,1--h1,)000 log Pa (4412741
2+11
exp
- Fa (
(98)
To calculate log Z, instead of normalizing each message to be valid distribution, we can think of normalizing each message so it's of the form exp (n, t(z)) }, i.c. remove all normalizers. Then to calculate the entire model's 's log Z, we must accrue normalizers that appear in predict steps (none appear in measurement steps) and simply add the normalizer for the last message.
Derivation of Smoother: Note that
g(2,) x mwrei. (2t) Ma (2t)
(99)
The second term is the filtered message computed in the filtering step. Thus we must simply compute the backwards messages. The message becomes:
((xx) Mawwe1 (zt) -mi
(100)
x ma (zt) / masliwe (zo- 1)
((z,, Ze+ 1)) }dz41
(101)
= (zt) / (mariizte (Z+11) max ma (#++1) expl 3d)
(102)
g(2(+1) exp(( Mzt, 2++1)) )24)
X Mayeee (2t)
(103)
Mossriser (2141)
29

---

Now let's evaluate the integral:
Jn,
2
/ exp
112
24 2e
XL
furr fet1le
] [ 1
dddee (104)
2 J22,
2++1 2
(P(+//T - Fet11e)] (2412241)
h1,t
2z
h2,
20+1
eexp {
2 J11.
2027
|>//
exp {{[ S1+2 + h2, + H2 (F Fetije + J22,a)
2+41
1) (105)
d
h1,t
21
[20+12++1
T
{{[
I1Le
[022"
=exp
-
+log f1+11r-f++111+ h2, S2, Fitirr Fetije + J22,t)
hi,
21
(106)
Again, let's evaluate the log partition function at the end, simplifying C2 = Fr+1/T Fe+1je + J22,1
1 2 (1+111 fetije + 22, + 222 'F fetije + h2, + FT2.21 log Ccp
(fetiT
(107)
We recover linear term
#f J12 ( Set1je + h22))
(108)
And quadratic term
1 T
2 ct J12, C J2
(109)
Combining these with the J11,t and h1, in Eq. (106) and the filtered distributions outside the integral in Eq. (103), we arrive at the update equations above. Derivation of Sampler: Moving backwards, want to sample from
(20121411) x Maswer (24)
t(ze, =++1)))
(110)
The of the derivation easy.
D.2 Parallel LDS Inference
This adaptation of the Kalman smoother algorithm is highly sequential and therefore inefficient on GPU architectures. We adapt the technique of Särkkä Fernáneeer [56] to parallelize our algorithm across time steps.
Let us define the kernel of (possibly un-normalized) exponential family distribution the inner product portion of the distribution:
ker(exp (n, t(=)) - cc)) exp((n, {()))}
(111)
Note that ker(f(z)) x f2). We define two sequences of functions:
(x(=11) = ker /
exp(( (00)),,,,(---) ((2-1, -
exp ((x))) )
(112)
(J J12, + J22,t) ( -1 M1,t + J12 RR, + J22.,) -1 (r. + h2)
[2 [2-
1)}
exp
(113)
21-1
exp {[[
[ 2-1 [1-2)111)
(114)
7t.
Thus ge 1=(= defined by natural parameters 74. For = we set I, = 00 m1 = there no time-step 0. Thus we claim gi (20) = 1 is uniform function.
30

---

Further define f((2-11221)
fc(2--122) exp((t(0:), (0tt((00),(((1111110 ((=-11, 2))}} exp(
(115)
2 (J1.e [t)
---12-1
(hiee 71)) J12,
24-1
}
exp
21-1 2t
(116)
1221 + +
(J22,e + Re)
2127
2 11.e
22-12211
-1,tt
24-1
exp
112,
1t-12 'P'
(117)
Ft
P222,
24
We similarly define f1 21) = fi (z1). or equivalently 211,1 = 1221 0 and Φ1.1 0 as there
no 2g. 22,1 ho ri and 222,r Jo + Rí to include the factors ho Jo from the initial state distribution p(21).
We define the associative operation:
(fi.g)) @ (fr.g) (fij.g9s)
(118)
Where
c) = ker J 9, (b),f (a, b)db
(119)
J (b) b)db
And
gi5 ker (9.(a) / 9, (b) b)db)
(120)
Proposition D.1. The operation @ is associative.
Proof. Särkkä & García-Fernández [56] define an operator
(f.. g.) x 9f)) = (f.s
(121)
Where
[s) J 99 (b) c),
(122)
I gg (b) b)bb
And
9. (a) 9;(a) / (b)a
(123)
and prove that this operator is associative. Further, any rescaling of inputs leads to a simple rescaling of outputs due the multi-lines nature of these functions. Thus for constants C1 C2, Ca Ca there exist constants c, Gg such that
(2.1) (cafs, (csffy,cciin
(124)
This gives us that
((i, 6 (fz. 9z)) & ((k.ge) (ker(ft;), ker(ār))) )) (Jk,gk)
(125)
= (cIfis, c20is) 8 (sk. 9k)
(126)
ker (ca c20r)) x (Jk, (k))
(127)
ker ((f)) x (fr, (fx.gh)
(128)
A similar decomposition shows that (f & ((f,9)) & (fk, ker (ff..g) (fyk. 95k))
(129)
ker (Jiy (is) x (Ju.
(130)
=(fr.g: ((f),g))) gk)
(131)
Where the second equality comes from the associativity of the operator.
Finally, we need to show that this operation gives us the right results:
31

---

Proposition D.2. Leta; = (fi, 9:). Then
01 02 ... 2 az = (fitt 91:)
(132)
where
fLa(Z)) = exp
(
felt F(ie
2t
(133)
2
is parameterized by the filtered messages Saj, Frlr-
Proof. We prove by induction that this holds when the operations are performed sequentially, namely
(((((a) @ 02) @ dg) 8 ....) @ az
(134)
These "left-justified" operations are simpler, and thus allow us easily demonstrate correctness. We then can make use of the associative property of @ to show that no matter the order of operations we get the correct result.
Recall that @1 (Zo) = 1 is the uniform function (as there is no 20) and thus has 11 = 0, n = 0. By induction, this must be true of all the cumulative kernels 9.y (=o), if g:(a) in Eq. (120) is independent of a, then (a) must be, too. Similarly, we recall that fi (zo: =1) = f1(21) is independent of its and Φ1.1
first argument and thus has 111,1 412,1 cumulative kernels f1=(20, =) must have their first three parameters equal to O; in Eq. (119)
Again, by induction all of the
fi b) is independent of then fig(a, c) will be, too. Thus fiæ(20, 22) f1=(2z).
Next we show that fu.((z)) = (Z), the filtered messages from the sequential filter. For 1, f11(21) = m=1 (21) by construction. We then see:
(Jz, I441 (2)) (241 ( ( (Ze) Ie+1((2) (()
S1a+1(2441) ker
(135)
ker
9e41 (2))
(136)
ker
(411))
mz (2 (2()
(137)
(xpp(AAc1,t(2441)) /
ker
((p((t((0011:))0)00 241)) mz (2) (2))
(138)
= ker
(xpp((A+2, m (ant)))(+1)) = =M (24+1)
(139)
Completing our proof!
This algorithm produces correct filtered messages but due to its associativity can be computed in parallel. To complete the filtering algorithm, we outline how evaluate (fs, (9,19,). The integrals in Eqs. (119), (120) can be evaluated as simple functions of the parameters of f., Jj. 9:. 9j:
Proposition D.3. Let
11,5
aa T
1 p
2
2 - 112J
Φ1,4
{ {{
}}
f((a.b) b) =cxp
4
a
f((b,)) cxx
bcT
d2,i
do,i
bbT
1 P22
coT
(140)
and
{([=
aa
g3(b) = exp { {(1=
66
(.(0) = exp
(141)
Yi
32

---

Then given the following definitions
C [ + 422,
(142)
is = (911.5 122 122,))
(143)
Nij -th, + 12 Cij (73 + 22,5)
(144)
P. Ii + 122, 111s
(145)
We can compute (fijsgis) = (fis 91) @ (fz. -;) as:
- ( 12 P p- 1/ 12, - r(i))
1j
Faa"
(01,1 12, P.j (02. - 11 + +77) + ni)
figla,c) =e0p=xxpp)) exp { {{
}
122 P 122
ac
(146)
2 + 2 P.j (02,
01 +75)
oc
z (22..
$ 12.j
4 12.5)
and
giy exp [[ -
[c
(147)
9: 7
The proof of this involves integrating the kernels of Gaussians, which is done many times in this supplement; it is left as an exercise for the reader.
The last task is to define the parallel smoother. The construction is very similar to the filter. Let
Ma (Ex)
el(2,, 20-1)
2411))
(148)
Muesi (21+1)
2 (J1,t + F(je)
2
(h1, Sea)
=exp
J12.1
(149)
h2.1 + rt+1 11441
R++1
(J22, Re+1 J Je+1/e+1).
24412
2 Fn1
2/4
E1,t
= =xx
E12,
2,2411
(150)
62,t
Ze+1
Ez2,c
20412/+1
Where we let (T 2T+1) = ( with E1,7 = Jryr, - -hrpp and remaining parameters equal (as there is 2T Then the associative operator @ is given by
(a, = e; b) @ (b c) ker (eel(a.)e)
(151)
Which can be computed (derivations omitted) as
Dij E22, + E11.j
(152)
1/4 (E E12, ET2.))
Taa a
(e1, E12,; D (e2, el
Ciy(a, =cxpp) exp {{
E12, E12.j
act
(153)
62. + E4 D (62,5 - (1.)
2 (E22.5
E|2, DJ E125)
Loc
Särkkä & Carefa-Fernánder [56] prove this operation associative when operating on full dis- tributions (i.e. without taking the kernel). The proof that we retain associativity when removing multiplicative constants mirrors the proof Prop. D. and omitted for brevity. Finally, we show that when these operations are done sequentially,
E, + (E:+1 + (e++1 @... (er-1 @ eT))))) q(x))
(154)
The base eT(2) = ((2)) is true by construction. Then, plugging the inductive hypothesis and Eq. (148) into the definition of ( yields Eq. (103), proving that ex(2) = ((2.) implies et-1 (2t-) = q(211).
33

---

D.3 HMM Inference
The discrete Markov model is time homogenous and requires only expected probabilities for the initial state 22 and expected probabilities for the transition matrix Tej. The "observations" of the HMM are (in the case of the SLDS) messages passed from the LDS chain, which we denote wy (k). The setup is fundamentally the same the LDS:
Input: g0) [t()]] Neoog
Forward Pass:
01 (k) Ellog m.. w1 (k))
(155)
for t
at (k) ..... Wy (k: logsumexp K [Ellog Tual + aa)1(t))
(156)
log Z = logsumexp K (k())
(157)
Backward Pass:
(r(k)=0 =0
(158)
for T- B.(k) logsumexp [ (0) Elog Tayl + Bet (7]
(159)
Return: log marginals (k()) + log (k):::
:: log 7
E
Evaluating the SVAE loss
Johnson et al. [29] decompose the loss into three terms
Prior KI.
Il p(0))) + E,io) (L(((2) ll p(z(0)) - Reconstruction
Local KL
(160)
Because p(0), ((: 4s (2|2; 7) are all exponential family distributions, the expectations here can be evaluated in closed form, with the exception of the reconstruction loss EE) (x|[)
The reconstruction loss can be estimated without bias using Monte Carlo samples and the reparameter- ization trick [32]. For this to work, must be able reparameterize the sample of q(2) such that can back propagate through the sampling routine, which is not possible for discrete random variables. As result, the latent variables which need be sampled to reconstruct (i.e. all variables which in the same mean field component as a variable being fed into the encoder) must be continuous.
The first term of the loss, the prior KL, is regularizer on the graphical model parameters and can be easily evaluated via Eq. (25). The rest this section is devoted to computing the KL divergence between the prior and variational factor of the local latent varibles z. We will start with the fully structured case where no mean field separation occurs.
Recalling that the optimal (z)) has parameters w given by Eq. (47), we can again invoke Eq. (25) to get Eg (L(((=) p(2(0))
Eg k(w t(0), t(0), (x))) log (ww) + log ((t())
(161)
(w- (((o)(((0)] Fegts) (=)) log (w) + log (y ((())
(162)
= E(t))(()) - log Z(w) + log (0) [tt))))
(163)
Namely, because log ((2) is defined as the expected log prior p(=[0) plus the recognition potentials, these two distributions only differ in log space by those recognition potentials. The log partition function of the graphical model can be obtained by belief propagation, and the expected log partition function from the prior can be generally known in closed form as the sum of individual partition functions at each vertex or factor in the corresponding graphical model.
Example: LDS The LDS SVAE only outputs recognition potentials on a subset of the statics of (2)). In particular, the encoder takes the observations x and outputs normal likelihood functions on each z independently. Thus the cross-tempora covariance statistics 2 have no corresponding recognition potentials. Further, we restrict the potentials on 2 to be diagonal, meaning our observational uncertainty factorizes across dimensions. Thus the only statisties in ((=) for which
(x) are non-zero are and =ū, for each component i of 1(z). As result, the inner product (x)()), E. (z)) sums over 2T terms for sequence with T timesteps and D-dimensional latent variable at each step.
34

---

Local KL with structured mean field When there is mean field separation. the prior and variational factor differ more than just the recognition potentials, as q(2) replaces some factors of p(z]0) with separable factors. We can rewrite
M (1 = Σ
kL(g(Zm)
(164)
The natural parameters of g(zm) differ from the natural parameters of p(zml Z<m 0) not only in the recognition potentials, but also for any factors which include a variable in "later" clusters 2ym-
Example: SLDS The SLDS SVAE has 2 mean field clusters, (kk)((=) for the discrete vari- ables h and continuous variables 2. we write our prior distribution p(@)p((k|))p<=)k, then KL(g(=) p(= p(= 0)) is computed exactly the same as in the LDS example above, these two distri- butions differ only in the recognition potentials. The new challenge is computing (L(g(k) p()(0) which contains no recognition potentials but differs based on mean field messages from ((2) which are present (k) but not in pkk)()).
In particular, the prior pkk ()) is a time series with no observations, while g(k) includes parameters every time step which encode the log probability of being in that state at that time, given by z.)]. The corresponding sufficient statistics are the marginal probabilities being in state at time Note that we never need to compute the sufficient statistics
Eq. (60): Ey [t(((rmm)) g (1(2-11,
corresponding to HMM transitions because we don't 't need them for the local KL nor do we need them to sample from the HMM (which we never do).
Surrogate Loss The surrogate loss differs from the true loss only in the reconstruction term, so it can similarly be written
Prior KL
Local KL.
-L(a(0))(2), o] = (g(((0) p(O))) + Eycor (L((2) || p= - (eg)()))
(165)
This is trivial evaluate, the local KL contains the negation of the last term (see Eq. (163) for the structured case. The mean field case contains these terms as well in addition to ones corresponding to "pulled apart" factors) and thus the surrogate loss can be computed by omitting the recognition- potential-timest simees-epecteddsttttieeee part of the local KL.
F
Natural gradients
Here we dissect two claims in the main paper more closely. First, we will show how natural gradients with respect to 7) can be easily computed. Defining, as in the main paper, My E.((..)) [()], we can use the chain rule to write:
ac = 0pm On On detach((r,))
ac 0pn +
OC
(166)
On
0c Where dy detach(H.) refers to the gradient of the loss with respect to 17, disconnecting the computa- tional graph at Me (i.e. the dependence of the loss on which is not accounted for by its dependence on (.)).
Proposition F.I.
OC
0
(167)
On detach(p),)
Proof. q(2) depends on only through the expected statistics H4, which define the parameters of the graphical model factors. Further, the local KL, outlined in Eq. (163), only depends q(0) through its expected statistics (noting that the log normalizers of p(=[0) are also staistics of q()). Thus the only place in the loss that v appears its is the prior KI.. Let no be the natural parameters of the prior p(0). Then the prior KI. as per Eq. (25), is given by:
(L(g(0) p(0)) (n 77. mn)) log ((n) log 2(7o)
(168)
35

---

Then we see:
(l (()) OKL(g(0) p(()) 0pm dn
+ Pn
log (n)
(169)
Op
On
dn
(KL(g(0) (()) 0pm
0pm
(1 - m1)
(170)
pp
Orj
Where we use Prop. C.I to cancel terms. Thus we conclude that the entire gradient of the loss with respect 7) passes through Moy
The main consequence of this theorem is that in all cases the gradient can be written
dc dc 0pe
(171)
On Pp Oq
And thus the natural gradient, which multiples the gradient by the inverse of Op By can be easily
computed withoout matrix arithmetic by skipping the portion of gradient computation corresponding to the 1) - My mapping. Here is some sample Jax code which takes function f and returns function which computes the forward pass but in the backward pass treats the gradient as identity:
def straightt.trrough())
def f :
zero = jax. max.lax.stor stop-gradient (x) return zero + jax. lax. wtw.lx..ttoppgrreeeoeee gra.op--grieettCeooo (x))
return straight_through._
f
Fisher information of transformed variables 17, the Fisher information matrix of 1) is given by:
Given a distribution g(0; 17) with natural parameters
① @logg((0)
(Dlog q(0) ]
(0)
(172)
On
On
Note that for an exponential family distribution ē Ôlo (() t(0) [t(e)[((0)]]
(173)
On
which clearly has expectation 0 under (()). Thus
Fú = (0) [(t(0) (o) [(()))) (t(0) En(o) [((O)
(174)
Varg(0) [{))]
(175)
Exponential family theory gives us further equalities:
8²
8
(ar,(o) Ön?
log 7(n)
Ön y (()))
(176)
Thus deriving desired property that F,,
8p δn
Now we consider general case, where " parameters rj and continuous invertible function f. Then we see:
( for some unconstrained, non-natural
log ((0) T Oij ( Oog g(0)
2 (log
]
F.j = Fy(
(177)
Oij
On
a log (())
On
(178)
On
Oîj
On
Om T
① log (())
① [llo
)
On
[Eq(8)
(179)
Jīj
On
On
Tīj
On T
On
Fn
(180)
Dij
We can use this to compute the natural gradient with respect to the unconstrained parameters:
Oc
0μ
(On
Om
ac
Om -T
Ori
T
Fr
F
=
= On
vuC) (181)
Oîj
On
On
0μ
36

---

Implementing natural gradients via automatic differentiation We argued in the main paper that implement natural gradients we replace the reverse-mode mode backpropagation through the - r map with forward-mode differentiation through the inverse operation. Assuming we have forward map and inverse function f _iinv, sample code is shown below:
_natgrad custtmusisssss
def f_natgrad_fwe fwd fud (input): return f finput),, f (input)
def f_natgrad_ bwd (resids, grads): return
(resids,), (grads,)) [1],)
f natgrard defvip (ftiaddffeieeeeeeeeeeeeee (f_natgrad_ [_aaggrad_bvd)
G Matrix-normal exponential family distributions
Recall that exponential family distributions are given by
p(z; 7) = exp (17, t(x)) - log (
(182)
The natural parameters often have coupled constraints, and so can be expressed as a simple function of canonical parameters which carry cestly-interpretable information about the distribution (e.g. the natural parameters of normal distribution can be expressed in terms of the mean and precision matrix of that distribution). Thus define each distribution, we provide (a) the natural parameters as function of canonical parameters, (b) the sufficient statistics t(r), (c) the expected sufficient statistics Ep( [t()), and (d) the log partition function log (()).
Note that sometimes we have symmetric positive definite (SPD) matrix-valued parameters. In every case, the corresponding statisties for those parameters are also symmetric. For example, given n x n. SPD matrix, distribution might have parameters and statistics:
Salt
x T 1 <
7
2 S.g <j
(xx)
(183)
We note that (n. t(x)) Σm= Σj- Sy or equivalently the sum of elements of S xxT (for column vector containing each 24). In the following sections, we will abuse notation for compactness by simply writing:
[S]
t(x)
xx
(184)
Constraints. To optimize these parameters via gradient descent, we must define constrained parameters as function of unconstrained parameters as discussed in Sec. Sec.nnaggadds] Further, to implement natural gradients we must implement the inverse operation. The table below lists all such transformations.
Notation
Size
Constraint
7
h-length vector
n/a
n/a
Rexm
n X m matrix
n/a
n/a
IR
n-length vector
All values positive
11 Softpluse
RT
n-length vector
All values greater than m
= Softplus
A 1
n-length vector
Non-negative, sums to 1
= SSof
S
n matrix symmetric positive definite
See below
Where Softmax made invertible by appending a 0 to the end 7j. For the SPD matrix parameters, we define S = diag(r) R diag(a) where E TR is vector of standard deviations and I is a correlation matrix, transformed and from unconstrained space by the tensorflow Correlation Cholesky bijector.
37

---

G.1 Multivariate Normal distribution
If (p, Σ) for (i) I. μ E R", and (ii) Σ E 4 then p(a; μr Σ) is given by:
Σ
((x) rx T
Ep(xn
(185)
Lap 14p-+3 +
The log partition function is given by:
1 q
1
log ((n) ~²
Σー
Σ
-1
log |
2 log(2
(186)
2
Here, x 3.14 is a constant.
G.2 Normal inverse Wishart (NIW) distribution
If μ. E ~ (SW m, X, v) for (i) E R", (ii) E, 54+, (iii) À E R4, and (iv) Ron-1. then p(u, E; S, m, A, 1/) given by:
S + Ammm 77
1
Am
v$ 3-1mm m
U1, E) =
Bl1 Σ) =
2 (n/ + vm T 5-1
V + n +
log ΣΣ|
1or (n log log |S| + Σnー 4( (
(187)
Where & is the digamma function. The log partition function is given by: (n log 2 log (S]) + log r With the multivariate gamma function I',.
log ((n) 2
2 +
(log(2m) - log X)
(188)
G.3 Matrix normal inverse Wishart (MNIW) distribution
If X, Σ ~ MNIW (S, M, V, v) for (i) X, X.M E Rnxm (ii) E,S E 54+, (iii) V E ST4, and (iv)
R Rya-1. 1, then pXX, E; M, V, is given by:
S + M
1E-1
1
MV
v$ všー
7=
{(1E)= = 1 x
E[((X,))
( + +
+ n m m+
J/ log |
n log 2
log S|| + >ズー v/ (
(189)
Where v is the digamma function. The log partition function is given by: V
(5)
V
n
n.m log(2
2 (n log 2 - log [S]) + log I' With the multivariate gamma function r,- For the SLDS, X E 2xn++ [A]b] where 2++1 ((x + 6, >.) so the MNIW distribution provides prior on the transition matrix, offset, and conditional noise.
log (n)
|
(190)
2
2
this MNIW distribution by defining
38


Reviews:
[{"review_id": "5ce3rLraXr", "reviewer": "Reviewer_JW8q", "summary": "The paper builds upon a previous work (Structured VAEs) which tries to integrate variational auto encoders with probabilistic graphical models resulting in more expressing latent variable models for data. The authors propose a new unbiased estimator of the gradients for optimization the ELBO (for SVAE) and show via experiments that the superior performance of their algorithm. In particular, they show their structured model outperforms transformers and RNN based models on sequence prediction tasks (time-series predictions).", "questions": "See the points made in the weaknesses. If my concerns are adequately addressed in the rebuttal I would be happy to improve my score. Here are few other minor concerns that could improve the readability of the paper. \n\n1. I am not sure I understand the distinction between the true loss and the surrogate loss. Is the point that while optimizing q(z) (the latent variables) we learn an additional inference network \\hat{l}(z | x) which parametrizes p(x | z) as an exponential distribution which is a conjugate to p(z)? Is \\hat{l}(z | x) replaced by p_{\\gamma}(x | z) while optimizing the outer loop of the ELBO with the locally optimal q(z)? Moreover, how is it ensured that \\hat{l}(z | x) will be a good approximation of p_{\\gamma}(x | z).\n\n2. In Fig.2 and Sec 3.1 what are the disentangled factors and residual edges? \n\n3. Lines 166-177 the capped implicit gradient estimator is not clear to me. If the Richardson iteration is run every step of the forward optimization. Does it mean that in every iteration of the inner loop when the \\omega gets updated (by block coordinate updates), the gradient in eq (8) gets computed and the parameter \\eta gets updated?", "limitations": "N/A", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper makes learning SVAEs much more feasible by making the training more efficient due to the unbiased natural gradient estimators proposed by the authors. \n\nThe authors further propose an interesting implicit differentiation method to bypass the need to back propagate through the block coordinate ascent iterations while optimizing q(z) in the inner optimization loop. This would result in tremendous memory gains.\n\nThe paper and the appendix is well-written for most part (see weaknesses).", "weaknesses": "1. The paper is very clear as to its optimization innovations for improving the training to SVAEs but unclear or ambiguous in other aspects. In particular, one of the claimed contributions of this work is the ability to learn discrete latents (it is in the title of the paper). However, only one small section (S3.2) is devoted to this and the contribution is not clear. The way it is written, the solution seems to be simply removing the discrete variables from input to the decoder/inference network. It is not clear why this is not needed in SVAEs? What approximation is being made? \n\n2. The authors claim that continuous relaxations of discrete distributions result in optimization procedures sensitive to the annealing schedule whereas their method is free from any such limitation since they do not need to backpropagate through discrete latents. I think an empirical comparison with VAEs that employ such relaxations like the Concrete VAE and the more popular VQ-VAE (this variant does not use a smooth relaxation of the discrete distribution but instead uses a straight-through estimator) would further strengthen this point. Without a comparison with contemporary methods it is hard to appreciate the benefit of SVAEs over these methods for optimizing over discrete latents.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your compliments and your feedback. We will aim to respond to each comment individually.\n\n1: The capacity of the SVAE to model discrete latent variables (without relaxation or bias) is a property of the SVAE model family, although it was not emphasized in the original paper which proposed the model. Our primary contribution is in improving and modernizing the SVAE so it can be tractably learned, but we argue it is a model family worth improving because of its discrete capabilities, which we demonstrate experimentally. \n\nRegarding the property itself, no approximation (other than the mean field approximation necessary to make inference tractable) is being used to allow for discrete variables; we simply need to specify our model such that the discrete latent variables are not direct inputs to the decoder network. For example, in the SLDS an encoder outputs variational factors on the continuous latent variables (Fig. 2), followed by block updating which estimates the posterior on both the continuous and discrete latent variables. This block updating causes the discrete distribution to impact the continuous variable and therefore the reconstructions, as the posterior on the continuous latent variables approximately marginalizes out the discrete switching states. We then sample continuous variables and decode. If in the Fig. 2 SLDS generative model, we had added $k_t$ as an input to the decoder network, we would have to sample $k$ and be in the standard VAE conundrum.\n\nWe agree that added discussion of this matter would improve the paper and plan to expand section 3.2, including in ways that address your following comment.\n\n2: For comments on the concrete VAE, please see the global response above. We agree that added discussion of the VQ-VAE is important, and plan to add it to discussion of related work in future revisions of the paper. Regarding experimental comparison, to our knowledge no one has successfully integrated either method into dynamical VAEs. \n\nQ1: The surrogate loss is only used to optimize $q(z)$, and the correct loss including $p_{\\gamma}(x|z)$ is used in the outer loop. Because the ELBO is a valid lower bound for any choice of $q(z)$, any routine which produces a $q(z)$ induces a trainable family of models. Therefore, the SVAE’s choice of encoding data to variational factors and optimizing a surrogate loss is simply a reparameterization of $q(z|x)$ compared to the standard amortized setup where an encoder outputs the distribution parameters directly. It is not clear a priori which method which have a smaller amortization gap, but the desirable property of the SVAE is parameter sharing between the generative and inference model, because surrogate optimization at inference time uses $E_{q(\\theta)}p(z|\\theta)$. \n\nWhen $q(z)$ is true loss-optimal, our ELBO becomes a tight bound and our loss improves. In the fully-amortized setup, an encoder with sufficient capacity will eventually output a good estimator of q(z) because doing so will produce the best loss. In the SVAE, $q(z)$ is true loss-optimal when the true loss and surrogate loss look very similar, which in turn only happens when the encoder network outputs a good approximation of $p_{\\gamma}(x|z)$. Just as training guides a standard encoder network to well-approximate the posterior, it guides the SVAE encoder to well-approximate the likelihood. \n\nQ2: In the true SLDS generative model, there is a joint distribution $p(z,k|\\theta)$. However, for tractable inference we are learning a factorized posterior $q(z) q(k)$. The residual edges, shown in red, depict the dependencies that have been removed (i.e. which connect continuous and discrete latent variables) in this factorization, which are replaced by disentangled factors shown in blue.\n\nQ3: We perform L block update steps to compute omega (as part of the forward pass), followed by up to L Richardson iteration steps to compute the gradient of the loss with respect to eta (as part of the backward pass). Although the number of forward and backward steps are the same, they are not interwoven, and eta is only updated after the entire loss computation and backpropagation is complete. "}, {"review_id": "dQKJbFK3Ha", "reviewer": "Reviewer_4krc", "summary": "The paper introduces a novel, efficient learning algorithm for structured variational autoencoder for any graphical model.\nIt uses a block updating algorithm with belief propagation to approximate the posterior and employs implicit differentiation for end-to-end learning that is more memory-efficient than an unrolled alternative.\nOn top of that it derives unbiased natural gradients to improve the convergence of the learning and a few practical tricks such as parallel inference and robust initialization to make the implementation scalable.\nOverall, the proposed algorithm makes SLDS models competitive to state-of-the-art time series models.\n", "questions": "If the straight-through estimator is used, is the natural gradient still unbiased as a whole?\n\nAre the time reported in Table the overall time of the 3-stage training or the number without the initialization (which seems to be expansive)?", "limitations": "Limitations are not discussed in the paper.\nSome of my questions above could be discussed in a dedicated limitation section/paragraph.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "### originality\nThe proposed learning algorithm is novel. \nAlthough each component seem to be well-established, combining them together is non-trivial.\n\n### quality\nThe paper is solid and the method is sound.\n\n### clarity\nThe paper is well-written and easy to follow.\n\n### significance\nThe proposed method makes learning and inference of a wide range of graphical models practical through SVAEs, which is very important to the filed of probabilistic ML.", "weaknesses": "### quality\n\nI think the experiments of the paper can be improved to incorporate more graphical models.\nFor now, the paper only tests the learning method on SLDS even though the method can be applied to a wide range of graphical models.\n\n### clarity\n\nThe arrows for inference and generative networks in Figure 2 are pretty small.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your compliments and questions. We will aim to address your questions and concerns individually.\n\nDespite the SVAE’s flexibility to handle arbitrary graphical models–one of its most exciting properties, in our opinion–we focused on time series models due to space limitations, and with the goal of providing controlled comparisons to existing models on a standardized benchmark. We believe exploration of the model on novel tasks is important future work. \n\nRegarding Figure 2, we will resize and make better use of whitespace to improve clarity of the visualization.  \n\nRegarding the straight-through gradient estimator, unlike prior works which make use of the estimator when gradients are intractable, in this case we are using it as a means of removing a Jacobian from the gradient computation. To compute the natural gradient (which is a preconditioned gradient and provably always a descent direction) we need to either use the straight-through gradient estimator or manually solve a system of linear equations induced by the Fisher Information matrix, which would prove much more expensive. Our resulting estimator of the natural gradient is unbiased.\n\nTable 3 does not include pre-training, but this initialization scheme is not expensive: because the standard VAE does not take into account time dependence (and therefore does not need to perform Kalman smoothing inference or RNN iteration) it is faster than any model discussed in the paper. Further, because models are trained for 200/100 epochs on Mocap/WSJ0 and we only pre-train for 10 epochs, the entire pre-training took <20 minutes on each data set. "}, {"review_id": "KCpIuHYIew", "reviewer": "Reviewer_FFkU", "summary": "This work presents a novel implicit optimisation approach for Structured Variational Autoencoders (SVAEs) allowing for the computation of natural gradients in an automatic way. This allows for both an efficient learning scheme that is robust and avoids the bias found in other optimisation approaches of models that comprise Discrete variables, e.g., continuous relaxations. Experimental evaluations of these new optimisation schemes lead to SVAEs with competitive performance when compared to SOTA time series models while resulting in interpretable learned representations.", "questions": "Please see the \"Weaknesses\" section.", "limitations": "Please see the \"Weaknesses\" section.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "SVAEs constitute a powerful generative approach with potentially rich latent representations. The authors propose a new optimisation scheme that could bypass their training complexity and improve research explorations.", "weaknesses": "This work misses comparison against other related methods and some of the claims should be further clarified.\n\n1) The continuous relaxation of the discrete variables do indeed produced a biased approximation, but have been shown to produce good \n     results in practice, even without annealing schedules (line 117). For example see [3,4,5]. These claims should be validated with \n     appropriate experimental evaluations in the considered setting.\n\n2) Since the SVAE framework resorts to maximisation of a surrogate ELBO, doesn't it introduce a different bias? the same applies to when using the Richardson iteration; this approximation does introduce additional bias in the estimation, especially considering the capped implicit gradient estimator that the authors propose. \n\n3) The related work is a bit lacking. The authors should discuss [2,6] and potentially compare against [6]. At the same time, there exists a newly published work by the same authors [7] that was accepted at ICML 2023, published before this submission. Additional discussions and comparisons should be included. \n\n4) In this context, in my view there should be a comparison with other optimisation schemes for discrete latent variables such as the Gumbel-Softmax [1] or the NES strategy [2]. Even though the framework improves over the baseline and performs many times better than the competition, it is important to quantify the impact of the optimisation compared to more recent approaches than the biased estimation of the original SVAE publication. \n\n5) What is the computational complexity of the framework in the various settings that the authors discuss?\n\n[1] Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables, ICLR 2017\n\n[2] Alon Berliner, Guy Rotman,Yossi Adi et al. Learning Discrete Structured Variational Autoencoder Using Natural Evolution Strategies, \nICLR 2022.\n\n[3] KP Panousis, S Chatzis, S Theodoridis, Nonparametric Bayesian Deep Networks with Local Competition, ICML 2019 \n\n[4] K Kalais, S Chatzis, Stochastic deep networks with linear competing units for model-agnostic meta-learning, ICML 2022\n\n[5] Sebastian Jaszczur, Aakanksha Chowdhery, Afroz Mohiuddin, et al., Sparse is Enough in Scaling Transformers, NIPS 2021\n\n[6] Yixiu Zhao, Scott W. Linderman, Revisiting the Structured Variational Autoencoder, NIPS BDL Workshop 2021\n\n[7] Yixiu Zhao, Scott W. Linderman, Revisiting Structured Variational Autoencoders, ICML 2023\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the added citations and comments. We will respond to each comment separately.\n\n1: Thank you for the citations; added discussion of this model family is appropriate. For a discussion of these methods as an experimental baseline, please see the global response above.\n\n2: The surrogate ELBO is used only to produce a variational posterior $q(z|x)$. Because we use a surrogate loss (and not the true loss), the posterior will be suboptimal. However, in amortized variational inference the posterior is always suboptimal! Thus the penalty for using a surrogate loss is simply an amortization gap, which is unavoidable in VAE learning. Unlike continuous relaxation-based methods, which relax the overall loss function and therefore perform descent on the wrong objective, all model parameters in the SVAE are optimized with respect to the true loss function.\n\nThe Richardson iteration, if run to convergence, can exactly compute the implicit gradient. By capping the number of iterations we do introduce the possibility of biased gradients. This yields a tradeoff between compute and accuracy, unlike relaxation-based approaches which are biased even in the large-compute limit. \n\n3: Thank you for the citations. We will add discussion of these papers to our related work. We would like to highlight that [7] was not published before our submission: it appeared on ArXiv on May 25th, a week after the NeurIPS submission deadline, and was published in the proceedings of ICML after that. We will add discussion of this work to section 7, and note that the paper differs from ours in 3 key respects. First, they show no experimental results on real data, only exploring the capacity of the SVAE to model synthetic and toy data. Second, because their experiments are restricted to the LDS graphical model (which requires no mean field factorization nor block updating) they do not need implicit differentiation, and do not explore the capacity of the SVAE to include discrete latent variables. Third, because they optimize point estimates of parameters instead of variational factors, they do not make use of natural gradients. In their point-estimate formulation, they can apply the parallel Kalman smoother of [1] off-the-shelf, whereas we derive a novel extension for our variational setting. Early experiments of ours revealed that using point estimates (and no natural gradients) barely hurts optimization of the LDS, but causes gradual pruning of discrete states in the SLDS until it has collapsed to a single mode. As this is the only difference between our implementations not already explored in our paper’s experiments, we would be happy to add results demonstrating this to our appendix in future revisions.\n\nRegarding [2], we would like to highlight that although it refers to “discrete structured VAEs”, these are unrelated to the SVAEs discussed in our paper and in [6-7]. [2] uses the concrete/gumbell softmax trick to model discrete latent variables.\n\n4: Please see our global response for discussion of Gumbel-Softmax relaxations. We further highlight that no related work (to our knowledge) has provided a formulation for applying either scheme to dynamical VAEs for time series models. \n\n5: Let D be the dimension of the continuous latent variable, and K be the dimension of the discrete latent variable. With a sequence of length T and L block updates, we have $O(LT(K^2+D^2 + KD))$ time complexity. Note that the $K^2 + KD$ factors can be parallelized to $\\log K + \\log D$ (they are matrix multiplies/element wise products followed by a sum along one axis), and our Kalman smoother can reduce the dependence on T to logT. Thus with optimal parallelization we have complexity $O(L \\cdot \\log T \\cdot (D^2 + \\log K))$. Because we do not store intermediate states our memory complexity is independent of L, capping out at $O(T(D^2+K^2))$.\n\n[1] Särkkä, S. and García-Fernández, Á. F. Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control, 2020\n\n[2, 6, 7] (we re-used your numbering to improved clarity of our response)\n"}, {"review_id": "79321mz9BX", "reviewer": "Reviewer_gyLw", "summary": "The authors revisit learning algorithms for structured variational autoencoders. \n\nThis paper proposes using the implicit function theorem to compute the implicit gradients, improving stability and memory efficiency. The authors further revisit the natural gradients studied in the original SVAE paper and propose an unbiased version. \n\nThe experimentation is carried out on human pose capture and audio spectrogram datasets. The evaluation is based on modelling performances (likelihood), visual quality (FID score) and missing data imputation.", "questions": "1. In the original SVAE paper, the authors explore the case of a simple latent GMM model, why not explore this simple case in this paper?\n2. Would your method scale to larger datasets as explored (for instance) in [\"Clockwork Variational Autoencoders\", Saxena et al., 2021]?\n3. Could you please detail how the FID scores are computed?", "limitations": "Technical limitations are not clearly discussed.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper is impressive by the depth and breadth of its theory (implicit gradients, unbiased natural gradients)\n2. A complete and well-detailed appendix featuring all derivations, although I was not in a position to proof-read all of it\n3. The experiment section shows a clear improvement of `SVAE-SLDS` (the method introduced in this paper) over the baseline `SVAE-SLDS-Bias` based on likelihood performances and FID", "weaknesses": "1. This paper lacks structure and clarity. More work is needed to synthesize and popularize the main contributions.\n2. This paper attempts to make too many contributions and therefore fails to present and support each of them sufficiently. \n3. The experiment section is not sufficient (only two rather small datasets).\n4. Section 4 claims proposing a more memory efficient algorithm: an experiment demonstrating the applicability to larger-scale datasets is missing.\n5. Section 3.2 seems to indicate that discrete latent variable models can only be optimized by reparameterizing discrete distributions. This is wrong, score function estimators (e.g. Reinforce) are viable alternatives.  [\"Monte Carlo Gradient Estimation in Machine Learning\", Mohamed et al., 2020](https://jmlr.org/papers/volume21/19-346/19-346.pdf) gives a great overview of gradient estimation methods and should, in my opinion, be discussed in the related work section. Reparameterization-free methods exist and work decently. For instance:\n\t- [\"Variational Inference for Monte Carlo Objectives\", Mnih et al., 2016](http://proceedings.mlr.press/v48/mnihb16.pdf)\n\t- [\"Revisiting Reweighted Wake-Sleep for Models with Stochastic Control Flow\", Anh et al, 2018](https://arxiv.org/abs/1805.10469) \n\t- [\"The Thermodynamic Variational Objective\", Masrani et al., 2019](https://arxiv.org/abs/1907.00031)\n\t- [\"Optimal Variance Control of the Score Function Gradient Estimator for Importance Weighted Bounds\", Lievin et al., 2020](https://arxiv.org/abs/2008.01998)\n\t- [\"Gradient Estimation with Discrete Stein Operators\", Shi et al., 2022](https://arxiv.org/abs/2202.09497)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough feedback. We will respond to your comments and questions individually.\n\n1 & 2: Could you provide any specifics for which contributions are not sufficiently supported or not sufficiently connected to the main structure of the paper?\n\n3: We chose our two datasets to replicate the comparisons made in [1], a recent survey paper of dynamical VAEs. The two datasets consist of 53,443 and 93,215 sequences respectively. This contrasts to related work, which has been restricted to toy and synthetic data: the original SVAE paper [2], the structured inference network paper [3], and recent concurrent work on extending the SVAE (see response to reviewer FFkU) [4] are all only evaluated on synthetic data (except for one experiment in [2] on not-publicly-available data). \n\n4: Greater memory efficiency does not impact the ability of the model to handle larger datasets, so long as mini-batch learning is used (the relationship between memory efficiency and mini-batch size is demonstrated in Table 1). This improvement does, however, enable easier application of the SVAE to larger models and/or longer sequences which have high memory overhead. Specifically our approach avoids storing *all* of the intermediate values computed within the inner optimization, which would otherwise be need for backpropagation. Even for moderately sized datasets, explicit backpropagation can induce prohibitively high memory costs as shown in table 1.\n\n5: Thank you for the citations; we will add more discussion of these broader methods to section 3.2. We would like to highlight that, to our knowledge, no one has successfully integrated these methods with the model family we are considering in this paper (dynamical VAEs). As a result, we believe these methods are more suitable for the related work section (as you recommend) than for experimental baselines.\n\nWith regard to score function estimators specifically, in preliminary experiments on related models we found these techniques had poor properties. For some published evidence on this front, see [5] which showed that score function estimators have orders of magnitude higher variance than reparameterization-based estimators. [6] compared natural gradient descent to score function estimators and drew the same conclusion.\n\nQ1: The original SVAE paper only explores this simple mixture model on synthetic data. We experimented with this model (and it is implemented in our codebase to be shared), and found that while our methods improve computation and memory efficiency, there was no substantial difference between the models learned by our method and the original SVAE paper. Given this, we felt comparisons on this model were less interesting than the ones included in our submission, but we’d be happy to add an experiment on this model family to our supplement.\n\nQ2: Our paper does not focus on long-range-dependencies (which is its own robust line of work, usually evaluated by temporal prediction tasks rather than the generative model quality/metrics), however our SVAE would be quicker to train on such data than RNN baselines: our derived parallel inference routine has time complexity logarithmic in the sequence length instead of linear.\n\nQ3: Thank you for the question. Details of our FID protocol (and architectures) were missing from our appendix and will be added. The architectures are described in a table (see attached pdf) and are detailed below. \n\nMocap: The InceptionV3 [7] network typically used to compute Frechet Inception Distance scores was trained on natural images and is therefore inappropriate for evaluating the generative performance of the H3.6M motion capture data. To resolve this, we compute the same score using a different underlying network. We trained a convolutional network to predict which of the 15 different motions was being performed in each sequence (note that in H3.6M, motions are distinct from actions/scenarios, which may include more subtle variations). We used a residual network architecture, where each residual block consists of 2 1-D convolutional layers with batch normalization and a skip connection. We chose 1-D convolutions (across time) as they are more appropriate for our temporal motion capture data. We used the Optuna library to search across model depths, activations and layer sizes to find a network within our design space that performed optimally on a validation dataset (25% of the original training data). As with the original FID score, we use the output of the final average pooling layer to compute the score.\n\nWSJ0: As with the H3.6M dataset, it is necessary to define a different base network to evaluate generative FID scores for the WSJ0 dataset. As the WSJ0 dataset does not have an appropriate set of labels to use to train a classifier, we chose to instead train a classifier on the CREMA-D dataset [9] which consists of similar speech clips. We preprocessed each audio sequence in CREMA-D identically to the WSJ0 dataset, and trained a classifier to predict one of 6 different emotions that was expressed. We used the same setup as H3.6M to choose a network architecture, train the model, and compute FID scores. \n\n[1] See [14] in our paper.\n\n[2] See [19] in our paper. \n\n[3] See [30] in our paper.\n\n[4] Yixiu Zhao, Scott W. Linderman, Revisiting Structured Variational Autoencoders. ICML 2023\n\n[5] Kucukelbir, A., Tran, D., Ranganath, R., Gelman, A., Blei, D., Automatic Differentiation Variational Inference. JMLR 2017\n\n[6] Ji, G., Sujono, D., and Sudderth, E. Marginalized Stochastic Natural Gradients for Black-Box Variational Inference. ICML 2021\n\n[7] Szegedy, C., Vanhoucke, V., Ioffe, S., Shlens, J., and Wojna, Z. Rethinking the inception architecture for computer vision. CVPR 2015\n\n[8] Cao, H., Cooper, D. G., Keutmann, M. K., Gur, R. C., Nenkova, A., and Verma, R. CREMA-D: Crowd-sourced emotional multimodal actors dataset. IEEE transactions on affective computing, 2014."}]

Author rebuttal:
We appreciate the feedback provided by all reviewers. In addition to the individual responses, we would like to highlight a point brought up by multiple reviewers: that we emphasize our model’s unbiased modeling of discrete latent variables but do not compare to the biased approach of the concrete (Gumbel-Softmax) VAE [1, 2].

To our knowledge, no one has successfully integrated concrete distributions into the model family discussed in this paper, dynamical VAEs. Thus, to make a comparison possible we have devised a new model which incorporates concrete distributions into a dynamical architecture, which we will call the “deep SLDS” (DSLDS) model. Results of this new comparison are provided below/attached.

The DSLDS has an identical generative model with the SLDS SVAE: switching linear dynamics in the latent space followed by a neural network decoder. However, instead of the SVAE surrogate loss trick, its inference is fully amortized and closely resembles inference in the DKS: a network independently encodes the observation at every time step, which are then used to iteratively sample $z_t$, $k_t$ in a Markovian fashion (i.e. k_t is sampled  given $k_{t-1}, z_{t-1}$ as well as $x_t$, and $z_t$ is sampled given $k_t, z_{t-1}$, and $x_t$). An intuitive visualization of the inference model is included in the attached figures. By combining the inference parameterization of the DKS and the generative parameterization of the SVAE SLDS, we can make a fair comparison between the two approaches. 

Our experiments produce a model with reasonable reconstructions but very poor generations (see attached pdf). The deep SLDS learns latent dynamics that quickly diverge, leading to unrecognizable generations. For comparison with the results in table 2: the estimated likelihood ($> \log p(x)$) for the DSLDS is **1.953** for Mocap and **1.278** for WSJ0. The resulting generations produced “inception” activations far from those of real data; attempting to estimate FID scores with these generations resulted in numerical issues (the numbers would be far larger than those for competing methods).

As VAEs do not generate at training time, a large amortization gap can lead to this kind of gap in generative and reconstruction performance. The SVAE resolves this issue by using generative parameters at inference time; this technique cannot be easily applied to models whose latent variables follow a concrete distribution because it is not an exponential family. Approximation errors may be magnified in dynamical VAEs, because small differences in parameters can compound over time to produce large changes in generations at later time steps.

There are other reasons to prefer the SVAE setup to concrete relaxations: [3] demonstrated that natural gradient descent outperforms concrete approximations for learning parameters of discrete distributions. In addition, the original concrete paper [1] focused on binary variables and noted difficulty scaling to discrete variables with a higher number of states (they demonstrate this by changing the variable arity from 4 to 8, whereas we use 50). They note that for K discrete states, the temperature must be less than 1/(1-K) to sufficiently approximate the discrete distribution, but also note that as K increases, higher temperatures are necessary to fight the increased peakiness of the concrete distribution. In our DSLDS experiments, we were only able to set the temperature as low as 0.2-0.3 before encountering numerical issues. We believe our results reflect the difficulty of the concrete distributions to handle 50-ary discrete random variables.

[1] Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete distribution: A continuous relaxation of discrete random variables. ICML 2017.

[2] Eric Jang, Shixiang Gu, and Ben Poole. Categorical Reparametrization with Gumbel-Softmax. In International Conference on Learning Representations (ICLR 2017), 2017.

[3] Ji, G., Sujono, D., and Sudderth, E. Marginalized Stochastic Natural Gradients for Black-Box Variational Inference. ICML 2021


Meta-review: The work adapts statistical properties of the exponential family and their natural representation to structural variational auto-encoders. The work ties together many components and with this, its readability can definitely be improved to have an impact to the community. Despite my excitement about this submission I would like to echo the discussion: "The primary purpose of a scientific paper is to communicate a novel idea and share advances made in the field. Simplifying an idea, method or experimental protocol is a difficult but necessary step to achieve this, and failure to do so would only hurt the long-term impact of the paper."
Decision: Accept (poster)