CMT: Mid-Training for Efficient Learning of Consistency, Mean Flow, and Flow Map Models (2509.24526v1)

Abstract: Flow map models such as Consistency Models (CM) and Mean Flow (MF) enable few-step generation by learning the long jump of the ODE solution of diffusion models, yet training remains unstable, sensitive to hyperparameters, and costly. Initializing from a pre-trained diffusion model helps, but still requires converting infinitesimal steps into a long-jump map, leaving instability unresolved. We introduce mid-training, the first concept and practical method that inserts a lightweight intermediate stage between the (diffusion) pre-training and the final flow map training (i.e., post-training) for vision generation. Concretely, Consistency Mid-Training (CMT) is a compact and principled stage that trains a model to map points along a solver trajectory from a pre-trained model, starting from a prior sample, directly to the solver-generated clean sample. It yields a trajectory-consistent and stable initialization. This initializer outperforms random and diffusion-based baselines and enables fast, robust convergence without heuristics. Initializing post-training with CMT weights further simplifies flow map learning. Empirically, CMT achieves state of the art two step FIDs: 1.97 on CIFAR-10, 1.32 on ImageNet 64x64, and 1.84 on ImageNet 512x512, while using up to 98% less training data and GPU time, compared to CMs. On ImageNet 256x256, CMT reaches 1-step FID 3.34 while cutting total training time by about 50% compared to MF from scratch (FID 3.43). This establishes CMT as a principled, efficient, and general framework for training flow map models.

• The paper introduces a mid-training stage (CMT) that bridges diffusion pre-training and flow map post-training, reducing instability.
• The method leverages teacher-generated trajectories and explicit regression targets to minimize gradient bias and accelerate convergence.
• Empirical results on CIFAR-10 and ImageNet show up to 98% lower training cost with few-step synthesis comparable to multi-step models.

Efficient Mid-Training for Flow Map Models: The CMT Framework

Introduction and Motivation

The paper introduces Consistency Mid-Training (CMT), a mid-training paradigm for efficient and stable learning of flow map models, including Consistency Models (CM) and Mean Flow (MF), in generative modeling. Flow map models directly learn the solution map of the probability flow ODE (PF-ODE) underlying diffusion models, enabling few-step generation. However, training these models is typically unstable, sensitive to hyperparameters, and computationally expensive. Existing approaches often initialize from pre-trained diffusion models, but this does not resolve the fundamental mismatch between infinitesimal-step diffusion weights and the large integrated jumps required for flow map learning, resulting in slow convergence and instability.

CMT addresses these issues by introducing a lightweight, trajectory-aligned mid-training stage between diffusion pre-training and flow map post-training. The CMT stage leverages trajectories generated by a pre-trained teacher model to provide explicit regression targets, yielding a stable and trajectory-consistent initializer for subsequent flow map training.

CMT Pipeline and Algorithmic Details

The CMT pipeline consists of three stages:

1. Pre-Training: Train a diffusion model or a lightweight flow map model to serve as a deterministic ODE sampler (teacher).
2. Mid-Training (CMT): Train a student model to map any point along a teacher-generated trajectory (from a prior sample) directly to the clean endpoint of that trajectory. The loss is a standard regression to a fixed target, avoiding stop-gradient surrogates and ad hoc heuristics.
3. Post-Training: Initialize the flow map model with CMT weights and perform final training using standard flow map objectives (e.g., ECT, ECD, MF).

The CMT loss for CM is:

$$\mathcal{L}_{\text{CMT}}(\theta) = \mathbb{E}_{i} \mathbb{E}_{x_T \sim p_{\text{prior}}} \left[ d(f_\theta(\hat{x}_{t_i}, t_i), \hat{x}_{t_0}) \right]$$

where $\{\hat{x}_{t_i}\}$ is a trajectory generated by the teacher ODE solver, and $d$ is a perceptual or $\ell_2$ distance.

For MF, the CMT loss generalizes to:

$$\mathcal{L}_{\text{CMT}}^{\text{MF}} = \mathbb{E}_{i>j} \mathbb{E}_{x_T \sim p_{\text{prior}}} \left[ \| f(\hat{x}_{t_i}, t_i, t_j) - \frac{\hat{x}_{t_i} - \hat{x}_{t_j}}{t_i - t_j} \|_2^2 \right]$$

