MixFlow Training: Alleviating Exposure Bias with Slowed Interpolation Mixture

TL;DR

The paper tackles exposure bias in diffusion and flow-based generative models by uncovering the Slow Flow phenomenon and introducing MixFlow, a training-time technique that leverages slowed interpolations across a range of timesteps. By sampling slowed timesteps and training with a mixed input $\mathbf{x}_{m_t}$ at each training step $t$, MixFlow reduces the discrepancy between training and sampling inputs. Empirically, it yields consistent improvements across SiT, REPA, and especially RAE on ImageNet, achieving state-of-the-art gFID scores at 256×256 and strong results at 512×512, with pronounced gains at fewer sampling steps. The approach is simple to implement (a few code changes) and broadly compatible with various interpolation schemes, demonstrating a practical and scalable path to better diffusion-based generation.

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.

Figures (20)

• Figure 1: Illustrating (1) the Slow Flow phenomenon during the sampling process: the timestep (y-axis), corresponding to the ground truth noisy data that is the nearest to the generated noisy data at the sampling timestep $t$ (x-axis), is slower (with higher noise), i.e., the shading area is under the line $x=y$; and (2) the effectiveness of MixFlow training: the range of slowed timesteps for (b) MixFlow training is smaller and closer to the sampling steps than (a) standard training, indicating that MixFlow training effectively alleviates the training-testing discrepancy. The boundary of the shading area in (b) is plotted as blue lines in (a). Note: x-axis - the sampling timestep at which the noisy data is generated; y-axis - the slowed timestep corresponding to the ground truth noisy data that is the nearest to the generated noisy data; shading area - the range (the vertical line) of slowed timesteps at each sampling step; noise corresponds to timestep $0$, and data corresponds to timestep $1$. The slowed timestep ranges are obtained from $20,000$ training images in ImageNet deng2009imagenet, $50$ sampling steps, and SiT-B ma2024sit. Details on how to plot the figures are provided in Appendix \ref{['appendix:detailsforfigure1']}.
• Figure 2: Illustrating MixFlow training. (a) MixFlow training. At the training timestep $t$, the input noisy data is the interpolation $\mathbf{x}_{m_t}$ at the slowed timestep $m_t$. It is different from standard training: the input noisy data is the interpolation $\mathbf{x}_t$ at the timestep $t$. (b) Slowed timestep. $m_t$ is a timestep lying in the range $[(1-\gamma)t, t]$. The illustration in (b) is for $\gamma = 1$.
• Figure 3: A toy example illustrating the advantage of the MixFlow training over the standard training. The distribution from the model with the MixFlow training fits better the ground truth distribution than the standard training. Details about this toy example are provided in Appendix \ref{['appendix:detailsforfigure3']}.
• Figure 4: Studies of mixture range coefficient $\gamma$ for sampling the slowed timestep: $m_t \sim \mathcal{U}[(1-\gamma)t, t]$.
• Figure 5: Example results illustrating the advantages of MixFlow with respect to counting, spatial relationship and object shape. Left: SD 3.5; Middle: SD 3.5-ft-20k; Right: MixFlow.