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

This paper introduces Consistency Mid-Training (CMT), a novel intermediate training stage that stabilizes and accelerates the learning of flow map models like Consistency Models and Mean Flow, achieving state-of-the-art few-step generation quality with significantly reduced training data and computational costs compared to existing methods.

The Gist

The Big Picture: The "Instant Travel" Problem

Imagine you want to travel from your house (a random cloud of noise) to a specific destination (a beautiful, high-resolution image of a cat).

The Old Way (Diffusion Models):
Think of the old method as a very cautious hiker. To get from the noise to the cat, the hiker takes thousands of tiny, slow steps. At every single step, they stop to check a map, adjust their shoes, and ask, "Am I getting closer?"

• Pros: They almost always arrive at the right place.
• Cons: It takes forever. If you want to generate an image in real-time, this is too slow.

The New Goal (Flow Map Models):
Researchers wanted to build a "teleporter." Instead of taking thousands of steps, the model should learn to jump directly from the noise to the cat in just one or two giant leaps. This is called a "Flow Map."

The Problem:
Building a teleporter is incredibly hard. If you try to teach a model to jump from point A to point Z immediately, it usually gets confused. It tries to guess the path, but because it hasn't learned the "terrain" (the journey in between), it often crashes, produces blurry images, or takes forever to learn.

The Solution: CMT (The "Scout" Strategy)

The authors introduce CMT (Consistency Mid-Training).

To understand CMT, imagine you are training a student to be a master navigator.

1. Phase 1: The Expert (Pre-training)
   You hire a world-class expert hiker (a standard Diffusion Model). This expert knows the terrain perfectly. They can walk from the noise to the cat, taking 35 slow, careful steps. They never get lost.

2. Phase 2: The Scout (Mid-Training / CMT)
   This is the paper's big innovation.
   Instead of throwing the student into the deep end immediately, you put them in a "scout" phase.

   • You take the Expert's path (the 35 steps).
   • You show the student: "Look, if you are at step 10 of the Expert's journey, the final destination is right here."
   • You teach the student to look at any point along the Expert's path and instantly know where the final destination is.
   • The Magic: The student isn't guessing anymore. They are learning a direct map based on a path that is already proven to work. They learn the "shape" of the journey without having to take the slow steps themselves yet.

3. Phase 3: The Teleporter (Post-Training)
   Now, you take that "Scout" student and train them to be the final teleporter. Because they already understand the terrain so well (thanks to the Scout phase), they learn to make the giant jump from noise to cat incredibly fast and accurately.

Why is this a Game-Changer?

The paper compares this new method to the old ways of training teleporters:

• Random Start: Trying to teach a teleporter from scratch is like teaching someone to fly by throwing them off a cliff. They crash.
• Expert Transfer: Trying to just copy the Expert's weights is like giving a hiker a teleporter suit without explaining how it works. They still stumble because the "physics" of a teleporter are different from a hiker.
• CMT (The Scout): This gives the student a "cheat sheet" of the terrain.

The Results (The "Wow" Factor):
The paper shows that using CMT is like upgrading from a bicycle to a supersonic jet:

• Speed: It reduces the training time by up to 98%. In some cases, what used to take 4,000 hours of computer time now takes only 400 hours.
• Quality: The images generated are sharper and more realistic (lower FID scores) than previous methods, even with fewer steps (1 or 2 steps instead of 35).
• Stability: It stops the training process from "diverging" (crashing or going crazy), which was a major headache for researchers before.

A Simple Analogy: Learning to Drive

• Diffusion Model: Learning to drive by practicing in a parking lot for 10,000 hours, moving the car 1 inch forward, stopping, checking mirrors, moving 1 inch forward. Safe, but slow.
• Old Flow Map Training: Trying to learn to drive a race car at 200mph immediately. You will crash.
• CMT:
  1. You watch a professional driver (the Expert) drive the track perfectly.
  2. Mid-Training: You sit in the passenger seat and learn: "If we are at this curve, the finish line is there." You learn the relationship between the curve and the finish line without actually driving fast yet.
  3. Post-Training: Now you get behind the wheel. Because you already know the relationship between the road and the finish line, you can drive the race car at full speed immediately without crashing.

Summary

The paper solves the problem of making AI image generators fast without making them bad. They did this by inserting a "middle school" phase (Mid-Training) where the model learns to read a map of the journey before trying to run the race. This makes the whole process cheaper, faster, and much more reliable.

Technical Summary

1. Problem Statement

Flow map models, such as Consistency Models (CM) and Mean Flow (MF), aim to enable few-step (or one-step) generation by learning the direct solution map (integration) of the Probability Flow Ordinary Differential Equation (PF-ODE) used in diffusion models. While these models promise fast inference, their training faces three critical challenges:

• Instability: Training is highly sensitive to hyperparameters and often prone to divergence.
• Bias and Noise: Existing methods rely on "stop-gradient" pseudo-targets (e.g., using the model's own previous step or a teacher model's output) which drift during training, introducing bias and unstable optimization signals.
• High Cost: Training from scratch or even initializing from pre-trained diffusion models requires massive computational resources (GPU hours) and data budgets. Pre-trained diffusion models capture infinitesimal movements, whereas flow maps must learn large integrated jumps; this mismatch makes direct initialization fragile and slow to converge.

