Amortized Variational EM

• Amortized Variational EM is a scalable framework that replaces per-data point variational optimization with a shared inference network.
• It integrates concepts from variational autoencoders, iterative inference, and deep dynamical models to efficiently handle large and complex datasets.
• By amortizing the E-step, the method achieves faster inference and improved convergence, benefiting applications in generative modeling and missing data imputation.

Amortized Variational EM (Expectation-Maximization) is a class of algorithms that apply variational inference within the framework of Expectation-Maximization but replace per-datum, per-step variational E-steps with global inference networks ("amortization"). These methods scale EM-style latent variable modeling to large or complex datasets and models, enable efficient inference, and remove the need for local optimization at each data instance and iteration. This paradigm unifies and extends variational autoencoders, iterative inference, recursive mixture estimation, and deep dynamical latent models, providing both practical computational benefits and novel avenues for modeling expressivity.

1. Conceptual Foundations and Motivation

The standard EM algorithm for latent variable models iterates between computing exact or variational posteriors for each data point (E-step) and maximizing expected complete-data likelihood (M-step). For deep or structured latent variable models—where posteriors are intractable and nonlinear—approximating the E-step via variational inference is standard. However, classical approaches require either computationally intensive per-datum optimization or are limited by the simplicity of the variational family (such as unimodal Gaussian encoders). Amortized Variational EM (AVE) replaces the per-datum local E-step optimization with a parameterized inference network, typically realized as a deep neural network. This enables scalable training and fast test-time inference—critical in modern settings such as deep generative modeling and time series (Marino et al., 2018, Kim et al., 2020, Cherifi et al., 22 Mar 2026).

The core structural change in AVE is that variational parameters for each data instance are output by a shared inference function (e.g., $q_\phi(z|x)$), so the cost of optimizing variational posteriors is "amortized" across all data. This approach provides computational tractability, enables applications to large datasets, and forms the basis of state-of-the-art generative modeling.

2. Mathematical Structure: AVE as Variational EM

The underlying mathematical formulation is based on the joint maximization of a variational lower bound (ELBO) on the observed data log-likelihood: $$L(\theta, q) = E_{q(z|x)}[\log p_\theta(x, z) - \log q(z|x)]$$ Classical variational EM alternates optimizing $q$ for each $x$ (E-step) and optimizing $\theta$ (M-step). Amortized approaches restrict $q$ to a global parametric family: $q_\phi(z|x)$, with parameters $\phi$ fitted across all data simultaneously.

In time-series latent variable models, the amortized E-step can be defined on a per-time-step basis using the filtering variational objective: $$\mathcal{F} = \sum_{t=1}^T \mathbb{E}_{q(z_{<t}|x_{<t})}[\mathcal{F}_t]$$ where each step's free energy

$$\mathcal{F}_t = -\mathbb{E}_{q(z_t|x_{\leq t}, z_{<t})}[\log p_\theta(x_t, z_t | x_{<t}, z_{<t}) - \log q(z_t | x_{\leq t}, z_{<t})]$$

is minimized by updating only the current variational factor $q(z_t | x_{\leq t}, z_{<t})$, holding the rest fixed (Marino et al., 2018).

In recursive mixture estimation for VAEs, the E-step adds or refines a mixture component $q_t(z|x)$ to better approximate the functional gradient of the ELBO, then recomputes the mixture weights for the updated variational posterior (Kim et al., 2020).

REM (Dieng et al., 2019) and IWAE (Dieng et al., 2019) can be interpreted as stochastic EM algorithms where the M-step is performed via (importance-weighted) expectation maximization using a proposal distribution parameterized by an amortized inference network.

3. Amortization Mechanisms: Inference Networks

Amortization can be realized in various ways:

• Direct parameter prediction: A global neural encoder predicts the parameters of $q_\phi(z|x)$ for each $x$.
• Iterative inference networks: As in Amortized Variational Filtering (AVF), inference is performed by a learned function $f_\phi$ that takes as input the current posterior parameters and gradient information and outputs refined posterior parameters in a small number of steps (Marino et al., 2018).
• Recursive mixtures: Amortized EM can construct a sequence of encoder networks $q_m(z|x)$ whose convex combination forms a more expressive variational family. Each new component is optimized to improve the ELBO and increase mode-coverage, and at test time the evaluation is still amortized—requiring only a fixed number of network evaluations (Kim et al., 2020).
• Wake–sleep algorithms: These alternate a "wake" M-step using samples from the approximate posterior and a "sleep" E-step that fits the recognition network to posteriors sampled from the generative model (Wenliang et al., 2020).

In all cases, amortization removes per-datum optimization, enabling constant inference-time per instance, in contrast to semi-amortized or non-amortized methods that require multiple inference optimization steps.

4. Filtering, Temporal, and Structured Inference

