Curvature-Adaptive Consistency Flow Matching: Autonomous Trajectory Optimization via Reinforcement Learning¶

Conference: ECCV 2026
arXiv: 2606.22394
Code: None
Area: Diffusion Models / Image Generation
Keywords: Consistency Distillation, Flow Matching, Reinforcement Learning, Trajectory Optimization, Curriculum Learning

TL;DR¶

CACFM reformulates consistency distillation as a dynamic decision-making process, using a lightweight RL agent to automatically detect and prioritize training on high-curvature bottleneck stages of the PF-ODE trajectory. Combined with flow-adapted DMD and adversarial consistency loss, it achieves state-of-the-art performance with 4-step inference on FLUX and SDXL, outperforming FLUX-schnell by more than 2 FID points.

Background & Motivation¶

Flow Matching and its consistency distillation variants have become the dominant paradigm for accelerating diffusion model inference. Methods such as LCM and PCM segment the PF-ODE trajectory and learn a direct mapping from any point to the origin, significantly reducing inference steps. However, existing distillation methods exhibit a clear static nature in their sampling strategies: they either sample sub-trajectory segments uniformly, or directly reuse the Logit-Normal distribution from Rectified Flow training (which assumes the middle steps are the most crucial). They never dynamically adjust the sampling focus based on the optimization dynamics of the distillation process itself.

This paper reveals a key contradiction: the optimization difficulty distribution of consistency distillation is completely different from that of standard iterative generation. By calculating the Oracle Consistency Error under a converged teacher, the authors discover that the actual optimization bottlenecks lie in the boundary stages (initialization and final refinement), exhibiting a U-shaped distribution rather than the bell-shaped distribution concentrated in the middle assumed by Logit-Normal. Consequently, static sampling strategies waste significant computational resources on low-curvature linear segments while underfitting high-curvature transition regions, leading to structural collapse and blurry textures at extremely low inference steps.

Core Idea: Formulate the decision of "which trajectory segment to train on" as a reinforcement learning task—utilizing a lightweight Q-learning agent as a geometric probe to automatically perceive the optimization difficulty of each segment and dynamically allocate the training budget. This allows a coarse-to-fine curriculum learning strategy to emerge without any manual scheduling.

Method¶

Overall Architecture¶

CACFM reformulates consistency distillation as a closed-loop dynamic system: the RL agent perceives the geometric bottlenecks of different trajectory segments (state) and selects the sub-trajectory segment most in need of optimization (action); the student model is updated on this segment using a hybrid loss (consistency + adversarial + DMD); and the resulting change in loss is fed back as a reward to drive the agent's next decision. The entire framework inserts a very lightweight decision layer into the standard PCM distillation loop—its parameter scale is negligible, yet it fundamentally alters the underlying logic of training resource allocation.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Data (z, c)"] --> B["Observe Geometric State<br/>Sort segment losses to get state"]
    B --> C["RL Agent Selects Sub-trajectory<br/>epsilon-greedy + Q-learning"]
    C --> D["Teacher Solver Step<br/>+ Student Consistency Prediction"]
    D --> E["Hybrid Loss<br/>PCM + DMD + Adversarial"]
    E --> F["Model Update<br/>Gradient Descent + EMA Target"]
    F --> G["Reward Calculation + Q-table Update"]
    G --> B

Key Designs¶

1. RL Agent as Geometric Probe: Modeling sub-trajectory selection as an MDP and solving it via tabular Q-learning

The fundamental limitation of static sampling (Uniform / Logit-Normal) is its complete lack of awareness regarding the curvature heterogeneity along the PF-ODE trajectory—all segments receive equal training opportunities or are allocated resources based on static priors, thereby wasting computational power on already well-learned smooth segments. CACFM's solution is to treat "which segment to train" as a decision-making problem.

Specific implementation: Discretize the continuous trajectory \(t \in [0,1]\) into \(M=4\) semantic phases (Initialization, Structural Formation, Texture Filling, Final Refinement). The state \(s_t\) is defined as the ranking of the current consistency losses across all segments (resulting in \(M! = 24\) possible states), and the action \(a_t \in \{1,\dots,M\}\) is the selected sub-trajectory segment for training. The discretization of the state space is a deliberate design choice—tabular Q-learning with 24 states can converge within extremely few iterations, entirely eliminating the need for a neural network policy and incurring virtually zero extra computational overhead.