2. Methodology: Consistency Mid-Training (CMT)

The authors propose CMT, a novel "mid-training" stage inserted between the initial pre-training of a diffusion model (or a smaller flow map model) and the final flow map post-training.

Core Concept

CMT treats the learning of the flow map as a trajectory-consistent regression problem. Instead of learning from noisy self-generated targets or random initialization, CMT trains a model to map any point along a high-quality, pre-computed solver trajectory directly to the clean endpoint of that same trajectory.

The Three-Stage Pipeline

1. Stage 1: Pre-Training (Teacher Sampler):
   • Utilizes an existing, high-quality pre-trained diffusion model (e.g., EDM, EDM2) or a smaller flow map model.
   • This teacher is equipped with a deterministic ODE solver (e.g., DPM-Solver++) to generate high-fidelity trajectories from a prior sample $x_T$ to a clean sample $x_0$.
2. Stage 2: Mid-Training (CMT):
   • Objective: Train a student model to predict the clean endpoint ($x_0$) or the average drift between any two points on the teacher's trajectory.
   • Loss Function:
     • For Consistency Models (CM): The model $f_\theta$ is trained to map any intermediate state $\hat{x}_{t_i}$ on the teacher's trajectory directly to the clean origin $\hat{x}_0$.
       $$L_{CMT-CM} = \mathbb{E}_{i, x_T} [d(f_\theta(\hat{x}_{t_i}, t_i), \hat{x}_0)]$$
     • For Mean Flow (MF): The model learns the average drift $h_\theta$ between two points on the trajectory.
       $$L_{CMT-MF} = \mathbb{E}_{i>j, x_T} \left[ \left\| h_\theta(\hat{x}_{t_i}, t_i, t_j) - \frac{\hat{x}_{t_i} - \hat{x}_{t_j}}{t_i - t_j} \right\|^2 \right]$$
   • Key Advantage: The targets ($\hat{x}_0$ or finite differences) are fixed and explicit (derived from the teacher), eliminating the need for stop-gradients, complex time-weighting schedules, or ad-hoc heuristics.
3. Stage 3: Post-Training:
   • The weights from the CMT stage are used to initialize the final flow map model (e.g., ECT, ECD, or MF).
   • Because the initialization is already "trajectory-aligned," the subsequent post-training converges rapidly and stably with minimal engineering tricks.

3. Key Contributions

• Novel Paradigm: Introduces the concept of mid-training for vision generation, bridging the gap between diffusion pre-training and flow map post-training.
• Theoretical Guarantee: The paper provides a theoretical analysis (Theorem 5.1) showing that CMT initialization significantly reduces the gradient bias between the practical loss and the oracle flow map loss compared to random initialization or standard diffusion initialization. It proves that CMT yields a trajectory-aligned initializer that minimizes the discrepancy between the student and the true flow.
• Architecture Agnostic: The method works for both Consistency Models (based on EDM) and Mean Flow (based on Flow Matching), and can utilize different teacher samplers (diffusion models or smaller flow map models).
• Simplicity: Removes the need for complex training tricks like $\Delta t$ annealing, custom time sampling, or stop-gradient pseudo-targets.

4. Experimental Results

The authors evaluated CMT on multiple benchmarks (CIFAR-10, ImageNet 64x64/256x256/512x512, AFHQv2, FFHQ, and MS-COCO T2I).

• State-of-the-Art (SOTA) Performance:
  • CIFAR-10: 2-step FID of 1.97 (surpassing the teacher EDM's 2.01).
  • ImageNet 64x64: 2-step FID of 1.32.
  • ImageNet 512x512: 2-step FID of 1.84.
  • ImageNet 256x256: 1-step FID of 3.34 (beating MF from scratch at 3.43).
• Efficiency Gains:
  • Data Efficiency: Achieves SOTA results using up to 98% less training data (images processed) compared to baselines like sCT and sCD.
  • Compute Efficiency: Reduces total training time (GPU hours) by 91.4% on ImageNet 512x512 and ~50% on ImageNet 256x256 compared to training from scratch.
  • Convergence: On ImageNet 256x256, CMT reaches low FID in half the GPU hours required by vanilla MF, and generates semantically meaningful images much earlier in training than random or diffusion-initialized baselines.
• Robustness: Demonstrated effectiveness even when using a weak, smaller teacher model (MF-B/4) to train a larger student (MF-XL/2), proving the method's ability to transfer trajectory knowledge effectively.

5. Significance

This work fundamentally shifts the training paradigm for few-step generative models. By introducing a principled, trajectory-consistent intermediate stage, CMT solves the long-standing issues of instability and high computational cost in flow map learning.

• Practical Impact: It makes training high-quality, few-step generative models feasible with significantly lower resource budgets, democratizing access to SOTA generation speeds.
• Theoretical Insight: It clarifies that the instability in previous methods stems from the mismatch between the initialization and the true flow map objective, and that a "trajectory-aligned" initialization is the key to stable optimization.
• Generalizability: The framework is not limited to specific architectures, offering a general recipe for improving the training of any ODE-based generative model.

In summary, CMT establishes a new standard for training flow map models, achieving superior quality with drastically reduced costs by leveraging a simple yet powerful mid-training strategy.