Anchored Variational Inference for Personalized Sequential Latent-State Models

Abstract: Sequential latent-variable models with subject-specific random effects provide a flexible framework for modeling temporally structured data with both local latent dynamics and stable between-subject heterogeneity. In such models, conditional inference for the local latent process is often tractable, but integrating over subject-specific random effects can be computationally demanding. We propose an anchored variational inference framework for efficient approximate inference in this setting. The central idea is to replace the full conditional posterior of the local latent process with its evaluation at a representative value of the subject-specific latent effect, called the anchor point, thereby preserving tractable local inference while substantially reducing computational cost. This approximation is especially appealing in sequential settings, where the posterior distribution of the random effect becomes increasingly concentrated as the sequence length grows. Under suitable conditions, we show that the posterior mean is a nearly optimal anchor point and that the resulting anchored variational EM (AVEM) algorithm approximately preserves the local monotonicity behavior of standard variational inference. We instantiate the framework in two representative classes of sequential latent-variable models, namely mixed hidden Markov models and mixed-effects state-space models, derive the corresponding AVEM algorithms, and use simulation studies to indicate that the resulting methods achieve accurate estimation with substantial computational gains. We also discuss a partially anchored variant of the framework, in which only the components of the subject-specific latent effect whose posteriors are well concentrated are anchored.

• The paper introduces an anchored variational framework that replaces the local latent process with its evaluation at the posterior mean, preserving tractable computations.
• It demonstrates substantial computational efficiency over MCEM and QEM by reducing integration complexity in both MHMMs and MESSMs.
• Theoretical guarantees and empirical evaluations show that AVEM converges with minimal error as sequence length increases, highlighting its scalability.

Anchored Variational Inference for Personalized Sequential Latent-State Models

Introduction and Motivation

The paper "Anchored Variational Inference for Personalized Sequential Latent-State Models" (2604.23454) introduces an anchored variational framework for scalable inference in sequential models that integrate local latent dynamics and subject-specific heterogeneity. Classical HMMs and SSMs offer tractable inference for latent states or paths, but extensions accommodating random effects across subjects lead to intractable posteriors requiring expensive numerical integration or iterative approximation. Standard mean-field variational approaches break intrinsic dependence among latent variables, degrading inference quality, while structured variational forms remain computationally burdensome due to repeated evaluation of the latent-path conditional on varying random effects.

Anchored Variational Inference: Formulation

The central contribution of the paper is the construction of an anchored variational family: for each subject, the conditional distribution of the local latent process is replaced by its evaluation at a representative anchor point of the random effect, usually the current posterior mean under the variational approximation. This preserves tractable local computations (e.g., forward-backward for HMM) and approximates only the marginal of the random effect.

The structured variational factorization

$q(U, f) = q(U \mid f) q(f)$

preserves dependence but is computationally expensive. The anchor family simplifies this to

$q(U, f) = p(U \mid f_0, D) q(f)$

where $f_0$ is the anchor, typically the mean of $q(f)$.

Lemma 1 formalizes that the optimal conditional $q(U \mid f)$ under fixed $q(f)$ is $p(U \mid f, D)$, and Lemma 2 characterizes the optimal anchor point as the minimizer of $$\mathbb{E}_{q(f)}[\mathrm{KL}(p(U \mid f_0, D) \Vert p(U \mid f, D))]$$. In practice, the mean anchor is nearly optimal for concentrated posteriors.

Figure 1: Schematic illustration of AVEM versus numerically implemented exact EM for a two-dimensional random effect $f_i \in \mathbb{R}^2$, showing AVEM's adaptation to posterior location and scale compared to fixed-grid quadrature.

Theoretical Guarantees

The anchored approximation is theoretically justified via KL-smoothness assumptions: as trajectory length $T$ increases, the posterior for the random effect $q(U, f) = p(U \mid f_0, D) q(f)$0 concentrates, implying that the normalized KL gap between mean-anchor and optimal-anchored structured variational objectives vanishes as $q(U, f) = p(U \mid f_0, D) q(f)$1. The monotonicity of the anchored ELBO is approximately preserved, with per-iteration decrement bounded by $q(U, f) = p(U \mid f_0, D) q(f)$2 for $q(U, f) = p(U \mid f_0, D) q(f)$3, ensuring convergence properties akin to standard variational EM.

AVEM Algorithms for Structured Sequential Models

