Title: Unbiased learning of deep generative models with structured discrete representations
Venue: NeurIPS 2023
Accepted: yes
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

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 SV AE 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 SV AE performs competitively while learning interpretable and structured discrete data representations.

1 Introduction

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 SV AEs.

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 SV AE 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 SV AEs 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].

Walking

Sitting

Posing

Taking photo

States predictive of

walking

States predictive of

sitting

States predictive of

posing

States predictive of

photo

Other common states

Uncommon states

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 SV AE 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 SV AE performance to state-of-the-art recurrent neural networkand 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 | θ ) is defined by a graphical model (as in Fig. 2) with parameters θ , and z is processed by a (deep) neural network with weights γ to compute the data likelihood p γ ( x | z ) .

Exact evaluation or simulation of the posterior p γ ( z, θ | 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):

Here, q ( θ ; η ) q ( z ; ω ) ≈ p γ ( z, θ | x ) are known as variational factors . We parameterize these distributions via arbitrary exponential families with natural parameters η, ω . This implies that

An exponential family is log-linear in its sufficient statistics t ( z ) , where the normalizing constant Z ( ω ) ensures it is a proper distribution (see App. C.1 for properties of exponential families). For models where p γ ( x | z ) has a restricted conjugate form (rather than a deep neural network), we can maximize Eq. (1) by alternating optimization of η, ω ; 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 ; ω ) to the latent variables in that batch, and updating q ( θ ; η ) by the resulting (natural) gradient.

Amortized VI. Because it is costly to optimize Eq. (1) with respect to ω for each batch of data, VAEs employ amortized VI [32, 47, 53] to approximate the parameters of the optimal q ( z ; ω ) via a neural network encoding of x . The inference network weights ϕ for this approximate posterior q ϕ ( 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 ( θ ) , the true optimizer (across all probability distributions) of Eq. (1) is given by q ( z ) ∝ p 0 ( z ; η ) p γ ( x | z ) , where p 0 ( z ; η ) ∝ exp { E q ( θ ; η ) [log p ( z | θ )] } 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 γ ( x | z ) . Rather than approximating

𝑧

!

𝑧

"

𝑧

#

𝑥

!

𝑥

"

𝑥

#

𝑧

!

𝑧

"

𝑧

#

𝑥

!

𝑥

"

𝑥

#

𝑧

!

𝑧

"

𝑧

#

𝑘

"

𝑘

#

𝑥

!

𝑥

"

𝑥

#

𝑘

"

𝑘

#

𝑧

!

𝑧

"

𝑧

#

𝑥

!

𝑥

"

𝑥

#

Linear	Dynamical	System	(LDS)

𝑘

!"#

𝑘

!

𝑘

!$#

𝑧

!"#

𝑧

!

𝑘

!"#

𝑘

!

𝑘

!$#

𝑧

!"#

𝑧

!

Switching	Linear	Dynamical	System	(SLDS)

𝑝(𝑧|𝜃)

𝑞

%

(𝑧|𝑥; 𝜂)

Algorithm 1

Structured Mean Field Variational Inference

Require:

graphical model potentials

η

, network poten-

tials

λ

=

encode

(

x

;

ϕ

)

defining

ˆ

ℓ

ϕ

(

z

|

x

)

.

(

µ

1

, . . . , µ

M

)

init_expected_stats

(

η, λ

)

while

not converged

do

for

m

1

to

M

do

ω

m

MF

(

µ

-

m

;

η

)

µ

m

BP

(

ω

m

;

η, λ

)

return

ω

=

concat

(

ω

1

, . . . , ω

M

)

this whole posterior as in amortized VI, the SV AE [29] only approximates the likelihood function, and explicitly multiplies it by the known expected prior to obtain an approximate posterior.

In more detail, the SV AE approximates the likelihood ℓ ( z ) = p γ ( x | z ) with a function parameterized by a neural network ˆ ℓ ϕ ( z | x ) . The optimal posterior given this approximation is then equal to

The posterior that optimizes the surrogate loss ˆ L [ · , ϕ ] of Eq. (3) is q ϕ ( z | x ; η ) ∝ p 0 ( z ; η ) ˆ ℓ ϕ ( z | x ) . If ˆ ℓ ϕ ( z | x ) is chosen to be conjugate to p ( z | θ ) , this multiplication and normalization is easy for many exponential families. Note that ˆ ℓ ϕ ( z | x ) does not need to be normalized, as any multiplicative factors would disappear when the posterior is normalized. The overall ELBO is then L [ q ( θ ; η ) q ϕ ( z | x ; η ) , γ ] .

If the encoder's approximate likelihood ˆ ℓ ϕ is (up to normalization) close to the true likelihood ℓ , the surrogate loss ˆ L will closely approximate the true loss L , and there will be little amortization gap. SVAE inference has the advantage that q ϕ ( z | x ; η ) depends on the learned posterior q ( θ ; η ) 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 V AE [57] and related hierarchical V AEs [59, 12] also incorporate generative parameters in amortized variational inference, but impose a restricted generative hierarchy that is less flexible and interpretable than the SV AE.

Example 1: Standard Normal. For a basic V AE, θ is fixed and z ∼ N (0 , I ) . The SVAE inference network outputs a Gaussian ˆ ℓ ϕ ( z | µ = µ ( x ; ϕ ) , τ = τ ( x ; ϕ )) ∝ N ( z ; µ, diag ( τ -1 )) with mean µ and inverse-variance (precision) τ . The product of this Gaussian with the standard normal prior has a simple form: q ϕ ( z | x ) = N ( z ; τ τ +1 µ, diag (( τ +1) -1 )) . This reparameterization of the standard V AE 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 ∼ N ( A t z t -1 + b t , Q t ) , z 1 ∼ N ( µ 1 , Σ 1 ) . In this case, we expand the exponential family distribution q ( z ; ω ) as a sum across time steps:

where ω = concat ( ω t ) and the prior p ( z | θ ) belongs to this family. The prior induces temporal dependence between the z t , but we assume the likelihood factorizes as p γ ( x | z ) = ∏ t p γ ( 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 | θ ) are learned through optimization of q ( θ ) , the inference routine shares those parameters, improving accuracy. These advantages were not present in the standard normal V AE, where the prior on z lacks structure and is not learned. Inference, normalization, and sampling of q ϕ ( z | x ; η ) 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 | θ ) 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 F :

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 ϕ ( z | x ; η ) 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 m =1 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 ϕ ( z m | x ; η ) for each mean field cluster by maximizing ˆ L [ q ( θ ; η ) ∏ m q ( z m ) , ϕ ] . The optimal q ϕ ( z m | x ; η ) inherit the structure of the joint optimizer q ϕ ( z | x ; η ) , 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 ω :

