Designing a Conditional Prior Distribution for Flow-Based Generative Models

TL;DR

This paper introduces a novel approach for flow-based generative models that designs a conditional prior distribution based on input conditions, significantly improving training efficiency and sample quality in conditional generation tasks.

Contribution

The work proposes leveraging the ability to design non-trivial prior distributions in conditional flow models, leading to faster training and higher quality conditional samples.

Findings

Improved training times and generation efficiency.

Higher quality samples with fewer sampling steps.

Enhanced alignment scores (FID, KID, CLIP) compared to baselines.

Abstract

Flow-based generative models have recently shown impressive performance for conditional generation tasks, such as text-to-image generation. However, current methods transform a general unimodal noise distribution to a specific mode of the target data distribution. As such, every point in the initial source distribution can be mapped to every point in the target distribution, resulting in long average paths. To this end, in this work, we tap into a non-utilized property of conditional flow-based models: the ability to design a non-trivial prior distribution. Given an input condition, such as a text prompt, we first map it to a point lying in data space, representing an ``average" data point with the minimal average distance to all data points of the same conditional mode (e.g., class). We then utilize the flow matching formulation to map samples from a parametric distribution centered around this point to the conditional target distribution. Experimentally, our method significantly improves training times and generation efficiency (FID, KID and CLIP alignment scores) compared to baselines, producing high quality samples using fewer sampling steps.