SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Harry Bendekgey, Gabriel Hope, Erik B. Sudderth
{hbendekg, hopej, sudderth}@uci.edu
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.
Advances in deep learning have dramatically increased the expressivity of machine learning models at great cost to their interpretability. This trade-off can be seen 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 variables associated with every observation. Johnson et 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 a 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 can 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 easily accommodates missing data, leading to accurate and multimodal imputations. We further improve training speed by generalizing prior work on parallel Kalman smoothers [56].
Figure 1: The SVAE-SLDS segments each sequence of human motion, which we display as a sequence of discrete colors. Discrete variables are interpretable: Below each segmentation, we show five segmentations of 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.
We begin in Sec. 2 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 be efficiently trained to their full potential. Although SVAEs may incorporate any latent graphical structure, we focus on temporal data. In Sec. 8, we are the first to compare SVAE performance to state-of-the-art recurrent neural network-and transformer-based architectures on time series benchmarks [21], and the first to demonstrate that SVAEs provide a principled method for multimodal interpolation of missing data.
2 Background: Graphical Models and Variational Inference
We learn generative models that produce complex data  $x$  via lower-dimensional latent variables  $z$ . The distribution  $p(z|\theta)$  is defined by a graphical model (as in Fig. 2) with parameters  $\theta$ , and  $z$  is processed by a (deep) neural network with weights  $\gamma$  to compute the data likelihood  $p_{\gamma}(x|z)$ .
Exact evaluation or simulation of the posterior  $p_{\gamma}(z,\theta |x)$  is intractable due to the neural network likelihood. Variational inference (VI [62]) defines a family of approximate posteriors, and finds the distribution that best matches the true posterior by optimizing the evidence lower bound (ELBO):
$$
\mathcal {L} [ q (\theta ; \eta) q (z; \omega), \gamma ] = \mathbb {E} _ {q (\theta ; \eta) q (z; \omega)} \left[ \log \frac {p (\theta) p (z | \theta) p _ {\gamma} (x | z)}{q (\theta ; \eta) q (z ; \omega)} \right] \leq \log p _ {\gamma} (x). \tag {1}
$$
Here,  $q(\theta ;\eta)q(z;\omega)\approx p_{\gamma}(z,\theta |x)$  are known as variational factors. We parameterize these distributions via arbitrary exponential families with natural parameters  $\eta ,\omega$  . This implies that
$$
q (z; \omega) = \exp \{\langle \omega , t (z) \rangle - \log Z (\omega) \}, \quad Z (\omega) = \int_ {z} \exp \{\langle \omega , t (z) \rangle \} d z. \tag {2}
$$
An exponential family is log-linear in its sufficient statistics  $t(z)$ , where the normalizing constant  $Z(\omega)$  ensures it is a proper distribution (see App. C.1 for properties of exponential families). For models where  $p_{\gamma}(x|z)$  has a restricted conjugate form (rather than a deep neural network), we can maximize Eq. (1) by alternating optimization of  $\eta, \omega$ ; these coordinate ascent updates have a closed form [62]. Stochastic VI [23] improves scalability (for models with exponential-family likelihoods) by sampling batches of data  $x$ , fitting a locally-optimal  $q(z; \omega)$  to the latent variables in that batch, and updating  $q(\theta; \eta)$  by the resulting (natural) gradient.
Amortized VI. Because it is costly to optimize Eq. (1) with respect to  $\omega$  for each batch of data, VAEs employ amortized VI [32, 47, 53] to approximate the parameters of the optimal  $q(z;\omega)$  via a neural network encoding of  $x$ . The inference network weights  $\phi$  for this approximate posterior  $q_{\phi}(z|x)$  are jointly trained with the generative model. A potentially substantial amortization gap exists [14, 36]: the inference network does not globally optimize the ELBO of Eq. (1) for all  $x$ .
Structured VAEs. For a fixed  $q(\theta)$ , the true optimizer (across all probability distributions) of Eq. (1) is given by  $q(z)\propto p_0(z;\eta)p_{\gamma}(x|z)$ , where  $p_0(z;\eta)\propto \exp \{\mathbb{E}_{q(\theta ;\eta)}[\log p(z|\theta)]\}$  is an expected prior on  $z$  (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_{\gamma}(x|z)$ . Rather than approximating
Figure 2: Left: Generative (above) and inference (below) graphical models for SVAE-LDS and SVAE-SLDS. For the SLDS, we show the prior and posterior as factor graphs [38].  $q_{\phi}(z|x)$  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 are propagated across residual edges between disentangled factors. Right: General form of iterative mean-field (MF) and belief-propagation (BP) updates for SVAE inference in a model where  $q(z)$  factorizes into  $M$  groups of latent variables. For the SVAE-SLDS,  $M = 2$ .
this whole posterior as 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  $\ell(z) = p_{\gamma}(x|z)$  with a function parameterized by a neural network  $\hat{\ell}_{\phi}(z|x)$ . The optimal posterior given this approximation is then equal to
$$
q _ {\phi} (z | x; \eta) = \underset {q (z)} {\arg \max } \hat {\mathcal {L}} [ q (\theta ; \eta) q (z), \phi ] = \underset {q (z)} {\arg \max } \mathbb {E} _ {q (\theta ; \eta) q (z)} \left[ \log \frac {p (\theta) p (z | \theta) \hat {\ell} _ {\phi} (z | x)}{q (\theta ; \eta) q (z)} \right]. \tag {3}
$$
The posterior that optimizes the surrogate loss  $\hat{\mathcal{L}}[\cdot, \phi]$  of Eq. (3) is  $q_{\phi}(z|x;\eta) \propto p_0(z;\eta)\hat{\ell}_{\phi}(z|x)$ . If  $\hat{\ell}_{\phi}(z|x)$  is chosen to be conjugate to  $p(z|\theta)$ , this multiplication and normalization is easy for many exponential families. Note that  $\hat{\ell}_{\phi}(z|x)$  does not need to be normalized, as any multiplicative factors would disappear when the posterior is normalized. The overall ELBO is then  $\mathcal{L}[q(\theta; \eta)q_{\phi}(z|x;\eta), \gamma]$ .
If the encoder's approximate likelihood  $\hat{\ell}_{\phi}$  is (up to normalization) close to the true likelihood  $\ell$ , the surrogate loss  $\hat{\mathcal{L}}$  will closely approximate the true loss  $\mathcal{L}$ , and there will be little amortization gap. SVAE inference has the advantage that  $q_{\phi}(z|x;\eta)$  depends on the learned posterior  $q(\theta ;\eta)$  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 VAEs [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 1: Standard Normal. For a basic VAE,  $\theta$  is fixed and  $z\sim \mathcal{N}(0,I)$ . The SVAE inference network outputs a Gaussian  $\hat{\ell}_{\phi}(z|\mu = \mu (x;\phi),\tau = \tau (x;\phi))\propto \mathcal{N}(z;\mu ,\mathrm{diag}(\tau^{-1}))$  with mean  $\mu$  and inverse-variance (precision)  $\tau$ . The product of this Gaussian with the standard normal prior has a simple form:  $q_{\phi}(z|x) = \mathcal{N}(z;\frac{\tau}{\tau + 1}\mu ,\mathrm{diag}((\tau +1)^{-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:  $z_{t} \sim \mathcal{N}(A_{t}z_{t - 1} + b_{t}, Q_{t})$ ,  $z_{1} \sim \mathcal{N}(\mu_{1}, \Sigma_{1})$ . In this case, we expand the exponential family distribution  $q(z; \omega)$  as a sum across time steps:
$$
q (z; \omega) = q \left(z _ {1}\right) \prod_ {t = 2} ^ {T} q \left(z _ {t} \mid z _ {t - 1}\right) = \exp \left\{\left\langle \omega_ {1}, t \left(z _ {1}\right) \right\rangle + \sum_ {t = 2} ^ {T} \left\langle \omega_ {t}, t \left(z _ {t - 1}, z _ {t}\right) \right\rangle - \log Z (\omega) \right\}, \tag {4}
$$
where  $\omega = \mathrm{concat}(\omega_{t})$  and the prior  $p(z|\theta)$  belongs to this family. The prior induces temporal dependence between the  $z_{t}$ , but we assume the likelihood factorizes as  $p_{\gamma}(x|z) = \prod_t p_{\gamma}(x_t|z_t)$ .
For models like the LDS, approximating the likelihood function with conjugate, time-independent Gaussian distributions is a much simpler task than approximating the temporally-coupled posterior. In addition, as the generative parameters  $A_{t}, b_{t}, Q_{t}$  of  $p(z|\theta)$  are learned through optimization of  $q(\theta)$ , the inference routine shares those parameters, improving accuracy. These advantages were not present in the standard normal VAE, where the prior on  $z$  lacks structure and is not learned. Inference, normalization, and sampling of  $q_{\phi}(z|x;\eta)$  in the LDS model is feasible via a Kalman smoothing algorithm [4] that efficiently (with cost linear in  $T$ ) aggregates information across time steps.
3 Structured Variational Inference
For complex graphical models, the distributions  $p(z|\theta)$  and  $q(z)$  typically factorize across subsets of the latent variables  $z$ , as illustrated in Fig. 2. We thus generalize Eq. (4) by partitioning  $z$  into local variable groups, and representing the dependencies between them via a set of factors  $\mathcal{F}$ :
$$
q (z; \omega) = \exp \left\{\sum_ {f \in \mathcal {F}} \langle \omega_ {f}, t \left(z _ {f}\right) \rangle - \log Z (\omega) \right\}. \tag {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 to compute marginals or normalize the  $q_{\phi}(z|x;\eta)$  defined in Sec. 2. 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  $\{z_m\}_{m=1}^M$  be a partition of the variables in the graphical model, chosen so that inference within each  $z_m$  is tractable. We infer factorized (approximate) marginals  $q_{\phi}(z_m|x;\eta)$  for each mean field cluster by maximizing  $\hat{\mathcal{L}}[q(\theta; \eta) \prod_m q(z_m), \phi]$ . The optimal  $q_{\phi}(z_m|x;\eta)$  inherit the structure of the joint optimizer  $q_{\phi}(z|x;\eta)$ , 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 to  $m$  via residual edges. These expectations in turn depend on their clusters' parameters, defining a stationary condition for the optimal  $\omega$ :
$$
\omega_ {m} = \operatorname {M F} (\mu_ {- m}; \eta), \quad \mu_ {m} = \operatorname {B P} (\omega_ {m}; \eta , \phi , x). \tag {6}
$$
Here, BP is a belief propagation algorithm which computes expected statistics  $\mu_{m}$  for cluster  $m$ , and the linear mean field function MF updates parameters of cluster  $m$  given the expectations of other clusters  $\mu_{-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  $\hat{\mathcal{L}}[q(\theta)\prod_{m}q(z_m)]$ .
3.2 Reparameterization and Discrete Latent Variables
While optimizing  $q_{\phi}(z_m|x;\eta)$  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  $\phi$  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 high-variance gradient estimates [50, 28]. Continuous relaxations of discrete variables [44, 27, 7] produce biased approximations of the true discrete ELBO, and are sensitive to annealing schedules for temperature hyperparameters.
SVAE training only requires reparameterized samples of those latent variables which are direct inputs to the generative network  $p_{\gamma}(x|z)$ . By restricting these inputs to 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  $z_{m}$ , and discrete variables by  $k_{m}$ .
Example 3: Gaussian Mixture. Consider a generalized VAE where the state is sampled from a mixture model:  $k \sim \mathrm{Cat}(\pi)$ ,  $z \sim \mathcal{N}(\mu_k, \Sigma_k)$ . The likelihood  $p_{\gamma}(x|z)$  directly conditions on only the continuous latent variable  $z$ . Variational inference produces disentangled factors  $q_{\phi}(z|x;\eta)q_{\phi}(k|x;\eta)$ , and we evaluate likelihoods by decoding samples from  $q_{\phi}(z|x;\eta)$ , without sampling  $q_{\phi}(k|x)$ .
Example 4: Switching Linear Dynamical System (SLDS). Consider a set of discrete states which evolve according to a Markov chain  $k_{1} \sim \mathrm{Cat}(\pi_{0}), k_{t} \sim \mathrm{Cat}(\pi_{k_{t-1}})$ , and a continuous state evolving according to switching linear dynamics:  $z_{0} \sim \mathcal{N}(\mu_{0}, \Sigma_{0}), z_{t} \sim \mathcal{N}(A_{k_{t}}z_{t-1} + b_{k_{t}}, Q_{k_{t}})$ . The transition matrix, offset, and noise at step  $t$  depends on  $k_{t}$ . Exact inference in SLDS is intractable [39], but structured VI [20] learns a partially factorized posterior  $q_{\phi}(z|x;\eta)q_{\phi}(k|x;\eta)$  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-backward message-passing to compute marginals of discrete states (see App. D). Let  $k_{tj} = 1$  if the SLDS is in discrete state  $j$  at time  $t$ ,  $k_{tj} = 0$  otherwise,
Method	Time of gradient step (ms)
B = 1	B = 32	B = 64	B = 128
Implicit + Parallel	603	922	1290	2060
Unrolled + Parallel	659	1080	n/a	n/a
Implicit + Sequential	2560	3160	3290	3530
Unrolled + Sequential	2660	3290	3980	n/a
Table 1: Time of ELBO backpropagation in an SVAE-SLDS 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 attempted to allocate more than 48GB of GPU memory.
and  $\bar{\theta}_j = \mathbb{E}_{q(\theta ;\eta)}[\theta_j]$  be the expected (natural) parameters of the LDS for discrete state  $j$ . Structured VI updates the natural parameters of discrete states  $\omega_{k_{tj}}$ , and continuous states  $\omega_{z_t,z_{t + 1}}$ , as follows:
$$
\omega_ {z _ {t}, z _ {t + 1}} = \sum_ {j} \mathbb {E} _ {q} \left[ k _ {t j} \right] \bar {\theta} _ {j}, \quad \omega_ {k _ {t j}} = \langle \bar {\theta} _ {j}, \mathbb {E} _ {q} [ t (z _ {t - 1}, z _ {t}) ] \rangle . \tag {7}
$$
4 Stable and Memory-Efficient Learning via Implicit Gradients
When  $q_{\phi}(z|x)$  is 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 \sim q_{\phi}(z|x;\eta)$  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. 1. 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  $\frac{\partial\omega}{\partial\eta}$ ,  $\frac{\partial\omega}{\partial\phi}$  without storing intermediate states. We focus on gradients with respect to  $\eta$  for compactness, but gradients with respect to  $\phi$  are computed similarly. Let  $\omega^{(1)},\ldots ,\omega^{(L)}$  be the sequence of  $\omega$  values produced during the "forward" pass of block coordinate ascent, where  $\omega^{(L)}$  are the optimized structured VI parameters. The IFT expresses gradients via the solution of a set of linear equations:
$$
\frac {\partial \omega^ {(L)}}{\partial \eta} = \left(\frac {\partial g (\omega ; \eta , \phi , x)}{\partial \omega}\right) ^ {- 1} \frac {\partial g (\omega ; \eta , \phi , x)}{\partial \eta}, \quad g (\omega) = \omega - \mathrm {M F} (\mathrm {B P} (\omega ; \eta , \phi , x); \eta). \tag {8}
$$
Here we apply the BP and MF updates in parallel for all variable blocks  $m$ , rather than sequentially as in Eq. (6). At a VI fixed point, these parallel updates leave parameters unchanged and  $g(\omega) = 0$ .
For an SLDS with latent dimension  $D$  and  $K$  discrete states,  $\omega$  has  $O(K + D^2)$  parameters at each time step. Over  $T$  time steps,  $\frac{\partial g}{\partial\omega}$  is thus a matrix with  $O(T(D^{2} + K))$  rows/columns and  $O(T^{2}D^{2}K)$  non-zero elements. For even moderate-sized models, this is infeasible to explicitly construct or solve.
We numerically solve Eq. (8) via a Richardson iteration [54, 64] that repeatedly evaluates matrix-vector products  $(I - A)v^{\prime}$  to solve  $A^{-1}v$ . 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.
$$
R i c h a r d s o n: \quad \frac {\partial \omega^ {(L)}}{\partial \eta} \approx - \sum_ {j = 0} ^ {J} \left(I - \frac {\partial g (\omega^ {(L)} ; \eta , \phi , x)}{\partial \omega}\right) ^ {j} \frac {\partial g (\omega^ {(L)} ; \eta , \phi , x)}{\partial \eta}. \tag {9}
$$
$$
\text {U n r o l l e d :} \quad \frac {\partial \omega^ {(L)}}{\partial \eta} \approx - \sum_ {\ell = 0} ^ {L} \left[ \prod_ {i = \ell} ^ {L} \left(I - \frac {\partial g \left(\omega^ {(i)} ; \eta , \phi , x\right)}{\partial \omega}\right) \right] \frac {\partial g \left(\omega^ {(\ell)} ; \eta , \phi , x\right)}{\partial \eta}. \tag {10}
$$
Lorraine et al. [43] tune the number of Richardson steps  $J$  to balance speed and accuracy. However, there is another reason to limit the number of iterations: when the forward pass is not iterated until convergence,  $\omega^{(L)}$  is not a stationary point of  $g(\omega)$  and therefore Eq. (9) is not guaranteed to converge as  $J \to \infty$ . For batch learning, waiting for all VI routines to converge to a (local) optimum might be prohibitively slow, so we might halt VI before  $\omega^{(L)}$  converges to 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, so that
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 a batch of 128 data points, and dashed lines show  $90^{\text{th}}$  percentiles. Conditioning: Convergence of SVAE-LDS negative-ELBO loss versus number of optimization steps (log-scale) for conventional (non-natural) gradients, biased natural gradients [29], and unbiased natural gradients computed via automatic differentiation (Sec. 5).
$J = L$ . In this regime, implicit gradient computation has a one-to-one correspondence to unrolled gradient computation, while requiring a small fraction of the memory. This can be thought of as a 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).