MixFlow Training: Alleviating Exposure Bias with Slowed Interpolation Mixture

TL;DR

MixFlow is a novel training method that addresses exposure bias in diffusion models by utilizing slowed interpolation mixtures, leading to improved image generation quality on benchmarks like ImageNet.

Contribution

The paper introduces MixFlow, a new training approach that leverages slowed interpolation mixtures to mitigate exposure bias in diffusion models.

Findings

Achieves state-of-the-art FID scores on ImageNet at 256x256 and 512x512 resolutions.

Effectively improves class-conditional and text-to-image generation.

Validates approach across multiple diffusion model architectures.

Abstract

This paper studies the training-testing discrepancy (a.k.a. exposure bias) problem for improving the diffusion models. During training, the input of a prediction network at one training timestep is the corresponding ground-truth noisy data that is an interpolation of the noise and the data, and during testing, the input is the generated noisy data. We present a novel training approach, named MixFlow, for improving the performance. Our approach is motivated by the Slow Flow phenomenon: the ground-truth interpolation that is the nearest to the generated noisy data at a given sampling timestep is observed to correspond to a higher-noise timestep (termed slowed timestep), i.e., the corresponding ground-truth timestep is slower than the sampling timestep. MixFlow leverages the interpolations at the slowed timesteps, named slowed interpolation mixture, for post-training the prediction network for each training timestep. Experiments over class-conditional image generation (including SiT, REPA, and RAE) and text-to-image generation validate the effectiveness of our approach. Our approach MixFlow over the RAE models achieve strong generation results on ImageNet: 1.43 FID (without guidance) and 1.10 (with guidance) at 256 x 256, and 1.55 FID (without guidance) and 1.10 (with guidance) at 512 x 512.