The reward design is critical for this MDP: different sub-trajectory segments naturally possess different loss scales, so directly using the raw loss reduction as the reward would introduce bias. CACFM adopts a baseline-advantage strategy: it maintains an EMA baseline \(B_t(m)\) of the consistency loss for each segment. The reward is defined as the improvement of the current loss relative to the baseline, \(r_t = \lambda_r \cdot (B_t(a_t) - \mathcal{L}_{\text{con}})\), where \(\lambda_r=100\) and the baseline update coefficient \(\beta=0.95\). The Q-learning update is calculated as \(Q(s_t, a_t) \leftarrow Q(s_t, a_t) + \alpha[r_t + \gamma \max_a Q(s_{t+1}, a) - Q(s_t, a_t)]\), with \(\alpha=0.1\), \(\gamma=0.9\), and the exploration strategy \(\epsilon\) linearly decaying from 1.0 to 0.1 over the first 20k steps.

Why it works: The baseline-advantage mechanism eliminates biases stemming from scale differences in segment losses, enabling the agent to learn "which segment yields the highest return" rather than simply "which segment has the highest loss." The negligible computational toll of tabular Q-learning allows the agent to be embedded into every single training iteration for real-time dynamic scheduling. Finally, using loss rankings rather than raw loss values as states provides normalization robustness across different batches.

2. Flow-adapted DMD + Score Estimation: Extracting score function from velocity field to achieve distribution-level alignment

Consistency losses constrain point-to-point trajectory mapping, which easily leads to mode collapse or blurry artifacts when operating under extremely low step counts (e.g., 4 steps)—an inherent limitation of single-point matching. CACFM introduces Distribution Matching Distillation (DMD) to force the student's generation distribution \(\mathbb{P}_{\boldsymbol{\theta}}^{\text{student}}(\mathbf{x}_0)\) closer to the teacher's distribution \(\mathbb{P}_{\boldsymbol{\phi}}^{\text{teacher}}(\mathbf{x}_0)\) at the distribution level, i.e., minimizing the KL divergence \(D_{\text{KL}}(\mathbb{P}_{\boldsymbol{\theta}}^{\text{student}} \| \mathbb{P}_{\boldsymbol{\phi}}^{\text{teacher}})\).

The gradient of DMD is approximated as \(\nabla_{\theta}\mathcal{L}^{\text{DMD}} \approx \mathbb{E}\left[(\boldsymbol{s}^{\text{teacher}}(\mathbf{x}_{\tau}) - \boldsymbol{s}^{\text{student}}(\mathbf{x}_{\tau})) \nabla_{\theta}\boldsymbol{f}_{\boldsymbol{\theta}}(\mathbf{x}_t, t)\right]\), which requires computing the score functions of both the teacher and the student. A key challenge is that Flow Matching models predict a velocity field \(\boldsymbol{v}\) rather than the noise \(\boldsymbol{\epsilon}\) used in diffusion models, meaning standard score estimation formulas from diffusion models cannot be directly applied. The authors leverage Tweedie's formula to derive the precise analytical relationship between score and velocity under the Flow Matching framework:

\[s_{\theta}(\mathbf{x}_t, t) = -\frac{\mathbf{x}_t + (1-t)\boldsymbol{v}_{\boldsymbol{\theta}}}{t}\]

This derivation is based on the optimal transport probability path \(p_t(\mathbf{x}|\mathbf{x}_0) = \mathcal{N}((1-t)\mathbf{x}_0, t^2\mathbf{I})\); the complete steps are provided in the paper's appendix. Additionally, the student's own score is estimated using the student model via the same formula, which achieves self-distillation and avoids introducing an extra discriminator network. The innovation in this step lies in seamlessly adapting DMD to the velocity-predicting Flow Matching framework, rather than simply replicating established practices from noise-predicting diffusion models.

3. Adversarial Consistency Loss: Maintaining manifold consistency via discriminator constraints to prevent self-distillation drift