In sequential latent variable models, the AVE framework is adapted to the filtering setting—called Amortized Variational Filtering EM (AVF). The generative process is factorized autoregressively, and the variational posterior is strictly causal: $$q(z_{1:T}|x_{1:T}) = \prod_{t=1}^T q(z_t | x_{\leq t}, z_{<t})$$ At each time $t$, the E-step minimizes the step-wise free energy with respect to the parameters of $q(z_t | x_{\leq t}, z_{<t})$ while holding the past fixed, and then the M-step updates the generative parameters using the accumulated gradients. The inference network $f_\phi$ is trained to efficiently output approximately optimal free-energy-minimizing posterior parameters. This methodology supports incremental, online, and real-time inference (Marino et al., 2018).

These ideas generalize to models with missing data, where an amortized inference network predicts the distribution over missing covariates given observed variables, as in AV-LR for logistic regression with missing data (Cherifi et al., 22 Mar 2026). The ELBO or its importance-weighted variant (IWELBO) is then optimized jointly over model and inference parameters.

5. Algorithmic Realizations and Empirical Properties

Multiple algorithmic patterns are instantiated within the AVE paradigm:

• AVF (Amortized Variational Filtering): Uses a learned iterative inference network to minimize each filtering E-step's free energy. Pseudocode and practical implementation directly couple time-step-wise E-step minimization with joint parameter updates (Marino et al., 2018).
• Recursive mixture estimation (RME): Adds new mixture components to the inference network, optimizing both direct data ELBO and a divergence-expansion penalty, updating mixture weights via a small neural network, and finally maximizing the decoder ELBO with the new mixture (Kim et al., 2020).
• REM: Alternates between updating an inclusive-KL-fitted recognition network proposal and reweighting model parameters using IS-EM (importance sampling EM), yielding decoupled and stable joint training (Dieng et al., 2019).
• Wake–sleep: Alternates expectation maximization steps using data and model-sampled pairs (Wenliang et al., 2020).

Empirically, these approaches achieve or surpass the per-time-step ELBO and test likelihoods of standard, non-amortized EM or hand-tuned inference filters. For example, AVF improves on VRNN, SRNN, and SVG baselines across speech, music, and video domains, with observed reductions in negative free-energy and more stable convergence (Marino et al., 2018). Recursive mixture VAEs (RME) consistently realize higher IWAE log-likelihoods and more flexible multimodal posterior approximations compared to single-component amortized, semi-amortized, or flow-based approaches, at a fraction of inference-time computational cost (Kim et al., 2020). REM achieves higher test log-likelihoods and prevents posterior collapse, maintaining richer variational approximations (Dieng et al., 2019).

6. Comparison to Classical and Analytical EM

Analytical EM, as feasible for specific deep generative networks with piecewise affine (CPA) structure, enables exact posteriors and closed-form E- and M-steps (Balestriero et al., 2020). However, its combinatorial integral costs preclude scaling to high-dimensional settings. Amortized VI in the VAE (and thus AVE) introduces a variational gap, amortization error, and approximation of the true, often multimodal, posterior by parametric encoder outputs. Insights from analytical EM suggest that richer, region-aware amortized variational families could better match true posteriors.

The AVE framework connects to the classical EM theory: IWAE is identified as a stochastic EM M-step using importance sampling (Dieng et al., 2019); recursive amortized methods can be viewed as functional EM procedures where the E-step is performed in function space via mixture growth or iterative refinement (Kim et al., 2020). Wake–sleep algorithms realize amortized E- and M-steps by neural function approximation, shifting away from per-datum optimization while preserving EM-like alternating objectives (Wenliang et al., 2020).

7. Applications, Limitations, and Empirical Considerations

Amortized Variational EM methods are widely adopted in:

• Deep generative sequence models (e.g., VRNN, SRNN, SVG) for real-time inference on speech, music, video (Marino et al., 2018).
• High-dimensional density estimation and representation learning, outperforming standard VAEs in log-likelihood, expressivity, and computational cost (Kim et al., 2020, Dieng et al., 2019).
• Missing data imputation and regression (e.g., AV-LR) with computational advantages (30–60 seconds training vs. thousands in classical EM, real-time inference) and accuracy comparable to or exceeding state-of-the-art algorithms, including under MNAR mechanisms (Cherifi et al., 22 Mar 2026).
• Generative models with discrete or non-Euclidean latents, simulation-based likelihoods, and settings where per-instance variational optimization is prohibitive (Wenliang et al., 2020).

Limitations include the persistence of the amortization gap—the error induced by approximating per-instance posterior optima with a global inference network—and the bounded flexibility of parametric or mixture inference networks. Empirical studies indicate that deeper/iterative inference networks, recursive mixtures, and the use of importance weighting or richer variational families diminish this gap, but do not eliminate it entirely. Analytical EM, where tractable, outperforms all amortized variants in likelihood, but is infeasible for high-dimensional or standard deep architectures (Balestriero et al., 2020).

A plausible implication is that future advances will focus on hybrid approaches: semi-amortized methods, iterative or recursive inference, and targeted expansion of the variational family, informed by structural properties of the model and data. These advances aim to further close the amortization gap, improve uncertainty quantification, and exploit the scalability of AVE methodology.