Here, BP is a belief propagation algorithm which computes expected statistics µ m for cluster m , and the linear mean field function MF updates parameters of cluster m 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 [ q ( θ ) ∏ m q ( z m )] .

3.2 Reparameterization and Discrete Latent Variables

While optimizing q ϕ ( z m | x ; η ) at inference time requires some computational overhead, it allows us to bypass the typical obstacles to training V AEs 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 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 γ ( 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 ∼ Cat ( π ) , z ∼ N ( µ k , Σ k ) . The likelihood p γ ( x | z ) directly conditions on only the continuous latent variable z . Variational inference produces disentangled factors q ϕ ( z | x ; η ) q ϕ ( k | x ; η ) , and we evaluate likelihoods by decoding samples from q ϕ ( z | x ; η ) , without sampling q ϕ ( k | x ) .

Example 4: Switching Linear Dynamical System (SLDS). Consider a set of discrete states which evolve according to a Markov chain k 1 ∼ Cat ( π 0 ) , k t ∼ Cat ( π k t -1 ) , and a continuous state evolving according to switching linear dynamics: z 0 ∼ N ( µ 0 , Σ 0 ) , z t ∼ 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 ϕ ( z | x ; η ) q ϕ ( k | x ; η ) 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,

and ¯ θ j = E q ( θ ; η ) [ θ j ] be the expected (natural) parameters of the LDS for discrete state j . Structured VI updates the natural parameters of discrete states ω k tj , and continuous states ω z t ,z t +1 , as follows:

4 Stable and Memory-Efficient Learning via Implicit Gradients

When q ϕ ( z | x ) is computed via closed-form inference, gradients of the SV AE ELBO may be obtained via automatic differentiation. This requires backpropagating through the encoder and decoder networks, as well as through reparameterized sampling z ∼ q ϕ ( z | x ; η ) 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 ∂ω ∂η , ∂ω ∂ϕ without storing intermediate states. We focus on gradients with respect to η for compactness, but gradients with respect to ϕ are computed similarly. Let ω (1) , . . . , ω ( L ) be the sequence of ω values produced during the 'forward' pass of block coordinate ascent, where ω ( L ) are the optimized structured VI parameters. The IFT expresses gradients via the solution of a set of linear equations:

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 ( ω ) = 0 . For an SLDS with latent dimension D and K discrete states, ω has O ( K + D 2 ) parameters at each time step. Over T time steps, ∂g ∂ω 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 matrixvector products ( I -A ) v ′ 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 SV AEs. The resulting algorithm resembles unrolled gradient estimation, but we repeatedly backpropagate through updates at the endpoint of optimization instead of along the optimization trajectory.

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, ω ( L ) is not a stationary point of g ( ω ) and therefore Eq. (9) is not guaranteed to converge as J →∞ . For batch learning, waiting for all VI routines to converge to a (local) optimum might be prohibitively slow, so we might halt VI before ω ( 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

2

5

10

20

50

Number of Block Update (Forward Pass) Steps

10

1

10 0

rMSE/No Solve rMSE

Implicit + Cap + Thresh

Implicit

Unrolled

2

5

10

20

50

Number of Block Update (Forward Pass) Steps

10

3

10

2

10

1

10 0

10 1

rMSE/No Solve rMSE

Implicit + Cap + Thresh

Implicit

Unrolled

10 0

10 1

10 2

10 3

10 4

Step

60000

50000

40000

30000

20000

Loss

Non-Nat

Unbiased Nat

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).

Experiments. In Fig. 3 (left, middle) we compute the accuracy of different approximate gradient estimators for training a SV AE-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 naïve 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 ϕ , suffering high memory overhead. They compute neither unrolled nor implicit gradients with respect to generative model parameters η ; 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 η of q ( θ ; η ) , given by:

Johnson et al. [29] demonstrate the advantages of natural gradients for the SV AE, drawing parallels to the natural gradients of stochastic VI (SVI [23]). SVI extends the variational EM algorithm to mini-batch learning: similar to the SV AE, it fits q ( z ) in an inner optimization loop and learns q ( θ ; η ) 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 L , whereas the SVAE uses a 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 SV AE.

Let µ η be the expected sufficient statistics of q ( θ ; η ) . Exponential family theory tells us that ∂µ ∂η = F η [30, 45], allowing Johnson et al. [29] to derive the natural gradients of the SV AE loss:

This gradient differs from the SVI gradient by the final term: because SVI's inner loop optimizes ω with respect to the true loss L , ∂ L ∂ω = 0 for conjugate models. Johnson et al. [29] train their SV AE 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 SV AE. First, in the structured mean field case ∂ω ∂η 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 η into an unconstrained space, and computing natural gradients with respect to those new parameters. Letting ˜ η be an unconstrained reparameterization of η , such as η = Softplus { ˜ η } = log(1 + e ˜ η ) for a positive precision parameter, we have:

See App. F for proof. This differs from the non-natural gradient in two ways. First, the Jacobian of the η → µ 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 ˜ η → η 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 ˜ η → η map with a forward-mode differentiation through the inverse η → ˜ η 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 SV AEs 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 ( θ ) instead of having fixed parameters ˆ θ , 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 γ ( 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.

Example Sequence

Transformer Generation

SVAE SLDS Generation (Ours)

SIN SLDS Generation

DKS Generation

SRNN Generation

DKS+Straight Through Generation

DKS+Concrete Generation

20

15

10

5

0

5

10

decibles

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 t | x t ; η ) 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.

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]), 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 SV AE.

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 SV AEs where the inference network output factorizes across (temporal) latent variables. Orthogonal to our contributions, Yu et al. [65] investigate the advantages of taking the SV AE 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 SV AE to include discrete latent variables. Second, because they optimize point-estimates ˆ θ of parameters instead of variational factors q ( θ ) , they do not make use of natural gradients. In this point-estimate formulation,