The self-distillation mechanism of DMD (where the student uses its own score for guidance) carries an inherent risk: if the student's score estimation is inaccurate, the optimization direction of the KL divergence can deviate from the true distribution, causing generated results to drift away from the natural image manifold. CACFM introduces an adversarial consistency loss as a regularization constraint: a discriminator \(\mathcal{D}\) conditioned on the text prompt \(\boldsymbol{c}\) applies a hinge loss to the student's prediction \(\tilde{\mathbf{x}}_s\) and the teacher's target \(\hat{\mathbf{x}}_s\) (both perturbed with noise):

\[\mathcal{L}^{\mathrm{adv}} = \operatorname{ReLU}(1 + \mathcal{D}(\tilde{\mathbf{x}}_s, \boldsymbol{c})) + \operatorname{ReLU}(1 - \mathcal{D}(\hat{\mathbf{x}}_s, \boldsymbol{c}))\]

This loss encourages the student's predictions to be indistinguishable from the teacher's from the perspective of the discriminator, essentially serving as an adversarial test of "whether the consistency mapping falls on the natural image manifold." This is complementary to DMD's distribution matching objective: while DMD brings distributions closer at the statistical level, the adversarial loss penalizes outliers at the sample level. Their combination effectively suppresses artifacts and distortions under extremely low step counts.

A Concrete Example¶

Taking FLUX 4-step distillation as an example: at the start of training, the initial consistency loss ranking for the 4 sub-segments (Phases 0-3) is [Phase 1, Phase 2, Phase 0, Phase 3] (where the state is one of the 24 possible permutations). The Q-table has not yet learned an effective policy, so a segment is selected randomly under \(\epsilon=1.0\). Assuming Phase 3 (Final Refinement) is chosen, the loss for this segment decreases from 0.8 to 0.65 after one training step, yielding a reward of \(100 \times (0.8 - 0.65) = 15\). After updating the Q-table, the Q-value for action = Phase 3 under this state increases.

By step 10k, \(\epsilon\) has decayed to approximately 0.5, and the agent has learned from the Q-table that the boundary stages (Phase 0 and Phase 3) yield systematically higher rewards than the middle stages. Consequently, the agent chooses boundary segments roughly 70% of the time. After 20k+ steps, the agent further manifests sequential preferences: it first conducts dense training on Phase 0 (global structure formation) and, once the structure stabilizes, automatically shifts its focus to Phase 3 (high-frequency detail refinement). The entire process operates without any manual curriculum scheduling, driven purely by the reward.

Loss & Training¶

The total loss is a weighted sum of three terms:

\[\mathcal{L}_{\text{total}} = \mathcal{L}^{\text{PCM}} + 0.1\mathcal{L}^{\mathrm{adv}} + 0.5\mathcal{L}^{\mathrm{DMD}}\]

where \(\mathcal{L}^{\text{PCM}}\) is the standard phased consistency loss (PCM loss), \(\mathcal{L}^{\text{DMD}}\) is the flow-adapted DMD loss, and \(\mathcal{L}^{\mathrm{adv}}\) is the adversarial consistency loss. The reward is calculated based solely on \(\mathcal{L}^{\text{PCM}}\) to prevent the high variance of the DMD and adversarial losses from interfering with the agent's decision-making. The target network \(\boldsymbol{\theta}^{-}\) is updated via EMA. RL hyperparameters: \(\alpha=0.1, \gamma=0.9\), \(\epsilon\) scales from \(1.0 \to 0.1\) (linearly decaying over the first 20k steps), \(\lambda_r=100\), and \(\beta=0.95\). Training is conducted on the LAION dataset, and evaluation is performed on a random 15k subset of CC3M.

Key Experimental Results¶

Main Results¶

Table 1: FID comparison on FLUX (CC3M 15K), lower is better

Methods	4-Step	8-Step	16-Step
Turbo	58.22	46.24	43.66
Hyper-SD	45.35	43.86	43.40
TDD	45.81	41.65	41.12
Schnell	41.19	40.47	39.89
CACFM (Ours)	39.19	36.96	37.62

On FLUX, CACFM achieves state-of-the-art results across all step configurations. It outperforms Schnell by 2.0 FID points on 4-step generation, and by 3.5 FID points on 8-step generation. The gap is most pronounced at extremely low step budgets (4 steps), validating the critical role of curvature-adaptive training in computationally constrained scenarios.

