Latent Stochastic Interpolants

TL;DR

Latent Stochastic Interpolants (LSI) enable joint learning of latent representations and generative models by extending the Stochastic Interpolants framework with an end-to-end trained encoder and decoder, improving efficiency and flexibility.

Contribution

This work introduces LSI, a novel method that combines latent space modeling with stochastic interpolants, allowing end-to-end training and better handling of high-dimensional data.

Findings

LSI effectively learns latent representations on ImageNet.

LSI outperforms traditional diffusion models in computational efficiency.

LSI maintains generative flexibility while reducing high-dimensional computational demands.

Abstract

Stochastic Interpolants (SI) is a powerful framework for generative modeling, capable of flexibly transforming between two probability distributions. However, its use in jointly optimized latent variable models remains unexplored as it requires direct access to the samples from the two distributions. This work presents Latent Stochastic Interpolants (LSI) enabling joint learning in a latent space with end-to-end optimized encoder, decoder and latent SI models. We achieve this by developing a principled Evidence Lower Bound (ELBO) objective derived directly in continuous time. The joint optimization allows LSI to learn effective latent representations along with a generative process that transforms an arbitrary prior distribution into the encoder-defined aggregated posterior. LSI sidesteps the simple priors of the normal diffusion models and mitigates the computational demands of applying SI directly in high-dimensional observation spaces, while preserving the generative flexibility of the SI framework. We demonstrate the efficacy of LSI through comprehensive experiments on the standard large scale ImageNet generation benchmark.