[Paper Review] Collapsed Amortized Variational Inference for Switching Nonlinear Dynamical Systems
This paper proposes collapsed amortized variational inference for switching nonlinear dynamical systems (SNLDS), enabling efficient end-to-end training via reparameterization by marginalizing discrete latent states and using an amortized inference network for continuous states. The method achieves superior segmentation of complex time series—such as videos and 3D human poses—compared to prior approaches, especially in modeling nonlinear dynamics with fewer modes.
We propose an efficient inference method for switching nonlinear dynamical systems. The key idea is to learn an inference network which can be used as a proposal distribution for the continuous latent variables, while performing exact marginalization of the discrete latent variables. This allows us to use the reparameterization trick, and apply end-to-end training with stochastic gradient descent. We show that the proposed method can successfully segment time series data, including videos and 3D human pose, into meaningful ``regimes'' by using the piece-wise nonlinear dynamics.
Motivation & Objective
- To develop an efficient inference method for switching nonlinear dynamical systems (SNLDS) with both discrete and continuous latent variables.
- To enable end-to-end training using stochastic gradient descent by applying the reparameterization trick to continuous variables after marginalizing discrete states.
- To address posterior collapse in SNLDS models by introducing a novel posterior regularization scheme that prevents overuse of a single mode.
- To improve segmentation accuracy on nonlinear time series, including video and 3D human pose sequences, compared to existing models like SLDS.
- To demonstrate the method’s effectiveness on real-world datasets such as CMU MoCap and synthetic motion sequences with piecewise nonlinear dynamics.
Proposed method
- The method uses a collapsed variational inference framework that exactly marginalizes discrete latent states $s_{1:T}$ using the forward-backward algorithm, enabling efficient gradient computation.
- An amortized inference network $q_{\boldsymbol{\phi}}(\mathbf{z}_{1:T}|\mathbf{v}_{1:T})$ is trained to approximate the posterior over continuous states $\mathbf{z}_{1:T}$, allowing reparameterization and differentiable training.
- The inference network employs bidirectional and unidirectional RNNs to encode observations $\mathbf{x}_{1:T}$ and propagate hidden states, respectively, with shared parameters across time.
- A Rao-Blackwellized particle filter is used as a proposal, though the method primarily relies on a single posterior sample, similar to VAEs.
- Posterior regularization is introduced via annealed entropy regularization and temperature scheduling to prevent mode collapse and stabilize training.
- The model is trained jointly using stochastic gradient descent on the evidence lower bound (ELBO), with early training using high entropy and temperature to encourage exploration of all modes.
Experimental results
Research questions
- RQ1Can collapsed amortized variational inference enable efficient and scalable training of SNLDS with nonlinear dynamics?
- RQ2How does marginalizing discrete states improve training stability and enable the use of reparameterization in hybrid systems?
- RQ3To what extent does posterior regularization prevent mode collapse in SNLDS models with flexible nonlinear dynamics?
- RQ4How does the proposed method compare to SLDS and other baselines in segmenting nonlinear time series with minimal mode usage?
- RQ5What is the impact of annealing schedules on model convergence and segmentation quality in unsupervised settings?
Key findings
- The SNLDS model significantly outperforms SLDS on nonlinear 2D motion tasks, achieving higher F1 scores with fewer states, especially when using a 5-frame tolerance for change detection.
- On the Dubins car path dataset, the SNLDS model achieved a relative negative log-likelihood improvement of over 10% compared to baselines when using optimal annealing schedules.
- In the Salsa dancing task using CMU MoCap data, the model successfully segmented 3D human motion into three interpretable modes: forward/backward motion, clockwise turning, and counter-clockwise turning.
- The F1 score for the Salsa dancing sequences improved by over 20% when using delayed annealing (starting at 50k steps), indicating that later regularization helps in learning more distinct modes.
- The method achieved better segmentation with fewer modes than SLDS, demonstrating that nonlinear models can capture complex dynamics more efficiently than piecewise linear approximations.
- The annealing schedule that delayed both entropy coefficient and temperature decay until 50k and 100k steps, respectively, led to the best final performance in terms of both likelihood and F1 score.
Better researchstarts right now
From reading papers to final review, dramatically reduce your research time.
No credit card · Free plan available
This review was created by AI and reviewed by human editors.