(SVAE-SLDS)

(SIN-SLDS)

(SRNN)

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 | θ ) 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 SV AE-SLDS states switch among multiple continuous modes, the VQVAE representation is purely discrete, limiting information encoded by a latent variable to log 2 K 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].

8 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 (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 Frechét 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 V AEs 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.

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 obs, we can encode to obtain q ( z obs ) and infer the missing latent variables by drawing from the generative prior p ( z missing | z obs , θ ) . Because baseline models parameterize p ( z | θ ) by one-directional neural networks, they can only condition z missing on z obs at previous time steps, leading to discontinuity at the end of the missing window. For further specifications and for details of the SV AE 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 SV AE-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).

9 Discussion

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.

References

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.

Amari, S.-I. Natural gradient works efficiently in learning. Neural computation , 10(2):251-276, 1998.

Ba, J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

Barber, D. and Chiappa, S. Unified inference for variational Bayesian linear Gaussian state-space models. In Advances in Neural Information Processing Systems , 2007.

Beal, M. J. Variational algorithms for approximate Bayesian inference . University College London, 2003.

Bengio, Y., Léonard, N., and Courville, A. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432 , 2013.

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.

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 .

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.

Catalin Ionescu, Fuxin Li, C. S. Latent structured models for human pose estimation. In International Conference on Computer Vision , 2011.

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.

Child, R. Very deep VAEs generalize autoregressive models and can outperform them on images. In International Conference on Learning Representations , 2021.

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.

Cremer, C., Li, X., and Duvenaud, D. Inference suboptimality in variational autoencoders. In International Conference on Machine Learning , 2018.

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.

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.

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.

Fréchet, M. Sur la distance de deux lois de probabilité. In Annales de l'ISUP , volume 6, pp. 183-198, 1957.

Garofolo, J., Graff, D., Paul, D., and Pallett, D. Csr-i (wsj0) sennheiser ldc93s6b. Web Download. Philadelphia: Linguistic Data Consortium , 1993.

Ghahramani, Z. and Hinton, G. E. Variational learning for switching state-space models. Neural Computation , 12:831-864, 2000.

Girin, L., Leglaive, S., Bie, X., Diard, J., Hueber, T., and Alameda-Pineda, X. Dynamical variational autoencoders: A comprehensive review. Foundations and Trends in Machine Learning , 15(1-2):1-175, 2021.

Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.08415 , 2016.

Hoffman, M. D., Blei, D. M., Wang, C., and Paisley, J. Stochastic variational inference. Journal of Machine Learning Research , 2013.

Hughes, M. C. and Sudderth, E. Memoized online variational inference for dirichlet process mixture models. In Advances in Neural Information Processing Systems , 2013.

Hughes, M. C., Stephenson, W. T., and Sudderth, E. Scalable adaptation of state complexity for nonparametric hidden markov models. In Advances in Neural Information Processing Systems , 2015.

Ionescu, C., Papava, D., Olaru, V., and Sminchisescu, C. Human3.6m: Large scale datasets and predictive methods for 3d human sensing in natural environments. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2014.

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with Gumbel-softmax. In International Conference on Learning Representations , 2017.

Ji, G., Sujono, D., and Sudderth, E. B. Marginalized stochastic natural gradients for black-box variational inference. In International Conference on Machine Learning , 2021.

Johnson, M. J., Duvenaud, D. K., Wiltschko, A., Adams, R. P., and Datta, S. R. Composing graphical models with neural networks for structured representations and fast inference. In Advances in Neural Information Processing Systems , 2016.

Khan, M. E. and Nielsen, D. Fast yet simple natural-gradient descent for variational inference in complex models. In 2018 International Symposium on Information Theory and Its Applications (ISITA) , 2018.

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015.

Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. In International Conference on Learning Representations , 2014.

Kingma, D. P. and Welling, M. An introduction to variational autoencoders. Foundations and Trends in Machine Learning , 12(4):307-392, 2019.

Krantz, S. G. and Parks, H. R. The implicit function theorem: history, theory, and applications . Springer Science & Business Media, 2002.

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.

Krishnan, R., Liang, D., and Hoffman, M. On the challenges of learning with inference networks on sparse, high-dimensional data. In AISTATS , pp. 143-151, 2018.

Krishnan, R. G., Shalit, U., and Sontag, D. Deep kalman filters. arXiv preprint arXiv:1511.05121 , 2015.

Kschischang, F. R., Frey, B. J., and Loeliger, H.-A. Factor graphs and the sum-product algorithm. IEEE Trans. Info. Theory , 47(2):498-519, February 2001.

Lerner, U. and Parr, R. Inference in hybrid networks: Theoretical limits and practical algorithms. In Conference on Uncertainty in Artificial Intelligence , 2001.

Li, C., Welling, M., Zhu, J., and Zhang, B. Graphical generative adversarial networks. In Advances in Neural Information Processing Systems , 2018.

Lin, W., Hubacher, N., and Khan, M. E. Variational message passing with structured inference networks. In International Conference on Learning Representations , 2018.

Lin, X., Bie, X., Leglaive, S., Girin, L., and Alameda-Pineda, X. Speech modeling with a hierarchical transformer dynamical vae. In IEEE International Conference on Acoustics, Speech and Signal Processing , 2023.

Lorraine, J., Vicol, P., and Duvenaud, D. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Statistics , 2020.

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.

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.