The AVEM procedure is instantiated for two canonical models: mixed hidden Markov models (MHMMs) and mixed-effects state-space models (MESSMs).

Mixed Hidden Markov Models

For MHMMs, the latent states are discrete, and random effects follow Gaussian distributions. AVEM employs Gaussian variational approximations for $q(U, f) = p(U \mid f_0, D) q(f)$4, requiring only one forward-backward HMM pass per anchor, whereas exact EM necessitates quadrature or Monte Carlo integrations over the random effect.

Figure 2: Estimation accuracy and computational performance of AVEM under Gaussian MHMM for varying sample sizes, sequence lengths, and random-effect variances.

AVEM updates variational parameters through closed-form or numerical optimization and updates model parameters via maximization of the anchored ELBO. This leads to significant computational efficiency, with theoretical guarantees on accuracy for concentrated posteriors.

Mixed-Effects State-Space Models

For MESSMs, AVEM utilizes Kalman filter/smoother routines for linear-Gaussian latent processes, anchored at the posterior mean of the random effect components. Variational factors for transition and loading matrices are updated via explicit closed-form expressions, and only one Kalman smoother pass is needed per subject per iteration.

Figure 3: Boxplots of RMSEs and normalized ELBO trajectories for MESSM estimation, showing AVEM's stability and approximate ELBO monotonicity.

Empirical Evaluation

Extensive simulation studies for Gaussian MHMMs demonstrate that AVEM maintains estimation accuracy with substantial computational gains as both $q(U, f) = p(U \mid f_0, D) q(f)$5 and $q(U, f) = p(U \mid f_0, D) q(f)$6 increase. AVEM is superior to Monte Carlo EM (MCEM) and quadrature-based EM (QEM) in runtime, with error rates comparable or lower, especially as the posterior for $q(U, f) = p(U \mid f_0, D) q(f)$7 concentrates with longer trajectories.

Figure 4: Estimation accuracy and runtime of AVEM as function of latent state number $q(U, f) = p(U \mid f_0, D) q(f)$8 and random-effect dimension $q(U, f) = p(U \mid f_0, D) q(f)$9.

AVEM's computational cost is primarily sensitive to the number of latent states, and only mildly sensitive to random-effect dimension, highlighting its scalability.

For non-Gaussian emission models (e.g., Bernoulli MHMM), AVEM remains faster and broadly competitive in estimation accuracy, although QEM may offer marginal gains for certain variance parameters.

Figure 5: Performance comparison of AVEM, MCEM, and QEM for Gaussian MHMMs across varying trajectory lengths and random-effect variances.

Figure 6: Performance comparison of AVEM, MCEM, and QEM for Bernoulli MHMMs under different configurations.

AVEM is shown to maintain monotonic ELBO convergence and robustness in the state-space setting, with population and subject-specific parameter RMSEs decreasing as $f_0$0 and $f_0$1 increase.

Partial Anchoring and Model Flexibility

The framework extends to partial anchoring: components of the random effect with well-concentrated posteriors are anchored, while diffuse components are integrated numerically, providing accuracy gains when only partial concentration is feasible.

Figure 7: Comparison of AVEM and PAVEM in localized-random-effect setting, showing estimation error for components and computational gains.

Practical and Theoretical Implications

Anchored variational inference achieves a balance between computational efficiency and posterior dependence preservation in sequential models with random effects. The theoretical guarantees and empirical results indicate AVEM's suitability for settings where each subject provides substantial time series data: posterior concentration of random effects is the norm, and anchored approximation is accurate.

In practice, AVEM enables scalable inference in complex latent-state models, including those with non-Gaussian emissions, while maintaining tractability and accuracy. The methodology is directly applicable to neuroscience, behavioral studies, and panel data analysis, and its computational design is compatible with large-scale applications.

From a theoretical standpoint, AVEM provides a controlled approximation for variational EM, with monotonicity and optimality errors quantifiable in terms of posterior concentration. This framework may extend to broader classes of graphical models with tractable conditional local inference and low-dimensional heterogeneity.

Conclusion

Anchored variational inference offers a tractable, theoretically justified, and empirically validated solution for inference in personalized sequential latent-state models, preserving essential posterior dependencies and achieving substantial computational gains. The AVEM framework generalizes to various structured models and emission types, and its theoretical foundation supports its robustness and accuracy. Future work may focus on diagnostics for anchoring adequacy, further extensions to hierarchical and multi-site models, and characterizing the local optima equivalence with structured variational objectives.