Curriculum Sampling: A Two-Phase Curriculum for Efficient Training of Flow Matching

1. Problem Statement

In Flow Matching (FM) models, the choice of the timestep sampling distribution $p(t)$ is a critical design decision. The standard practice has increasingly favored static, middle-biased distributions (such as Logit-Normal), which concentrate sampling density around the middle of the trajectory ($t \approx 0.5$).

The authors identify a fundamental speed–quality trade-off in these static strategies:

• Speed: Middle-biased sampling accelerates early convergence by focusing on the "transport" phase where the vector field is smoother.
• Quality: However, this approach systematically under-samples the boundary regimes ($t \to 0$ and $t \to 1$), leading to persistent errors in generating fine details and asymptotic fidelity. Consequently, models trained with middle-biased sampling often achieve worse final generation quality (higher FID) compared to Uniform sampling, despite converging faster initially.

The core problem is that treating $p(t)$ as a fixed hyperparameter fails to account for the evolving difficulty of the learning task across the trajectory.

2. Methodology: Curriculum Sampling

The authors propose Curriculum Sampling, a dynamic, two-phase training schedule that adapts the sampling distribution $p(t)$ over time to match the learning needs of the model.

Phase 1: Structure Learning (Rapid Convergence)

• Distribution: Logit-Normal (LN) sampling.
• Goal: Accelerate the learning of the global structure of the vector field.
• Mechanism: By concentrating samples in the middle regime ($t \approx 0.5$), the model quickly learns the primary transport dynamics. The authors found that allocating more updates toward the noise-side regime ($t \to 1$) during this phase yields the best convergence speed.

Phase 2: Boundary Refinement (High Fidelity)

• Distribution: Uniform sampling.
• Goal: Refine residual errors and resolve fine details.
• Mechanism: Once the global structure is established, the schedule switches to Uniform sampling. This ensures that the boundary regimes ($t \to 0$ and $t \to 1$), which correspond to the data and noise endpoints, are revisited with sufficient frequency to resolve the "U-shaped" difficulty profile where prediction errors are highest.

Switching Time ($T_s$)

The transition between phases is controlled by a hyperparameter $T_s$. The authors empirically determined that switching at approximately 40% of the total training budget (e.g., 60k steps out of 150k) yields optimal results, balancing structural acceleration with necessary boundary refinement.

3. Key Contributions

1. Identification of the Trade-off: The paper empirically demonstrates that middle-biased sampling (Logit-Normal) creates a performance ceiling due to the under-sampling of boundary conditions, while Uniform sampling, though slower initially, achieves superior asymptotic fidelity.
2. Loss Landscape Analysis: Through per-timestep loss monitoring, the authors characterize a U-shaped difficulty profile. The loss is lowest in the middle ($t \approx 0.5$) and peaks at the boundaries ($t \to 0, 1$), indicating that learning the final sharp details and initial noise departure is the most challenging part of the objective.
3. Curriculum Sampling Framework: They introduce a simple yet effective two-phase curriculum that transitions from a structure-focused distribution to a coverage-focused distribution, treating timestep sampling as an evolving curriculum rather than a static hyperparameter.
4. Efficiency Gains: The method achieves higher quality with fewer training steps, effectively decoupling the speed of convergence from the final quality ceiling.

4. Experimental Results

The method was evaluated on CIFAR-10 using a 55M parameter U-Net architecture for class-unconditional image generation.

• Performance (FID):
  • Uniform Baseline: Best FID of 3.85 (achieved at 150k steps).
  • Curriculum Sampling (LN $\to$ Uniform): Best FID of 3.22.
  • Improvement: This represents a 16.4% relative improvement in FID over the standard Uniform baseline and a 33% improvement over static Logit-Normal baselines.
• Convergence Speed:
  • The Curriculum Sampling method reached its peak performance at 100k steps, whereas the Uniform baseline required 150k steps.
  • This results in a 33% reduction in training steps to reach the best possible quality.
• Ablation Findings:
  • Switching too early (e.g., 50k) fails to leverage the acceleration of the middle-biased phase.
  • Switching too late (e.g., 70k+) leaves insufficient time for boundary refinement, preventing the FID from reaching its lowest floor.
  • Transitioning to a U-shaped Mode distribution in Phase 2 was less effective than transitioning to Uniform, suggesting that broad coverage is more critical than over-sampling boundaries at the expense of the transport path.

5. Significance and Impact

• Paradigm Shift: The paper challenges the notion that timestep sampling should be a fixed hyperparameter. Instead, it argues for dynamic, adaptive training protocols where the sampling distribution evolves to match the changing difficulty of the learning task.
• Practical Efficiency: The proposed curriculum offers a "free lunch" in terms of efficiency, delivering higher quality images in significantly less time. This is particularly valuable for large-scale generative model training where compute costs are high.
• Broader Implications: The findings suggest that inference-time timestep allocation could also be optimized non-uniformly, dedicating more resolution to time regimes where the learned vector field is hardest to approximate, potentially improving generation quality without retraining.

In conclusion, Curriculum Sampling resolves the speed-quality trade-off in Flow Matching by intelligently sequencing the learning process: first mastering the global structure via middle-biased sampling, then refining the details via uniform coverage, resulting in state-of-the-art efficiency and fidelity.