This mid-training stage is architecture-agnostic and can be applied to any flow map parameterization.

Theoretical Analysis

The paper provides a rigorous analysis of gradient bias and variance for different initialization schemes. The main result is that CMT initialization yields a gradient bias to the oracle flow map loss of $\mathcal{O}(\varepsilon + \Delta t^2)$, where $\varepsilon$ is the mid-training error and $\Delta t$ is the solver discretization. In contrast, diffusion-based initialization incurs additional irreducible bias terms due to the mismatch between the PF-ODE solution and the posterior mean, as well as the scaling of noise schedules. Random initialization is even less favorable, leading to slow and unstable convergence.

The SGD analysis shows that CMT initialization achieves the lowest excess risk and fastest convergence to the oracle solution among all considered schemes, with no need for ad hoc tricks or complex hyperparameter schedules.

Empirical Results

CMT demonstrates strong empirical performance across a range of datasets and resolutions:

• CIFAR-10 (32×32): CMT (w/ ECT) achieves 2-step FID of 1.97, outperforming all prior flow map models and matching the teacher diffusion model with 35 steps.
• ImageNet (64×64): CMT (w/ ECT/ECD) attains 2-step FID of 1.32, using only 19.2M training images (mid + post), compared to 102.4M for vanilla ECD and up to 819.2M for sCT.
• ImageNet (512×512): CMT (w/ ECD) achieves 2-step FID of 1.84 at 93% lower cost than previous sCD, with only 28.8M training images and 403 GPU hours.

  Figure 1: Two-step generated images by CMT (w/ ECD) on ImageNet 512×512, achieving FID 1.84 at 93% lower cost than previous sCD.

• ImageNet (256×256, MF): CMT initialization enables MF-XL/2 to reach FID 3.34 in 760 GPU hours, cutting training time by 50% compared to MF from scratch (FID 3.43).

CMT consistently improves post-training stability, accelerates convergence, and enhances final generation quality, while reducing training data and compute by up to 98% compared to baselines.

Implementation Considerations

• Teacher Trajectory Generation: Use high-order ODE solvers (e.g., DPM-Solver++) with moderate NFEs (e.g., 16) to generate reference trajectories. This balances supervision quality and computational cost.
• Loss Metric: Employ perceptual losses (LPIPS, ELatentLPIPS) for pixel and latent space models, respectively, to align with human visual fidelity. For MF, use squared $\ell_2$ loss due to the generality of the mapping.
• Batching and Data Efficiency: CMT leverages multiple intermediate points per trajectory, yielding high data efficiency and fast convergence.
• Hyperparameter Robustness: CMT eliminates the need for custom time sampling, weighting schedules, or learning rate annealing, simplifying engineering and improving reproducibility.
• Architecture Agnosticism: CMT can be applied to any flow map backbone (EDM, FM, etc.) and supports flexible teacher-student configurations.

Practical and Theoretical Implications

CMT establishes a principled and general framework for efficient flow map learning in generative modeling. By bridging the gap between diffusion pre-training and flow map post-training, CMT enables stable, fast, and high-fidelity few-step generation. The approach is broadly applicable to ODE-based generative models and can be extended to other domains (e.g., audio, video, latent spaces).

Theoretically, CMT provides a strong initialization that closely aligns with the oracle flow map, minimizing gradient bias and variance, and ensuring robust optimization dynamics. Practically, it dramatically reduces training cost and engineering complexity, making flow map models viable for large-scale, high-resolution synthesis.

Future Directions

Potential future developments include:

• Extending CMT to stochastic or non-deterministic teacher trajectories (e.g., SDE-based models).
• Applying CMT to conditional and multimodal generative tasks.
• Investigating CMT for other ODE-based architectures beyond vision (e.g., molecular generation, time series).
• Exploring automated teacher selection and trajectory generation for optimal supervision.

Conclusion

CMT introduces a mid-training stage that efficiently and stably initializes flow map models for generative modeling. By leveraging teacher-generated trajectories and explicit regression targets, CMT achieves state-of-the-art sample quality with minimal training cost and engineering overhead. The framework is theoretically justified, empirically validated, and broadly applicable, representing a significant advance in the practical training of few-step generative models.