Latent Stochastic Interpolants

1. Problem Statement

Stochastic Interpolants (SI) are a powerful generative modeling framework capable of transforming arbitrary probability distributions (bridging a prior $p_0$ and a target $p_1$) using continuous-time stochastic processes. However, standard SI has two critical limitations that prevent its application in modern, efficient generative modeling:

1. Requirement for Observed Samples: SI requires direct access to samples from both the prior ($x_0$) and the target ($x_1$) distributions to construct the interpolant. This makes it unsuitable for jointly optimized latent variable models, where the target distribution (the aggregated posterior) is not fixed but is learned alongside an encoder and decoder.
2. Computational Inefficiency: Applying SI directly in high-dimensional observation spaces (e.g., raw pixels) is computationally expensive. While latent variable models (like VAEs) reduce dimensionality, existing methods often decouple the learning of the latent representation from the generative process, leading to suboptimal alignment.

The core challenge is to enable joint end-to-end training of an encoder, a decoder, and a generative model entirely within a latent space, while retaining the flexibility of SI to handle arbitrary priors and avoiding the need for simulation-heavy training.

2. Methodology: Latent Stochastic Interpolants (LSI)

The authors propose Latent Stochastic Interpolants (LSI), a framework that generalizes SI to operate within a learned latent space.

A. Theoretical Foundation: Continuous-Time ELBO

The core innovation is the derivation of a principled Evidence Lower Bound (ELBO) directly in continuous time.

• Model Dynamics: The generative process is modeled as a Stochastic Differential Equation (SDE) in the latent space: $dz_t = h_\theta(z_t, t)dt + \sigma(z_t, t)dw_t$.
• Variational Approximation: To train the model, a variational posterior $Q$ is constructed. Crucially, the authors construct a Diffusion Bridge between the prior $p_0(z_0)$ and the aggregated posterior $p_\theta(z_1) = \int p_\theta(z_1|x)p(x)dx$.
• Simulation-Free Sampling: A major hurdle in latent SDE training is the computational cost of simulating the SDE for the variational posterior at every step. LSI overcomes this by assuming linear SDE dynamics (linear drift and additive noise). This allows the conditional density $p(z_t | z_0, z_1)$ to be derived analytically as a Gaussian. Consequently, latent samples $z_t$ can be generated directly via reparameterization without numerical integration (simulation-free), similar to standard diffusion models.

B. The Training Objective

The derived ELBO consists of two main terms:

1. Reconstruction Term: $-\ln p_\theta(x_1|z_1)$, ensuring the decoder can reconstruct the input from the latent representation.
2. Dynamics Mismatch Term: A penalty on the difference between the variational drift and the model drift, derived from the path-space KL divergence.

The final loss function (using the InterpFlow parameterization) is:
$$\mathcal{L} = \mathbb{E} \left[ -\ln p_\theta(x_1|z_1) + \frac{\beta_t}{2} \left\| \sigma\sqrt{\frac{t}{1-t}}\epsilon + \sqrt{1-t}(z_1 - z_0) + \sqrt{t}z_t - \hat{h}_\theta(z_t, t) \right\|^2 \right]$$
Where:

• $z_0 \sim p_0$ (Prior), $z_1 \sim p_\theta(z_1|x)$ (Encoded latent).
• $z_t$ is the interpolant constructed as $z_t = \eta_t \epsilon + \kappa_t z_1 + \nu_t z_0$.
• $\beta_t$ is a weighting term balancing reconstruction and generative fidelity.

C. Key Architectural Components

• Joint Optimization: The encoder, decoder, and latent SI model ($h_\theta$) are trained simultaneously. This allows the latent space to evolve specifically to support the generative process, unlike standard VAEs where the latent space is fixed before generative training.
• Flexible Priors: Unlike standard diffusion models that rely on simple Gaussian priors, LSI supports arbitrary prior distributions ($p_0$) in the latent space, leveraging the SI framework's ability to bridge any two distributions.
• Sampling: The framework supports flexible sampling strategies, including deterministic sampling (Probability Flow ODE) and stochastic sampling (SDE), as well as Classifier-Free Guidance (CFG).

3. Key Contributions

1. Latent Stochastic Interpolants (LSI): A novel framework enabling the joint end-to-end training of encoders, decoders, and continuous-time latent generative models.
2. Principled ELBO: Derivation of a continuous-time ELBO that provides likelihood control and enables simulation-free training by constructing a variational posterior via a diffusion bridge.
3. Unifying Perspective: Demonstrates that SI can be integrated into latent variable models, bridging the gap between flow-based methods, diffusion models, and VAEs.
4. Efficiency & Flexibility: The method avoids the computational burden of high-dimensional SI while retaining the ability to use diverse priors and achieve high-fidelity generation.

4. Experimental Results

The authors evaluated LSI on the ImageNet (2012) benchmark at resolutions of 64x64, 128x128, and 256x256.

• Performance: LSI achieves competitive FID scores comparable to observation-space SI models of similar parameter counts. For example, at 128x128, LSI achieved an FID of 3.12 (vs. 3.46 for observation space), demonstrating that joint learning in latent space does not sacrifice quality.
• Computational Efficiency:
  • FLOP Savings: While total parameters are similar, LSI is significantly cheaper during sampling. The encoder is unused, and the decoder is used only once. The latent model (which is smaller than the full observation model) is run repeatedly.
  • Results: Sampling 128x128 images with 100 steps resulted in a 29.7% reduction in FLOPs, and 256x256 images saw a 64.9% reduction compared to observation-space baselines.
• Joint Learning Benefits:
  • Varying the loss weight $\beta$ showed that jointly optimizing the encoder for the generative task (higher $\beta$) significantly improves FID (from 4.53 to 3.75) compared to independent training.
  • The framework is robust to "capacity shifts," maintaining performance even when model capacity is moved from the latent generator to the encoder/decoder.
• Prior Flexibility: Experiments showed LSI works with diverse priors (Uniform, Laplacian, Gaussian Mixture), though Gaussian priors performed best. This confirms the method's ability to break free from the "simple prior" constraint of standard diffusion models.
• Sampling Flexibility: The paper demonstrates successful Classifier-Free Guidance (CFG) and inversion (encoding an image to a latent, then decoding it back), proving the model's controllability.

5. Significance and Impact

• Bridging the Gap: LSI successfully unifies the flexibility of Stochastic Interpolants (arbitrary priors, continuous-time formulation) with the efficiency of latent variable models.
• Scalability: By moving the generative process to a lower-dimensional latent space while maintaining a rigorous ELBO, LSI offers a scalable path for high-resolution image generation without the massive computational cost of pixel-space diffusion.
• Theoretical Rigor: The derivation of the ELBO in continuous time for latent variables provides a solid theoretical foundation for future research in joint latent generative modeling, moving beyond heuristic multi-stage training approaches.
• Practical Utility: The ability to use diverse priors and the significant reduction in sampling FLOPs make LSI a promising candidate for deployment in resource-constrained environments or for applications requiring specific distributional properties in the latent space.

In summary, Latent Stochastic Interpolants represent a significant advancement in generative modeling, offering a mathematically principled, efficient, and flexible alternative to both standard diffusion models and traditional VAEs.