Table 2: FID comparison on SDXL (CC3M 15K), lower is better

Steps	Lightning	Turbo	LCM	Hyper-SD	PCM	InstaFlow	TDD	TCD	Ours
4	37.49	52.90	45.57	39.43	37.26	38.13	41.75	46.40	35.29
8	38.28	65.25	43.67	41.63	39.30	35.60	46.00	49.51	34.42
16	40.22	77.13	43.33	44.12	40.47	34.43	51.22	54.68	33.49

It similarly secures SOTA performance across all steps on SDXL, with its 4-step FID being 1.97 lower than the strongest baseline, PCM. Notably, CACFM steadily improves as the step count increases (e.g., 39.19 → 36.96 → 37.62 on FLUX), whereas baselines like Turbo often degrade at higher steps. This indicates that RL-guided optimization successfully "straightens" the probability flow trajectory rather than merely overfitting to a specific step configuration.

In human preference and subjective quality evaluations (HPSv2 / Aesthetic / PickScore), CACFM ranks in the top 3 across all step configurations on both backbones, often securing 1st place. The advantage is particularly prominent in Aesthetic scores (e.g., FLUX 4-step: 5.80 vs. Schnell: 5.52), demonstrating that the emerged curriculum learning for the final refinement stage directly enhances perceived detail quality.

Ablation Study¶

Table 3: Ablation study on SDXL 4-step

Configuration	HPS	Aesthetic	PickScore	Avg Rank
Naive CFM (Uniform)	0.236	5.432	20.69	4.11
CFM Logit-Normal	0.232	5.458	20.70	4.89
CFM Loss-Aware (EMA)	0.236	5.436	20.77	4.00
CACFM w/o DMD	0.239	5.485	20.78	3.78
CACFM w/o RL	0.235	5.430	20.71	3.22
CACFM (full)	0.276	5.893	21.22	1.00

Key Findings: - RL agent contributes the most: Removing RL (CACFM w/o RL) drops the Avg Rank from 1.00 to 3.22, and the Aesthetic score plunges from 5.893 to 5.430—demonstrating that geometric-aware scheduling is the core driving force. - Logit-Normal performs worse than Uniform (Rank 4.89 vs. 4.11): This directly validates the core finding of this work—the difficulty distribution of consistency distillation is U-shaped rather than bell-shaped, meaning that blindly carbon-copying the sampling prior of Rectified Flow is counterproductive. - DMD contributes significantly to aesthetic quality: Without DMD, the Aesthetic score drops from 5.893 to 5.485, confirming the critical role of distribution alignment in fine-grained details. - Superlinear gains of RL: The improvements from DMD alone (without RL) or RL alone (without DMD) are modest; however, their combination (full CACFM) produces gains far exceeding their sum—indicating a robust synergy between geometric-aware scheduling and distribution alignment.

Key Findings¶

Wall-clock Efficiency: CACFM incurs about 18% more overhead per step compared to PCM (due to the RL agent + discriminator forward pass). However, it consistently leads under a fixed time budget: 24 hours of CACFM training (FID 40.82) already outperforms 36 hours of PCM training (FID 41.65), demonstrating that the data efficiency gains from RL-guided curriculum learning far outweigh the computational overhead of the extra modules.
Emergent Curriculum Learning: The selection frequency of the RL agent aligns perfectly with the U-shaped oracle difficulty distribution (correlation coefficient \(\rho > 0.95\)). It spontaneously transitions from focusing on Phase 0 (global structure formation) in early training to focusing on Phase 3 (high-frequency detail refinement) in late training, without any manual scheduling.
Zero-shot Step Generalization: The RL policy trained with \(M=4\) continues to yield quality improvements when the inference step count is increased to 8 or 16 steps. This indicates that the optimization indeed straightens the trajectory rather than merely overfitting to the designated training steps—a key advantage over pure distillation methods that typically overfit to fixed steps.

Highlights & Insights¶