Marcel, S. and Rodriguez, Y. Torchvision the machine-vision package of torch. In Proceedings of the 18th ACM international conference on Multimedia , 2010.

Mnih, A. and Gregor, K. Neural variational inference and learning in belief networks. In International Conference on Machine Learning , 2014.

Pearl, J. Probabilistic Reasoning in Intelligent Systems . Morgan Kaufman, San Mateo, 1988.

Rajeswaran, A., Finn, C., Kakade, S. M., and Levine, S. Meta-learning with implicit gradients. In Advances in Neural Information Processing Systems , 2019.

Ranganath, R., Gerrish, S., and Blei, D. Black Box Variational Inference. In Proceedings of the Seventeenth International Conference on Artificial Intelligence and Statistics , 2014.

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.

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.

Rezende, D. J., Mohamed, S., and Wierstra, D. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning , 2014.

Richardson, L. F. Ix. the approximate arithmetical solution by finite differences of physical problems involving differential equations, with an application to the stresses in a masonry dam. Philosophical Transactions of the Royal Society of London. Series A, Containing Papers of a Mathematical or Physical Character , 210(459-470):307-357, 1911.

Rybkin, O., Daniilidis, K., and Levine, S. Simple and effective vae training with calibrated decoders. In International Conference on Machine Learning , 2021.

Särkkä, S. and García-Fernández, Á. F. Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control , 2020.

Sønderby, C. K., Raiko, T., Maaløe, L., Sønderby, S. K., and Winther, O. Ladder variational autoencoders. In Advances in Neural Information Processing Systems , 2016.

Szegedy, C., Vanhoucke, V., Ioffe, 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.

Vahdat, A. and Kautz, J. NVAE: A deep hierarchical variational autoencoder. In Advances in Neural Information Processing Systems , 2020.

Van Den Oord, A., Vinyals, O., et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems , 2017.

Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. In Advances in Neural Information Processing Systems , 2017.

Wainwright, M. J., Jordan, M. I., et al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning , 1, 2008.

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.

Young, D. M. Iterative Solution of Large Linear Systems . Academic Press, New York, 1971.

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.

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.

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 25hz. 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 1mm 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)

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.

SVAE-SLDS

SVAE-SLDS-Bias

DKS (zero fill)

SRNN (zero fill)

z

1

z

2

0

20

40

60

80

DKS (from prior)

SRNN (from prior)

0

50

100

150

200

Time

State

Discrete Latent Variable

0

10

20

30

40

50

60

70

80

Time

0

10

20

30

40

50

60

70

80

Time

0

10

20

30

40

50

60

70

80

Time

0

10

20

30

40

50

60

70

80

Time

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:

Here λ θ k ( z t -1 , k t -1 ) , µ θ z ( z t -1 , k t ) and µ γ ( z t ) are neural networks, while σ θ z and σ θ x are learnable variance parameters. µ θ x ( z t ) uses the same decoder architecture as our implementation of the standard DKS model. λ θ k ( z t -1 , k t -1 ) , µ θ z ( z t -1 , k t ) use the following architecture in our experiments:

The inference model can be written as:

Here g ϕ ( x t : T ) is same encoder network used by our DKS implementation (see Table 4), while λ ϕ k ( · ) and µ ϕ k ( · ) are networks with the same architecture as λ θ k ( · ) and µ θ k ( · ) shown above. We note that the output of g ϕ ( 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 θ k ( k t | z t -1 , k t -1 ) and q ϕ k ( k t | 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 V AE 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 η of q ( θ ) .

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 missing , x 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 obs = x 0 . 0 -0 . 2 ∪ x 0 . 8 -1 . 0 (first 20% and last 20% as in our results) leaving x missing = x 0 . 2 -0 . 8 .

To interpolate missing data we begin by encoding x obs to obtain q ( z obs ) (in the DKS/transformer/SRNN case) or ˆ ℓ ϕ ( z obs | x obs ) (in the SVAE/SIN case), where z 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 infilled 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 missing ∪ z 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

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 ϕ ( z | x ) as in training time, except setting the inference network potentials at missing timesteps to 0 (thus setting ˆ ℓ ϕ ( 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.

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

̸

B.2 Notation in this appendix

C Exponential families and mean field objective

To rigorously define the SVAE inference routine, in Sec. C.1 we (re-)introduce exponential family and factor graph notation for defining distributions. Then, in Sec. C.2 we derive the optimizer of ˆ L [ . . . ] 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 z ∈ Z , an exponential family is a family of distributions characterized by a statistic function t ( z ) : Z → R n with probabilities given by

Where η ∈ R n is a parameter vector known as the natural parameters , and the log partition function log Z ( η ) is given by

and ensures that the distribution properly normalizes (to simplify notation, here we focus on continuous z but this definition can be easily extended to discrete observations by replacing the integral with a sum and and restricting the output of t ( z ) to some discrete set). A single assignment of values to η characterizes a member of this exponential family, and so searching over a family means searching over valid values of η ∈ 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 is its expected statistics:

Proof.

Proposition C.2. Given two distributions p ( z ; η 1 ) , p ( z ; η 2 ) in the same exponential family with parameters η 1 and η 2 , the KL divergence between the two distributions is given by

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 ( θ ; η )

we say that p ( z | θ ) is conjugate to p ( θ ; η ) if one distribution's parameters is the other distribution's statistics, namely

Where ˆ t ( θ ) is the first dim ( t ( z )) entries of t ( θ ) and c is a constant.

Note that t ( θ ) and t ( z ) 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 ( θ ) , p ( z ) is understood to refer to different densities.

Further note that the sufficient statistics of p ( θ ) either become parameters of p ( z | θ ) or its normalizers. This can be notated two ways: either by appending 1s to the end of t ( z ) to align with the normalizers, or simply fold them into the log partition function. An example is provided below.

Example: NIW distribution. Anormal distribution p ( z ; µ, Σ) has natural parameters and sufficient statistics given by:

We bend notation here by concatenating vectors to 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:

Meanwhile, the Normal Inverse Wishart (NIW) distribution has sufficient statistics:

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 a simple form, namely:

This follows easily from the fact that p ( θ | z ) ∝ p ( θ ) p ( z | θ ) which are both log-linear in t ( θ ) .

Factor Graphs. The exponential family framework can be applied to complicated distributions over many variables, in which case it can be helpful to rephrase the distribution as a collection of factors on subsets of the variables. We represent the variable z as a graphical model with vertex set V such that the vertices partition the components of z . Let the factor set F be a set of subsets of V such that z f ⊆ z, f ∈ F . Then we write our distribution

These two ways of writing the distribution are trivially equal by defining η as the union of all η f and combining t ( z f ) appropriately. We demonstrate where this notation can be useful in two examples.

C.2 Mean field Optima

Proposition C.4. Given a fixed distribution q ( a ) and a loss function of the form

for some non-negative, continuous function f ( a, b ) with finite integral (such as an un-normalized density), the partially optimal q ( b ) across all possible distributions is given by

Proof. This is a simple extension of a proof in Johnson et al. [29]. Let ˜ p ( b ) be the distribution which satisfies ˜ p ( b ) ∝ exp { E q ( a ) [log f ( a, b )] } . We can drop terms in the loss which are constant with respect to q ( b ) and write:

This is applicable to ELBO maximization problems, as well as the surrogate maximization problem arising in the SV AE:

Example: Exact Inference in the SVAE Applying proposition C.4 to the SVAE's surrogate objective, we see:

Recall that we have chosen ˆ ℓ ϕ ( z | x ) to be a conjugate likelihood function to p ( z | θ ) . Let λ ϕ ( x ) be the outputs of the recognition network such that

Demonstrating that the optimal q ( z ) is in the same exponential family as p ( z | θ ) , with parameters ω = E q ( θ ; η ) [ t ( θ )] + λ ϕ ( x ) . These parameters are the sum of expected parameters from the prior and the parameters of the fake observations.

Example: Linear Dynamical System. Consider the linear dynamical system, where p ( z 1 ) = N ( µ 0 , Σ 0 ) and p ( z t | z t -1 ) = N ( A t z t -1 + b t , Σ t ) (in practice, we use a time-homogenous model in which A,b, Σ do not depend on the time step t ). We can write this distribution

Where t ( z 1 ) , t ( θ 1 ) are given by the NIW example in Sec. C.1, and for t = 2 , . . . , T ,

Note that both t ( θ t ) and t ( θ t +1 ) introduce parameters corresponding to the statistics z t . In the context of the SVAE, we add observations to each time step,

which is conjugate to the likelihood p ( z | θ ) so that the mean 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 z m ⊂ z, m ∈ M . We outline the math in this case by noting that conjugacy allows us to break down statistics on groups of vertices into statistics on individual vertices.

Then

i.e. this expression is linear in the statistics of each argument. This proves important in deriving the optimal factor q ( z m ) given the other factors q ( z -m ) . By applying prop C.4 to the mean field surrogate objective with q ( a ) = q ( θ ) ∏ -m q ( z m ) and q ( b ) = q ( z m ) , we see:

The expected-prior part of this expression can be rewritten (where z f,m = z f ∩ z m and z f, -m = z f ∩ z C m ):

Note that because the joint sufficient statistics of z f, -m , θ f are linear in each input, we can distribute the expectations through the statistic function and write

Where MF ( . ) is a mean field function which takes the statistics of θ f , z f, -m and produces the joint statistics t ( θ f , z f, -m ) . 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 generative model: p ( k 2 ) = Cat ( π 0 ) , p ( k t ) = Cat ( π k t -1 ) , p ( z 1 ) = N ( µ 0 , Σ 0 ) , p ( k 2 ) = Cat ( π 0 ) and p ( z t | z t -1 ) = N ( A k t z t -1 + b k t , Σ k t ) .

In an exponential family form, the model consists of the following factors: (i) an initial-state factor on z 1 matching the NIW example in Eqs. (10-12), (ii) an initial categorical factor on k 2 (see Sec. G for the form), (iii) a transition factor for p ( k t | k t -1 ) which has parameters π ij and sufficient statistics 1 { ( k t = i )&( k t +1 = j ) } for i, j = 1 , . . . , K , and (iv) transition factors corresponding to p ( z t | z t -1 , k t ) , which we will examine further.

Recall from Eq. (49) the expression for the parameters of p ( z t | z t -1 ) in the LDS. Then we can write:

We need to rewrite this to be an exponential family distribution of k t as well (instead of only appearing as an index). Note that the sufficient statistics of a categorical variable is simply its one-hot encoding. If we define t ( θ trans ) to be a matrix whose k th row is given by t ( θ k ) , and log Z ( θ trans ) to be a vector with components log Z ( θ k ) , we write:

Thus the conditional distribution can be described by two factors, one which takes the outer product of the one-hot encoding t ( k t ) 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 in (again, by multiplying a one-hot encoding with a vector of log probabilities).

Given observations, exact inference in the SLDS is intractable. Thus we divide the graphical model via structured mean field inference into q ( z ) q ( k ) . This separates out the joint factors ⟨ t ( θ trans ) , t ( k t ) · t ( z t -1 , z t ) T ⟩ . However, as noted in Eq. (51) we can reframe this as a factor on each mean field cluster:

Thus, applying Eq. (55), the optimal q ( k ) inherits the (expected) initial-state factors from p ( k 2 ) and transition factors from p ( k t | k t -1 ) , as well as the normalizer factors ⟨-log Z ( t ( θ trans )) , t ( k t ) ⟩ from Eq. (58), but replaces the joint factors of Eq. (58) which include z t -1 , z t with

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 ( z 1 ) but all of its transition factors are replaced by

Because t ( k t ) is a one-hot encoding of k t , E q ( k t ) [ t ( k t )] is simply a probability vector of what value k t takes. Thus the new parameter for t ( z t -1 , z t ) is simply a convex combination of the rows of t ( θ trans ) , yielding a time-inhomogeneous 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 E q ( z ) [ t ( z )]

For each t , compute E q ( θ ) [ t ( θ trans )] · E q ( z ) [ t ( z t -1 , z t )] as the new parameters in q ( k ) corresponding to t ( k t )

Run a belief propagation algorithm to get E q ( k ) [ t ( k )]

For each t , compute E q ( θ ) [ t ( θ trans )] T · E q ( k ) [ t ( k t )] as the new parameters in q ( z ) corresponding to t ( z t -1 , z t )

. . . and repeat until convergence.

D Belief propagation in time series

D.1 Sequential LDS Inference

Consider again the generative model

Recall from Eq. (47) that the q ( z ; ω ) 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:

Where h 0 , -1 2 J 0 are the expected statistics of an NIW prior. Further, we place a transition factor on each adjacent pair of states, given by Eq. (49). To simplify notation, we write

For the SLDS, each of these parameters is outputted by the mean field message passing function and is a convex combination of the corresponding parameters for each cluster, weighted by E q ( k t ) [ t ( k t )] . For the LDS, these values are time-homogeneous.

Note that there is also a factor outputted by the recognition network output at each time step:

Because the parameters are expectations, there may be no values of ˆ A t , ˆ Q t , ˆ b t which satisfy

And as a result, it might not be equivalent to any single LDS with fixed parameters. This prevents us from using general Kalman smoother algorithms, so we must return to the definition of exponential families and belief propagation.

The goal of inference is to compute the marginals q ( z t ) , draw samples, and compute log Z . Using the belief propagation algorithm, we know q ( z t ) is proportional to the product of incoming messages from factors which involve z t . In belief propagation, messages pass from factors to variables and back. Thus we write

Where m ω,z ( z ) is the message from the factor with parameter ω to variable z . We similarly define m z,ω ( z ) to be the message from variable to factor. Note that all m .,. ( z t ) are provably Gaussian, and can therefore each message function can be minimally described by the Gaussian natural parameters.

Messages from variables to factors are the product of incoming messages, e.g.

While messages from factors to variables are computed by integrating over incoming messages times the factor potential, e.g.

We attempt to match our notation to the typical Kalman smoother algorithm, by defining filtered messages J t | t , h t | t as the natural parameters of m z t ,ω t +1 ( z t ) , i.e. the distribution which has accumulated the first ω 1: t and λ 1: t , which correspond to the first t 'observations". We similarly define predicted messages J t +1 | f , h t +1 | f as the natural parameters of m ω t +1 ,z t +1 ( z t +1 ) , which corresponds to a distribution over z t +1 given the first t 'observations". Finally we describe the smoothed messages J t | T , h t | T to be the natural parameters of q ( z t ) (note that for simplicity the natural parameters for each distribution are actually -1 2 J . | . , h . | . and the negative one half is assumed to be factored out).

The message passing algorithm is outlined below: Input: ω t , λ t . Initialize :

Filter: for t = 1 , . . . , T - 1 Predict:

Smoother : for t = T -1 , . . . , 1

Sample: (the N below are taking natural parameters as input)

Return: marginal params f t | T , F t | T ; log Z ; and samples x 1: T .

Measurement:

Finalize:

From the log marginals, we can recover the expected statistics:

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:

What is that log partition function term at the end?

Simplifying P t = F t | t + J 11 ,t , we recover normalization term

Linear term

And quadratic term

Thus the entire message is:

To calculate log Z , instead of normalizing each message to be a valid distribution, we can think of normalizing each message so it's of the form exp {⟨ η, t ( z ) ⟩} , i.e. remove all normalizers. Then to calculate the entire model's log Z , we must accrue all 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

The second term is the filtered message computed in the filtering step. Thus we must simply compute the backwards messages. The message becomes:

Now let's evaluate the integral:

And quadratic term

Combining these with the J 11 ,t and h 1 ,t 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, we want to sample from

The rest of the derivation is 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ä & García-Fernández [56] to parallelize our algorithm across time steps.

Let us define the kernel of a (possibly un-normalized) exponential family distribution as the inner product portion of the distribution:

Note that ker ( f ( z )) ∝ f ( z ) . We define two sequences of functions:

Thus g t ( z t -1 ) is defined by natural parameters Γ t , γ t . For t = 1 , we set Γ 1 = 0 , γ 1 = 0 as there is no time-step 0. Thus we claim g 1 ( z 0 ) = 1 is a uniform function.

Further define f t ( z t -1 , z t )

We similarly define f 1 ( z 0 , z 1 ) = f 1 ( z 1 ) , or equivalently Φ 11 , 1 = Φ 12 , 1 = 0 and ϕ 1 , 1 = 0 as there is no z 0 . ϕ 2 ,t = h 0 + r 1 and Φ 22 ,t = J 0 + R 1 to include the factors h 0 , J 0 from the initial state distribution p ( z 1 ) .

We define the associative operation:

Where

And

Proposition D.1. The operation ⊗ is associative.

Proof. Särkkä & García-Fernández [56] define an operator

Where

And and prove that this operator is associative. Further, any rescaling of inputs leads to a simple rescaling of outputs due to the multi-linear nature of these functions. Thus for constants c 1 , c 2 , c 3 , c 4 there exist constants c 5 , c 6 such that

This gives us that

A similar decomposition shows that

Where the second equality comes from the associativity of the × operator.

Finally, we need to show that this operation gives us the right results:

Proposition D.2. Let a i = ( f i , g i ) . Then

is parameterized by the filtered messages f t | t , F t | t .

Proof. We prove by induction that this holds when the operations are performed sequentially, namely

These 'left-justified" operations are simpler, and thus allow us to 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 g 1 ( z 0 ) = 1 is the uniform function (as there is no z 0 ) and thus has Γ 1 = 0 , γ 1 = 0 . By induction, this must be true of all the cumulative kernels g ij ( z 0 ) , as if g i ( a ) in Eq. (120) is independent of a , then g ij ( a ) must be, too. Similarly, we recall that f 1 ( z 0 , z 1 ) = f 1 ( z 1 ) is independent of its first argument and thus has Φ 11 , 1 = Φ 12 , 1 = 0 and ϕ 1 , 1 = 0 . Again, by induction all of the cumulative kernels f 1: t ( z 0 , z t ) must have their first three parameters equal to 0; in Eq. (119) if f i ( a, b ) is independent of a then f ij ( a, c ) will be, too. Thus f 1: t ( z 0 , z t ) = f 1: t ( z t ) .

Next we show that f 1: t ( z t ) = m z t ,ω ( z t ) , the filtered messages from the sequential filter. For t = 1 , f 1:1 ( z 1 ) = m z 1 ,ω ( z 1 ) by construction. We then see:

(

∫

)

(138)

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 to evaluate ( f i , f j ) ⊗ ( g i , g j ) . The integrals in Eqs. (119), (120) can be evaluated as simple functions of the parameters of f i , f j , g i , g j :

Proposition D.3. Let

and where

= ker exp

{⟨

λ

t

+1

, t

(

z

t

+1

)

⟩} ·

z

t

{⟨

exp

t

(

θ

t

+1

)

, t

(

z

t

, z

t

+1

)

⟩} ·

m

z

t

,ω

(

z

t

)

dz

t

Then given the following definitions

We can compute ( f ij , g ij ) = ( f i , g i ) ⊗ ( f j , g j ) as:

and

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

Where we let e T ( z T , z T +1 ) = e T ( z T ) with E 11 ,T = J T | T , ϵ 1 ,T = -h T | T and remaining parameters equal to 0 (as there is no z T +1 . Then the associative operator ⊕ is given by

Which can be computed (derivations omitted) as

Särkkä & García-Fernández [56] prove this operation is associative when operating on full distributions (i.e. without taking the kernel). The proof that we retain associativity when removing multiplicative constants mirrors the proof of Prop. D.1 and is omitted for brevity. Finally, we show that when these operations are done sequentially,

The base e T ( z t ) = q ( z T ) is true by construction. Then, plugging the inductive hypothesis and Eq. (148) into the definition of ⊕ yields Eq. (103), proving that e t ( z t ) = q ( z t ) implies e t -1 ( z t -1 ) = q ( z t -1 ) .

D.3 HMMInference

The discrete Markov model is time homogenous and requires only expected probabilities for the initial state k 2 and expected probabilities for the transition matrix π ij . The 'observations" of the HMMare (in the case of the SLDS) messages passed from the LDS chain, which we denote ω t ( k ) . The setup is fundamentally the same as the LDS: E

Input: q ( θ ) [ t ( θ )] , η recog .

Forward Pass:

Backward Pass:

Return: log marginals α t ( k ) + log β t ( k ) ; log Z

E Evaluating the SVAE loss

Johnson et al. [29] decompose the loss into three terms

Because p ( θ ) , p ( z | θ ) , q ( θ ; η ) , q ϕ ( z | x ; η ) are all exponential family distributions, the expectations here can be evaluated in closed form, with the exception of the reconstruction loss E q ( z ) [ p γ ( x | z )] . The reconstruction loss can be estimated without bias using Monte Carlo samples and the reparameterization trick [32]. For this to work, we must be able to reparameterize the sample of q ( z ) such that we can back propagate through the sampling routine, which is not possible for discrete random variables. As a result, the latent variables which need to be sampled to reconstruct x (i.e. all variables which are 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 a regularizer on the graphical model parameters and can be easily evaluated via Eq. (25). The rest of 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 q ( z ) has parameters ω given by Eq. (47), we can again invoke Eq. (25) to get

Namely, because log q ( z ) is defined as the expected log prior p ( z | θ ) 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 q ( z ) . In particular, the encoder takes the observations x and outputs normal likelihood functions on each z independently. Thus the cross-temporal covariance statistics z t z T t +1 have no corresponding recognition potentials. Further, we restrict the potentials on z t z T t to be diagonal, meaning our observational uncertainty factorizes across dimensions. Thus the only statistics in t ( z ) for which λ ϕ ( x ) are non-zero are z t and z 2 ti for each component i of t ( z ) . As a result, the inner product ⟨ λ ϕ ( x ) , E q ( z ) t ( z ) ⟩ sums over 2 TD terms for a sequence with T timesteps and D -dimensional latent variable at each step.

Local KL with structured mean field When there is mean field separation, the prior and variational factor differ in more than just the recognition potentials, as q ( z ) replaces some factors of p ( z | θ ) with separable factors. We can rewrite

The natural parameters of q ( z m ) differ from the natural parameters of p ( z m | z <m , θ ) not only in the recognition potentials, but also for any factors which include a variable in 'later" clusters z >m .

Example: SLDS The SLDS SVAE has 2 mean field clusters, q ( k ) q ( z ) for the discrete variables k and continuous variables z . If we write our prior distribution p ( θ ) p ( k | θ ) p ( z | k, θ ) then KL ( q ( z ) ∥ p ( z | k, θ )) is computed exactly the same as in the LDS example above, as these two distributions differ only in the recognition potentials. The new challenge is computing KL ( q ( k ) ∥ p ( k | θ )) which contains no recognition potentials but differs based on mean field messages from q ( z ) which are present in q ( k ) but not in p ( k | θ ) .

In particular, the prior p ( k | θ ) is a time series with no observations, while q ( k ) includes parameters at every time step which encode the log probability of being in that state at that time, given by Eq. (60): E q ( θ ) [ t ( θ trans )] · E q ( z ) [ t ( z t -1 , z t )] . The corresponding sufficient statistics are the marginal probabilities of being in state k at time t . Note that we never need to compute the sufficient statistics corresponding to HMM transitions because we don'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

This is trivial to evaluate, as 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 recognitionpotential-times-expected-statistics 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 η can be easily computed. Defining, as in the main paper, µ η = E q ( θ ; η ) [ t ( θ )] , we can use the chain rule to write:

Where ∂ L ∂η detach ( µ η ) refers to the gradient of the loss with respect to η , disconnecting the computational graph at µ η (i.e. the dependence of the loss on η which is not accounted for by its dependence on µ η ).

Proposition F.1.

Proof. q ( z ) depends on η only through the expected statistics µ η which define the parameters of the graphical model factors. Further, the local KL, outlined in Eq. (163), only depends on q ( θ ) through its expected statistics (noting that the log normalizers of p ( z | θ ) are also staistics of q ( θ ) ). Thus the only place in the loss that η appears on its own is the prior KL. Let η 0 be the natural parameters of the prior p ( θ ) . Then the prior KL, as per Eq. (25), is given by:

Then we see:

Where we use Prop. C.1 to cancel terms. Thus we conclude that the entire gradient of the loss with respect to η passes through µ η

The main consequence of this theorem is that in all cases the gradient can be written

And thus the natural gradient, which multiples the gradient by the inverse of ∂µ η ∂η , can be easily computed withoout matrix arithmetic by skipping the portion of gradient computation corresponding to the η → µ η mapping. Here is some sample Jax code which takes a function f and returns a function which computes f in the forward pass but in the backward pass treats the gradient as identity:

def straight_through(f):

def straight_through_f(x):

zero

=

x - jax.lax.stop_gradient(x)

return zero +

jax.lax.stop_gradient(f(x))

return straight_through_f

Fisher information of transformed variables . Given a distribution q ( θ ; η ) with natural parameters η , the Fisher information matrix of η is given by:

Note that for an exponential family distribution

which clearly has expectation 0 under q ( θ ) . Thus

Exponential family theory gives us further equalities:

Thus deriving our desired property that F η = ∂µ η ∂η .

Now we consider a more general case, where η = f (˜ η ) for some unconstrained, non-natural parameters ˜ η and continuous invertible function f . Then we see:

We can use this to compute the natural gradient with respect to the unconstrained parameters:

Implementing natural gradients via automatic differentiation We argued in the main paper that to implement natural gradients we replace the reverse-mode backpropagation through the ˜ η → η map with a forward-mode differentiation through the inverse operation. Assuming we have a forward map f and inverse function f_inv , sample code is shown below:

G Matrix-normal exponential family distributions

Recall that exponential family distributions are given by

The natural parameters often have coupled constraints, and so can be expressed as a simple function of canonical parameters which carry easily-interpretable information about the distribution (e.g. the natural parameters of a normal distribution can be expressed in terms of the mean and precision matrix of that distribution). Thus to define each distribution, we provide (a) the natural parameters η as a function of canonical parameters, (b) the sufficient statistics t ( x ) , (c) the expected sufficient statistics E p ( x ; η ) [ t ( x )] , and (d) the log partition function log Z ( η ) .

Note that sometimes we have symmetric positive definite (SPD) matrix-valued parameters. In every case, the corresponding statistics for those parameters are also symmetric. For example, given a n × n SPD matrix, a distribution might have parameters and statistics:

We note that ⟨ η, t ( x ) ⟩ = ∑ n i =1 ∑ n j =1 S ij x i x j or equivalently the sum of elements of S ⊗ xx T (for column vector x containing each x i ). In the following sections, we will abuse notation for compactness by simply writing:

Constraints. To optimize these parameters via gradient descent, we must define constrained parameters as a function of unconstrained parameters as discussed in Sec.[natgrads]. Further, to implement natural gradients we must implement the inverse operation. The table below lists all such transformations.

Where Softmax is made invertible by appending a 0 to the end of ˜ η . For the SPD matrix parameters, we define S = diag ( σ ) · R · diag ( σ ) where σ ∈ R n + is a vector of standard deviations and R is a correlation matrix, transformed to and from unconstrained space by the tensorflow Correlation Cholesky bijector.

G.1 Multivariate Normal distribution

If x ∼ N ( µ, Σ) for (i) x, µ ∈ R n , and (ii) Σ ∈ S n ++ , then p ( x ; µ, Σ) is given by:

The log partition function is given by:

Here, π ≈ 3 . 14 is a constant.

G.2 Normal inverse Wishart (NIW) distribution

If µ, Σ ∼ NIW ( S, m, λ, ν ) for (i) µ, m ∈ R n , (ii) Σ , S ∈ S n ++ , (iii) λ ∈ R + , and (iv) ν ∈ R >n -1 , then p ( µ, Σ; S, m, λ, ν ) is given by:

Where ψ is the digamma function. The log partition function is given by:

With the multivariate gamma function Γ n .

G.3 Matrix normal inverse Wishart (MNIW) distribution

If X, Σ ∼ MNIW ( S, M, V, ν ) for (i) X,M ∈ R n × m , (ii) Σ , S ∈ S n ++ , (iii) V ∈ S m ++ , and (iv) ν ∈ R >n -1 , then p ( X, Σ; S, M, V, ν ) is given by:

Where ψ is the digamma function. The log partition function is given by:

With the multivariate gamma function Γ n . For the SLDS, we use this MNIW distribution by defining X ∈ R n × n +1 = [ A | b ] where z t +1 ∼ N ( Az t + b, Σ) so the MNIW distribution provides a prior on the transition matrix, offset, and conditional noise.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\n\nWEAKNESSES:\n1. 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)\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nTechnical limitations are not clearly discussed.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nSVAEs 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.\n\nWEAKNESSES:\nThis 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\nQUESTIONS:\nPlease see the \"Weaknesses\" section.\n\nLIMITATIONS:\nPlease see the \"Weaknesses\" section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n### 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.\n\nWEAKNESSES:\n### 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.\n\nQUESTIONS:\nIf 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)?\n\nLIMITATIONS:\nLimitations are not discussed in the paper.\nSome of my questions above could be discussed in a dedicated limitation section/paragraph.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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).\n\nSTRENGTHS:\nThis 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).\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\nSee 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?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-10"}