SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[1] Akiba, T., Sano, S., Yanase, T., Ohta, T., and Koyama, M. Optuna: A next-generation hyperparameter optimization framework. In Proceedings of the 25th ACM SIGKDD International Conference on 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, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450, 2016.
[4] Barber, D. and Chiappa, S. 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 or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 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://github.com/google/jax.
[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 Ionescu, 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):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. A 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, I. Jukebox: A generative model for music. arXiv preprint arXiv:2005.00341, 2020.
[16] Dieleman, S., van den Oord, A., and Simonyan, K. The challenge of realistic music generation: modelling raw audio at scale. In Advances in Neural Information Processing Systems, 2018.
[17] Fraccaro, M., Sønderby, S. K., Paquet, U., and Winther, O. Sequential neural models with stochastic layers. In Advances in Neural Information Processing Systems, 2016.
[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, L., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems, 2017.
[62] Wainwright, M. J., Jordan, M. I., et 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. In 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, L., 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.
A Experimental Protocol
A.1 Implementation
Code can be found at https://github.com/hbendekgey/SVAE-Implicit. All methods were implemented with the JAX library [8].
A.2 MOCAP Dataset
Our motion capture experiments were 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  $25\mathrm{hz}$ . Joint positions are recorded in meters relative to 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  $1\mathrm{mm}$  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 set 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-1, which were removed.
We model the likelihood  $p(x|z)$  of joint coordinates using independent Gaussian distributions. Following the work of [55] we fit a single global variance parameter to 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 InceptionV3 [58] network typically used to compute Frechet Inception Distance [18] 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/scenerios, 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 [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 WSJ0 Dataset
The WSJ0 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 513-dimensional observations at each time step. We modeled the spectrogram data via a Gamma distribution with fixed concentration parameter 2, equivalent to modeling the complex Fourier transform of the data with a mean 0 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.
(SVAE-SLDS)
(SIN-SLDS)
(SRNN)
(SVAE-LDS)
(SVAE-SLDS-bias)
(Transformer)
(DKS+Concrete)
(DKS+Straight-Through)
(DKS)
Figure 6: 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 (see App. A.6 for methodology). Segmentation colorbars are shown below models which use discrete latent variables; all non-SVAE-SLDS models collapse to a single discrete state throughout training. Amortized VI models draw missing data conditioned on previous data, introducing discontinuities at the end of the missing window.
Residual Block	Downsampling block	Mocap Classifier
Activation: Relu
Conv. window: 3
Conv.
Batch Norm
Activation
Conv.
Batch Norm
Activation
Add input	Conv.
Batch Norm
Activation
Conv. (stride 2)
Batch Norm
Activation
Add downsampled input	Residual block 128
Downsampling block 128
Residual block 128
Residual block 128
Global average pool
Dense
Table 3: Summary of classifier network architecture used to compute FID scores for Mocap data.
FID Scores 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. For architecture, we used a efficientnet_v2_s model from the torchvision package [46] on the data in log space.
Figure 7: Segmentation and Imputation. An SLDS SVAE is trained on synthetic data with feature dimension 100, sequence length  $T = 250$  and 4 distinct dynamics. We show the SVAE's latent representation of the 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 we show model imputations when time steps 20 to 70 are not provided. For the DKS and SRNN, we show one imputation using the infill-with-zeros method, and 2 imputations drawn from the learned latent variable prior, conditioned on the preceding latent variables for which a posterior is properly defined.
A.4 Discrete Deep Kalman Smoother Baselines
To our knowledge, time series VAEs that utilize Concrete [44] or straight-through gradient estimators have not been previously evaluated in literature. In order to perform a fair comparison to these approaches to gradient estimation with discrete latent variables, we have devised a new model that extends the Deep Kalman Smoother (DKS) [37] to incorporate both discrete and continuous latent variables. We can define the generative model as follows:
$$
p _ {\theta_ {k}} \left(k _ {t} \mid z _ {t - 1}, k _ {t - 1}\right) = \operatorname {C a t} \left(\lambda_ {\theta_ {k}} \left(z _ {t - 1}, k _ {t - 1}\right)\right) \tag {14}
$$
$$
p _ {\theta_ {z}} \left(z _ {t} \mid z _ {t - 1}, k _ {t}\right) = \mathcal {N} \left(\mu_ {\theta_ {z}} \left(z _ {t - 1}, k _ {t}\right), \sigma_ {\theta_ {z}}\right) \tag {15}
$$
$$
p _ {\gamma} \left(x _ {t} \mid z _ {t}\right) = \mathcal {N} \left(\mu_ {\gamma} \left(z _ {t}\right), \sigma_ {\theta_ {x}}\right) \tag {16}
$$
Here  $\lambda_{\theta_k}(z_{t - 1},k_{t - 1}),\mu_{\theta_z}(z_{t - 1},k_t)$  and  $\mu_{\gamma}(z_{t})$  are neural networks, while  $\sigma_{\theta_z}$  and  $\sigma_{\theta_x}$  are learnable variance parameters.  $\mu_{\theta_x}(z_t)$  uses the same decoder architecture as our implementation of the standard DKS model.  $\lambda_{\theta_k}(z_{t - 1},k_{t - 1}),\mu_{\theta_z}(z_{t - 1},k_t)$  use the following architecture in our experiments:
$\lambda (\cdot)$
Dense 64
Gelu
$\mu (\cdot)$
The inference model can be written as:
$$
q _ {\phi_ {k}} \left(k _ {t} \mid x _ {t: T}, z _ {t - 1}, k _ {t - 1}\right) = \operatorname {C a t} \left(\lambda_ {\phi_ {k}} \left(g _ {\phi} \left(x _ {t: T}\right), z _ {t - 1}, k _ {t - 1}\right)\right) \tag {17}
$$
$$
q _ {\phi_ {z}} \left(z _ {t} \mid x _ {t: T}, z _ {t - 1}, k _ {t}\right) = \mathcal {N} \left(\mu_ {\phi_ {z}} \left(g _ {\phi} \left(x _ {t: T}\right), z _ {t - 1}, k _ {t}\right), \sigma_ {\phi_ {z}}\right) \tag {18}
$$
Here  $g_{\phi}(x_{t:T})$  is the same encoder network used by our DKS implementation (see Table 4), while  $\lambda_{\phi_k}(\cdot)$  and  $\mu_{\phi_k}(\cdot)$  are networks with the same architecture as  $\lambda_{\theta_k}(\cdot)$  and  $\mu_{\theta_k}(\cdot)$  shown above. We note that the output of  $g_{\phi}(x_{t:T})$  is shared by both parts of the inference model.
In our experiments we match the latent dimension and number of discrete states to the corresponding SVAE SLDS model. For "DKS + Straight-Through" we use straight-through estimates of the gradient of samples of  $k_{t}$ . In this case we can directly compute the KL-divergence between  $p_{\theta_k}(k_t \mid z_{t-1}, k_{t-1})$  and  $q_{\phi_k}(k_t \mid x_{t:T}, z_{t-1}, k_{t-1})$  in closed form.
For "DKS + Concrete" we replace the Categorical distribution in equations 14 and 17 with a Concrete distribution. As the KL-divergence between concrete distributions does not have a 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 1. 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 a batch size of  $B = 128$  and learning rate  $10^{-3}$  (the transformer DVAE uses learning rate  $10^{-4}$ , which improved performance). We train all methods for 200 epochs on MOCAP and 100 epochs on WSJ0 (including VAE-pre-training 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, we 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 as implemented in bnpy [24, 25], to initialize the parameters  $\eta$  of  $q(\theta)$ .
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 "teacher-forcing" LSTM 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 framework).
For each model, we 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  $x_{\mathrm{missing}}$ ,  $x_{\mathrm{obs}}$  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  $x_{\mathrm{obs}} = x_{0.0 - 0.2} \cup x_{0.8 - 1.0}$  (first  $20\%$  and last  $20\%$  as in our results) leaving  $x_{\mathrm{missing}} = x_{0.2 - 0.8}$ .
To interpolate missing data we begin by encoding  $x_{\mathrm{obs}}$  to obtain  $q(z_{\mathrm{obs}})$  (in the DKS/transformer/SRNN case) or  $\hat{\ell}_{\phi}(z_{\mathrm{obs}}|x_{\mathrm{obs}})$  (in the SVAE/SIN case), where  $z_{\mathrm{obs}}$  is the subset of  $z$  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 inflated data into the encoders.
For the DKS and its extensions, we then sample the remaining  $z$  from the generative prior given the last encoded latent variable that precedes the missing chunk. In the regime above, we therefore sample from  $p(z_{0.2 - 0.8}|z_{0.0 - 0.2})$  as we cannot easily condition on later timesteps of a forward deep recurrent model. Finally, we decode  $z_{\mathrm{missing}} \cup z_{\mathrm{obs}}$  to obtain the reconstruction and interpolation.
For the SRNN and the transformer, due to their autoregressive nature we sample a value of  $z$  at a single timestep (i.e. the first timestep in  $z_{0.2 - 0.8}$ ) and decode to obtain a new  $x$ . We then repeat this
Method	Encoder arch.	Enc. Params	Decodder arch.	Dec. Params	Latent params	Runtime
SVAE-SLDS	Dense 256Gelu [22]Layer Norm [3]Dense 128GeluLayer NormDense 16	57,504	Dense 128GeluLayer NormDense 256GeluLayer NormDense	57,640	30,668(50 states)	16.4 GPU hours(MOCAP)12.5 GPU hours(WSJ0)
SVAE-LDS	Dense 256GeluLayer NormDense 128GeluLayer NormDense 16	57,504	Dense 128GeluLayer NormDense 256GeluLayer NormDense	57,640	716	9 GPU hours(MOCAP)4.6 GPU hours(WSJ0)
SIN-SLDS	Dense 256GeluLayer NormDense 128GeluLayer NormDense 16	57,504	Dense 128GeluLayer NormDense 256GeluLayer NormDense	57,640	34,018(50 states)	7.3 GPU hours(MOCAP)3.6 GPU hours(WSJ0)
SIN-LDS	Dense 256GeluLayer NormDense 128GeluLayer NormDense 16	57,504	Dense 128GeluLayer NormDense 256GeluLayer NormDense	57,640	1516	7.3 GPU hours(MOCAP)3.6 GPU hours(WSJ0)
Transformer	Dense 256GeluLayer NormPositional Enc.Self AttentionLayer NormDense 256Gelu + SkipLayer NormDense	360,480	Dense 256GeluLayer NormPositional Enc.Self AttentionLayer NormCross AttentionLayer NormDense 256GeluLayer NormDense	532,136	518,688	2 GPU hours(MOCAP)1.6 GPU hours(WSJ0)
SRNN	Dense 256GeluLayer NormDense 128Shared LSTM 128Reverse-LSTM 128Dense	252,288	Shared LSTM 128Dense 256GeluLayer NormDense	190,888	87,680	3.7 GPU hours(MOCAP)2.9 GPU hours(WSJ0)
DKS(+ Discrete)	Dense 256GeluLayer NormReverse-LSTM 128Dense 192	244,160	Dense 128GeluLayer NormDense 256GeluLayer NormDense	57,640	38,144(43,724)	3.3 GPU hours(MOCAP)2.4 GPU hours(WSJ0)
Table 4: Summary of network architectures used for human motion capture and speech experiments. "Latent parameters" are those which specify  $p(z|.)$ . For the SRNN model, one LSTM layer is shared between inference and generation; we denote this layer "shared LSTM" and count its parameters among those for the decoder. SVAE-SLDS-Bias uses the same architecture as the SVAE-SLDS but takes 36.5 GPU hours MOCAP and 28.5 GPU hours on WSJ0 to train.
process so as to update the hidden state  $h_t$  while holding  $z_t$  in the observed and already-interpolated region constant.
For the SVAE and SIN methods, we do inference in  $q_{\phi}(z|x)$  as in training time, except setting the inference network potentials at missing timesteps to 0 (thus setting  $\hat{\ell}_{\phi}(z|x)$  as uniform over the reals). For models with discrete random variables the block updating routine is prone to local optima, so we initialize according to the following heuristic: we encode  $x_{0.0 - 0.2}$  and  $x_{0.8 - 1.0}$  and perform block updating to obtain  $q(k_{0.0 - 0.2})$  and  $q(k_{0.8 - 1.0})$ . We then sample the endpoints  $k_{0.2}, k_{0.8}$  from these variational factors. We perform HMM inference to sample from  $p(k_{0.2 - 0.8}|k_{0.2}, k_{0.8})$ , 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  $z$  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. 1 we show the discrete states are a 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 an extended version of Fig. 5, showcasing the performance of all models on interpolating human motion capture data.
In Fig. 7 we compare our model's interpolation ability to several baselines on synthetic data. The data is constructed by evaluated a Laplacian distribution on a 100-dimensional grid, where observations evolve over time by either increasing or 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 the RNN-based models, if we simply infill the missing observations with 0, the encoder-decoder structure simply reconstructs the 0s. Drawing the missing region from the model's prior distribution yields a 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.
In this section we provide a glossary of notation used in the main paper and the appendix.
B.1 Notation in main paper