Aligning Latent Spaces with Flow Priors

TL;DR

This paper introduces a flow-based prior method for aligning learnable latent spaces with arbitrary target distributions, improving efficiency and flexibility in generative modeling.

Contribution

It proposes a novel alignment loss that simplifies latent space regularization by leveraging flow models, eliminating the need for likelihood evaluations and ODE solving.

Findings

Effective alignment of latent spaces demonstrated on ImageNet

Alignment loss closely approximates negative log-likelihood

Theoretical proof of the surrogate objective's validity

Abstract

This paper presents a novel framework for aligning learnable latent spaces to arbitrary target distributions by leveraging flow-based generative models as priors. Our method first pretrains a flow model on the target features to capture the underlying distribution. This fixed flow model subsequently regularizes the latent space via an alignment loss, which reformulates the flow matching objective to treat the latents as optimization targets. We formally prove that minimizing this alignment loss establishes a computationally tractable surrogate objective for maximizing a variational lower bound on the log-likelihood of latents under the target distribution. Notably, the proposed method eliminates computationally expensive likelihood evaluations and avoids ODE solving during optimization. As a proof of concept, we demonstrate in a controlled setting that the alignment loss landscape closely approximates the negative log-likelihood of the target distribution. We further validate the effectiveness of our approach through large-scale image generation experiments on ImageNet with diverse target distributions, accompanied by detailed discussions and ablation studies. With both theoretical and empirical validation, our framework paves a new way for latent space alignment.