SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 implicit proposal to other gradient estimators. Ground truth gradients are computed as the implicit gradient at the optimum, and we compare the root-mean-squared-error (rMSE) of various gradient estimators to that of the naive No-Solve solution, which replaces the inverted matrix in Eq. (8) with Identity.
We consider two models: one 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 to converge. We compare the memory-intensive unrolled estimator (Unrolled) to 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 a 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+Thresh): 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 is 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  $\phi$ , suffering high memory overhead. They compute neither unrolled nor implicit gradients with respect to generative model parameters  $\eta$ ; in practice they set the gradient of the inner optimization to 0, yielding a biased training signal. Our innovations instead enable memory-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_{\eta}$  of  $q(\theta; \eta)$ , given by:
$$
F _ {\eta} = \mathbb {E} _ {q (\theta ; \eta)} \left[ \left(\nabla_ {\eta} q (\theta ; \eta)\right) \cdot \left(\nabla_ {\eta} q (\theta ; \eta)\right) ^ {T} \right]. \tag {11}
$$
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(z)$  in an inner optimization loop and learns  $q(\theta; \eta)$  in an outer loop by natural gradient descent. The key difference between SVI and the SVAE is that SVI's inner optimization is done with respect to the true loss function  $\mathcal{L}$ , whereas the SVAE uses a surrogate  $\hat{\mathcal{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  $\mu_{\eta}$  be the expected sufficient statistics of  $q(\theta ;\eta)$ . Exponential family theory tells us that  $\frac{\partial\mu}{\partial\eta} = F_{\eta}$  [30, 45], allowing Johnson et al. [29] to derive the natural gradients of the SVAE loss:
$$
\frac {\partial \mathcal {L}}{\partial \eta} F _ {\eta} ^ {- 1} = \frac {\partial \mathcal {L}}{\partial \mu} \frac {\partial \mu}{\partial \eta} F _ {\eta} ^ {- 1} = \frac {\partial \mathcal {L}}{\partial \mu}, \quad \frac {\partial \mathcal {L}}{\partial \mu} = \overbrace {\eta_ {0} + \mathbb {E} _ {q _ {\phi} (z | x ; \eta)} [ t (z) ] - \eta} ^ {\text {S V I u p d a t e}} + \overbrace {\frac {\partial \mathcal {L}}{\partial \omega} \cdot \frac {\partial \omega}{\partial \eta}} ^ {\text {c o r r e c t i o n t e r m}}. \tag {12}
$$
This gradient differs from the SVI gradient by the final term: because SVI's inner loop optimizes  $\omega$  with respect to the true loss  $\mathcal{L}$ ,  $\frac{\partial\mathcal{L}}{\partial\omega} = 0$  for conjugate models. Johnson et 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  $\frac{\partial\omega}{\partial\eta}$  involves computing an implicit or unrolled gradient, as addressed by our numerical 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 natural parameters are guaranteed to be preserved, such as positivity or positive-definiteness.
We resolve this issue by reparameterizing  $\eta$  into an unconstrained space, and computing natural gradients with respect to those new parameters. Letting  $\tilde{\eta}$  be an unconstrained reparameterization of  $\eta$ , such as  $\eta = \text{Softplus}\{\tilde{\eta}\} = \log(1 + e^{\tilde{\eta}})$  for a positive precision parameter, we have:
$$
\frac {\partial \mathcal {L}}{\partial \tilde {\eta}} F _ {\tilde {\eta}} ^ {- 1} = \frac {\partial \mathcal {L}}{\partial \mu} \cdot \frac {\partial \eta}{\partial \tilde {\eta}} ^ {- T} = \left(\frac {\partial \tilde {\eta}}{\partial \eta} \cdot \nabla_ {\mu} \mathcal {L}\right) ^ {T}. \tag {13}
$$
See App. F for proof. This differs from the non-natural gradient in two ways. First, the Jacobian of the  $\eta \rightarrow \mu$  map is dropped, as before. Unlike Johnson et al. [29], we do not hand-derive the solution; we employ a straight-through gradient estimator [6] to replace this Jacobian with the identity. Then, the Jacobian of the  $\tilde{\eta} \rightarrow \eta$  map is replaced by its inverse-transpose. This new gradient can be computed without any matrix arithmetic by noting that the inverse of a Jacobian is the Jacobian of the inverse function. Thus Eq. (13) can be computed by replacing the reverse-mode backpropagation through the  $\tilde{\eta} \rightarrow \eta$  map with a forward-mode differentiation through the inverse  $\eta \rightarrow \tilde{\eta}$  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ä & García-Fernández [56] developed a method to parallelize the usually-sequential Kalman smoother algorithm across time for jointly Gaussian data. Their algorithm is not directly applicable to our VI setting where we take expectations over  $q(\theta)$  instead of having fixed parameters  $\hat{\theta}$ , but we derive an 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 a 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-dimensional data. We therefore propose a three-stage training routine: a basic VAE is trained to initialize  $p_{\gamma}(x|z)$ , 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.
Figure 4: Unconstrained generation of 513-dim. speech spectrogram data over  $T = 500$  time-steps (horizontal; models are trained on data with  $T = 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 a solid colorbar meaning a 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  $\hat{\ell}_{\phi}(z_t|x_t;\eta)$  to be uniform at a particular timestep  $t$ , our posterior estimate of  $z_{t}$  is only guided by the prior, which aggregates information across time to produce a smooth estimate of the posterior on  $z_{t}$ .
While discriminative methods may be explicitly trained to impute time series, we use imputation performance as a 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.
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]), is the highest-performing model in their survey; it models data via one-step-ahead prediction, producing probabilities  $p(x_{t}|z,x_{t - 1})$ . This model therefore reconstructs  $x$  using more information than is encoded in  $z$  by skipping over the latent state and directly connecting ground truth to reconstruction, reducing the problem of sequence generation to a 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-performing 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, or interpolating with exponentially decayed observations [11], effectively require training to learn these imputation heuristics. RNNs must thus be trained on missing-ness that is similar to test missing-ness, unlike the SVAE.
Transformers [61] have achieved state-of-the-art generative performance on sequential language modeling. However, Zeng et al. [66] argue that their permutation-invariance results in weak performance 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 a collection of partially-permutable tokens. Lin et al. [42] propose a dynamical VAE with encoder  $q(z_{t}|x_{1:T})$ , decoder  $p(x_{t+1}|x_{1:t}, z_{1:t+1})$ , and latent dynamics  $p(z_{t+1}|x_{1:t}, z_{1:t})$  parameterized by transformers.
Structured VAEs. We, as in Johnson et al. [29], only consider SVAEs where the inference network output factorizes across (temporal) latent variables. Orthogonal to our contributions, Yu et al. [65] investigate the advantages of taking the SVAE beyond this restriction, and building models where the recognition network outputs proxy-likelihood 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 to include discrete latent variables. Second, because they optimize point-estimates  $\hat{\theta}$  of parameters instead of variational factors  $q(\theta)$ , they do not make use of natural gradients. In this point-estimate formulation,
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 use a single discrete state, and thus cannot produce diverse imputations. The SRNN [17] produces varied sequences, but autoregressive generation is sometimes unstable and unrealistic, 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 a novel extension for our variational setting. Finally, their experimental results are confined to toy and synthetic data sets.
The most directly comparable model to the SVAE is the Stochastic Inference Network (SIN [41]), which employs a graphical model prior  $p(z|\theta)$  but estimates  $q(z)$  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-Quantized VAEs (VQ-VAEs [60]) are another family of deep generative models that make use of discrete latent representations. Rather than directly outputting an estimated posterior, the encoder of a VQ-VAE outputs a point in an embedding space which is quantized to one of  $K$  learnable quantization points. The encoder is trained using a straight-through estimator [6]. While discrete SVAE-SLDS states switch among multiple continuous modes, the VQ-VAE representation is purely discrete, limiting information encoded by a latent variable to  $\log_2K$  bits. In order to generate plausible and diverse samples, VQ-VAEs require large values of  $K$ , structured collections of discrete variables, and/or post-hoc training of autoregressive priors [51, 52, 16, 15].
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 (WSJ0 [19]); see Table 2. MOCAP has 84-dimensional data and training sequences of length  $T = 250$ . WSJ0 has 513-dimensional data and training sequences of length  $T = 50$ . See App. A.2 for further details.
Generation quality is judged via a modified Frechet inception distance (FID [18]) metric. We replace the InceptionV3 network with appropriate classifiers for motion capture and speech data (see App. A.3). SVAE-SLDS-Bias runs the SVAE as presented by Johnson et al. [29], with unrolled gradients, 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; see App. A.5 for details.
To demonstrate the SVAE's capacity to handle discrete latent variables in a principled manner, we compare to two DVAE baselines which incorporate discrete variables via biased gradients: the straight-through estimator [6] and the concrete (Gumbel-softmax) distribution [27, 44]. To our knowledge, no 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.
Method	log p(x) ≥ (↑)	Sample FID (↓)	Interpolation FIDs (↓)
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 ± 0.02
SVAE-SLDS-Bias [29]	2.36	34.6 ± 0.7	28.8 ± 0.2	25.8 ± 0.3	6.71 ± 0.12
SVAE-LDS	2.28	34.0 ± 0.3	19.3 ± 0.2	21.9 ± 0.2	7.90 ± 0.13
SIN-SLDS [41]	2.36	33.7 ± 0.4	12.38 ± 0.12	8.97 ± 0.08	3.27 ± 0.05
SIN-LDS [41]	2.33	65.2 ± 1.4	18.3 ± 0.2	15.5 ± 0.2	6.24 ± 0.09
Transformer [42]	2.82	421 ± 11	234 ± 9	228 ± 5	113 ± 5
SRNN [17]	2.94	62.7 ± 0.7	43.5 ± 0.7	24.2 ± 0.6	14.2 ± 0.3
DKS [35]	2.31	136 ± 6	46.7 ± 1.7	33.3 ± 1.1	9.0 ± 0.3
DKS+Concrete	1.70	144 ± 3	88 ± 3	89 ± 2	34.0 ± 1.4
DKS+Straight-Through	2.09	22 ± 3	22.6 ± 0.3	17.15 ± 0.14	13.8 ± 0.2
Audio Spectrogram (WSJ0)
SVAE-SLDS	1.54	9.61 ± 0.15	7.5 ± 0.2	8.14 ± 0.12	4.88 ± 0.08
SVAE-SLDS-Bias	1.45	18.6 ± 0.2	15.0 ± 0.2	15.2 ± 0.2	7.6 ± 0.12
SVAE-LDS	1.56	19.1 ± 0.3	17.9 ± 0.2	16.6 ± 0.3	7.2 ± 0.3
SIN-SLDS	1.53	20.0 ± 0.4	17.2 ± 0.3	14.9 ± 0.3	9.5 ± 0.2
SIN-LDS	1.54	17.8 ± 0.2	17.21 ± 0.11	13.2 ± 0.4	10.1 ± 0.2
Transformer	1.88	10.0 ± 0.2	12.0 ± 0.3	8.2 ± 0.2	5.7 ± 0.4
SRNN	1.94	23.6 ± 0.3	19.4 ± 0.5	17.4 ± 0.3	12.7 ± 0.4
DKS	1.55	12.9 ± 0.2	10.8 ± 0.2	10.8 ± 0.14	7.7 ± 0.05
DKS+Concrete	1.45	16.6 ± 0.2	12.8 ± 0.2	11.3 ± 0.2	8.2 ± 0.2
DKS+Straight-Through	1.48	15.51 ± 0.13	10.07 ± 0.11	9.02 ± 0.18	6.29 ± 0.13
Table 2: Comparison of model performance on log-likelihood (higher is better), FIDs of unconditionally generated samples (lower is 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. 1 we show the SVAE-SLDS's learned discrete encoding of several joint tracking sequences. While "sitting" sequences are dominated by a single dynamic mode, walking is governed by a 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 we show example generated sequences from each model of audio data. Like true speech, the SVAE-SLDS moves between discrete structures over time representing individual sounds. In contrast, the SIN-SLDS [41] and DKS+Concrete baselines collapse to a single discrete modality, blending together continuous dynamics. While the DKS+Straight-Through model does not collapse, it uses discrete states too coarsely to inform the high-frequency dynamics of speech.
Interpolation. Amortized VI cannot easily infer  $q(z_{t})$  at time steps where the observations  $x_{t}$  are missing. Thus, given observations at a subset of times  $x_{\mathrm{obs}}$ , we can encode to obtain  $q(z_{\mathrm{obs}})$  and infer the missing latent variables by drawing from the generative prior  $p(z_{\mathrm{missing}}|z_{\mathrm{obs}},\theta)$ . Because baseline models parameterize  $p(z|\theta)$  by one-directional neural networks, they can only condition  $z_{\mathrm{missing}}$  on  $z_{\mathrm{obs}}$  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 in 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 a single discrete state in training, resulting in uniform imputations that lack diversity. Example imputations for all models are provided in App. Fig. 6.
Transformers for time series. The permutation-invariance of transformers is visible in its varied performance on these two tasks. A 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 transformer-generated samples unrealistic (see Fig. 6 and Table 2).
The SVAE is uniquely situated at the intersection of flexible, high-dimensional modeling and interpretable 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.
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.