Turning "where to learn" into a learnable problem: While most acceleration methods focus on "how to learn" (better losses, better architectures), CACFM shifts the focus to "where to learn"—using an extremely lightweight RL agent (24-state Q-table, virtually zero extra overhead) to achieve a major leap in training efficiency. This paradigm can be generalized to any training scenario with trajectory/stage selection degrees of freedom, such as multi-stage VLM training or stage allocation in RLHF reward model training.
Clever Baseline-Advantage Reward Design: Different sub-segments naturally have distinct loss scales; directly using raw losses for rewards would degrade the agent's policy into a greedy strategy of "always selecting the segment with the highest loss." The EMA baseline eliminates this inherent bias, allowing the agent to focus on "marginal gains" rather than absolute difficulty—a highly reusable RL reward design trick.
Logit-Normal is Wrong in Distillation: This is a major "counter-intuitive" finding—the community habitually reuses the Logit-Normal sampling prior from Rectified Flow training, but this work proves via oracle consistency error experiments that the distillation difficulty distribution is U-shaped rather than bell-shaped. This insight itself serves as a crucial warning for the consistency distillation research direction.

Limitations & Future Work¶

Upper bound limited by the teacher model: The authors candidly acknowledge that the maximum generation quality of CACFM is capped by the teacher's vector field manifold and cannot exceed it—a shared limitation of all distillation approaches.
Single-step generation remains challenging: When extremely compressed to 1 step, highly curved topological flows cannot be perfectly approximated by a single linear step, still leading to minor artifacts. The authors did not provide 1-step experimental results.
Granularity selection of M requires manual tuning: Appendix experiments indicate that \(M=4\) is the optimal trade-off point (\(M=3\) is too coarse, while \(M=6\) yields a state space of 720, rendering exploration highly inefficient). However, this choice was determined strictly via empirical search without theoretical guidelines—requiring re-searching on larger models or different tasks.
No comparison with more complex RL algorithms: This work only utilizes tabular Q-learning and does not compare against deep RL methods like PPO under larger segment divisions (\(M>6\)). As the number of segments scales up, tabular methods suffer from state-explosion, which warrants neural network policies for future exploration.

vs. PCM / TCD / Hyper-SD (Static Consistency Distillation): These methods focus on improving "how to match" (finer partitioning schemes, more stable training tricks) but uniformly utilize static sampling priors. The core distinction of CACFM is that it treats sampling as a learning problem—letting RL decide which segment to train rather than assuming a fixed difficulty distribution.
vs. DDPO / DPOK (RL for Alignment): These methods use RL to optimize alignment with human/aesthetic preferences (what to generate), whereas CACFM uses RL to optimize the training dynamics of the distillation process itself (sub-trajectory selection). While they both fall under the broad scope of "RL for generative models", they operate on completely different levels—one targets "what to generate", the other targets "how to learn to generate".
vs. Turbo / Lightning (Adversarial Distillation): They utilize adversarial losses to boost few-step generation quality. While CACFM also incorporates an adversarial consistency loss, it plays an auxiliary role (regularization + stabilizing DMD self-distillation). The main driver remains the RL geometric probe—removing RL and relying solely on DMD + adversarial loss yields significantly lower improvements.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ First to introduce RL into Consistency Distillation training scheduling, identifying the U-shaped difficulty distribution and designing adaptive sampling accordingly. Both the problem formulation and the solution are novel, rather than incremental updates to existing frameworks.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Double backbones (FLUX + SDXL), multiple inference steps (4/8/16), multiple evaluation metrics (FID + HPS + Aesthetic + PickScore), wall-clock efficiency analysis, visualization of RL policies (difficulty distribution + curriculum emergence heatmaps), and comparisons of ablation + sampling strategies (Loss-Aware / MAB baseline). The appendix further provides sensitivity analyses for \(M\) and Q-learning convergence.
Writing Quality: ⭐⭐⭐⭐ The core insight (U-shaped difficulty vs. Logit-Normal bell shape) is consistently threaded throughout the paper; the strategy visualization in Figure 6 is very compelling. The DMD score derivation is technical but mathematically sound. A minor drawback is that the FIDs of some baseline methods vary greatly across tables (e.g., Turbo's 16-step FID reaches 77 on SDXL) without sufficient discussion.
Value: ⭐⭐⭐⭐⭐ The concept of "adaptive training scheduling" features strong transferability—any learning process with discrete stages and dynamic difficulty variations (e.g., multi-task learning, curriculum learning, multi-stage RLHF training) can potentially benefit. Furthermore, the approach incurs negligible computational cost, making it highly practical.