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. It utilizes a lightweight RL agent to automatically detect and prioritize high-curvature bottleneck segments along the PF-ODE trajectory for training. Combining flow-adapted DMD and adversarial consistency loss, it achieves state-of-the-art results on FLUX and SDXL with 4-step inference, delivering an FID over 2 points lower than FLUX-schnell.

Background & Motivation¶

Flow Matching and its consistency distillation methods have become the mainstream paradigm for accelerating diffusion model inference. Methods like LCM and PCM segment the PF-ODE trajectory and learn a direct mapping from any point on it to the origin, dramatically reducing inference steps. However, existing distillation methods exhibit a distinct static nature in their sampling strategies: they either sample sub-trajectory segments uniformly or directly reuse the Logit-Normal distribution from Rectified Flow training (assuming middle steps are most important), failing to dynamically adjust sampling priorities according to the optimization dynamics of the distillation process itself.

This work uncovers a key challenge: the optimization difficulty distribution in consistency distillation differs fundamentally from that of standard iterative generation. By computing the Oracle Consistency Error under a converged teacher, the authors find that the actual optimization bottlenecks lie in the boundary phases (initialization and final refinement), exhibiting a U-shaped distribution rather than the bell-shaped distribution peaked in the middle assumed by Logit-Normal. Consequently, static sampling strategies waste significant compute on low-curvature linear segments while underfitting high-curvature transition zones, leading to structural collapse and blurry textures at extremely low step regimes.

Core Idea: Reformulate "which trajectory segment to train on" as a decision-making problem learnable via RL. A lightweight Q-learning agent is employed as a geometric probe to automatically perceive optimization difficulties across different segments and dynamically allocate the training budget, allowing a coarse-to-fine curriculum learning strategy to emerge without manually designed schedulers.

Method¶

Overall Architecture¶

CACFM reformulates consistency distillation as a closed-loop dynamic system: the RL agent perceives the geometric bottlenecks of each trajectory segment (states), selects the sub-trajectory segment most worth optimizing (actions), and updates the student model on this segment using a hybrid loss (consistency + adversarial + DMD). The resulting shift in loss is fed back as a reward to drive the agent's next decision. This framework inserts an extremely lightweight decision-making layer into the standard PCM distillation loop—bringing negligible parameter overhead while rewriting the underlying logic of training resource allocation.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入数据 (z, c)"] --> B["观测几何状态<br/>对各段 Loss 排序得 state"]
    B --> C["RL Agent 选择子轨迹<br/>epsilon-greedy + Q-learning"]
    C --> D["Teacher Solver 步<br/>+ Student 一致性预测"]
    D --> E["混合损失<br/>PCM + DMD + Adversarial"]
    E --> F["模型更新<br/>梯度下降 + EMA Target"]
    F --> G["Reward 计算 + Q-table 更新"]
    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 of PF-ODE trajectories. Because all segments are trained with equal probability or according to a fixed prior, computational resources are wasted on smooth segments that are already well-learned. CACFM addresses this by formulating "where to train" as a decision-making problem.

Specifically, the continuous trajectory \(t \in [0,1]\) is discretized into \(M=4\) semantic phases (Initialization, Structural Formation, Texture Filling, Final Refinement). The state \(s_t\) is defined as the rank-ordering of the current consistency loss across all phases (yielding \(M! = 24\) possible states), and the action \(a_t \in \{1,\dots,M\}\) selects the sub-trajectory phase for training. The discretization of the state space is highly intentional—tabular Q-learning with just 24 states converges within very few iterations, eliminating the need for a neural network policy and incurring virtually zero computational overhead.

The design of the reward is crucial to this MDP: different sub-trajectory segments naturally possess varying scales of loss, so using raw loss reduction directly as a reward would introduce bias. CACFM adopts a baseline-advantage strategy: maintaining an EMA baseline \(B_t(m)\) of the consistency loss for each phase, 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 decay factor \(\beta=0.95\). Tabular Q-learning is updated via \(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 rate \(\epsilon\) is linearly decayed from 1.0 to 0.1 over the first 20k steps.

Why it works: The baseline-advantage mechanism removes the bias stemming from differing scale levels of loss across phases, enabling the agent to learn "which choice offers higher relative gains" instead of simply targeting "which phase has the highest absolute loss." The minimal computational cost of tabular Q-learning allows the agent to be embedded into every training iteration for real-time dynamic scheduling. Utilizing state rankings rather than raw loss values as states provides normalization robustness across different batches.

2. Flow-adapted DMD + Score Estimation: Analytical derivation of the score function from velocity fields for distribution-level alignment

Consistency losses constrain point-to-point trajectory mapping, which often suffers from mode collapse or blurry artifacts in extreme low-step regimes (e.g., 4 steps)—an intrinsic limitation of point-wise matching. CACFM incorporates Distribution Matching Distillation (DMD) to force the student generation distribution \(\mathbb{P}_{\boldsymbol{\theta}}^{\text{student}}(\mathbf{x}_0)\) to match the teacher distribution \(\mathbb{P}_{\boldsymbol{\phi}}^{\text{teacher}}(\mathbf{x}_0)\) at the distribution level, 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 teacher and student. The key difficulty is that the Flow Matching model predicts a velocity field \(\boldsymbol{v}\) rather than the noise \(\boldsymbol{\epsilon}\) predicted in standard diffusion models, meaning existing score estimation formulas cannot be directly applied. Utilizing Tweedie's formula, the authors derive a 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})\), with the complete proof provided in the appendix. Additionally, the student's own score is estimated using the student model through the exact same equation to achieve self-distillation, avoiding the need for an extra discriminator network. The novelty lies in seamlessly adapting DMD to the velocity-predicting Flow Matching framework rather than simply copying existing practices designed for noise-predicting diffusion models.

3. Adversarial Consistency Loss: Utilizing discriminator constraints to maintain manifold consistency and prevent self-distillation drift

The self-distillation mechanism of DMD (where the student uses its own score predictions for guidance) poses an inherent risk: if the student's score estimation is inaccurate, the optimization direction of the KL divergence may drift away from the true distribution, causing generated samples to fall outside the natural image manifold. To mitigate this, CACFM introduces an adversarial consistency loss as a regularization constraint, employing a discriminator \(\mathcal{D}\) conditioned on the text prompt \(\boldsymbol{c}\) to enforce a hinge loss on the student prediction \(\tilde{\mathbf{x}}_s\) and the teacher 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 predictions to be indistinguishable from the teacher under the discriminator, essentially performing an adversarial check on whether the consistency mapping falls onto the natural image manifold. This is highly complementary to DMD's distribution matching objective: while DMD aligns distributions at a statistical level, the adversarial loss penalizes outliers at the sample level. Achieving both simultaneously effectively suppresses artifacts and distortions under extremely few steps.

A Complete Example¶

Taking FLUX 4-step distillation as an example: at the start of training, the initial consistency losses of the 4 sub-phases (Phase 0-3) are ranked as [Phase 1, Phase 2, Phase 0, Phase 3] (representing one of the 24 possible permutation states). Because the Q-table has not yet learned an effective policy, exploration proceeds randomly with \(\epsilon=1.0\). Assume Phase 3 (Final Refinement) is selected; after one training step, the loss for this phase decreases from 0.8 to 0.65, yielding a reward of \(100 \times (0.8 - 0.65) = 15\). Following the Q-table update, the Q-value for action=Phase 3 under this state increases.

By step 10k, \(\epsilon\) decays to around 0.5, and the agent has learned from the Q-table that rewards for boundary segments (Phase 0 and Phase 3) are systematically higher than those for intermediate segments. Consequently, the agent chooses boundary segments for training about 70% of the time. Beyond 20k+ steps, the agent further manifests phase preferences: starting with intensive training on Phase 0 (global structure formation), and then automatically shifting focus to Phase 3 (high-frequency detail refinement) once the structure stabilizes. This entire process occurs without manual curriculum scheduling, driven purely by the reward signal.

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 solely based 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 using EMA. RL hyperparameters: \(\alpha=0.1, \gamma=0.9, \epsilon: 1.0 \to 0.1\) (linearly decayed over the first 20k steps), \(\lambda_r=100, \beta=0.95\). Training data uses LAION, and evaluation is conducted 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 evaluated inference steps. It surpasses Schnell by 2.0 FID in 4-step generation and by 3.5 FID in 8-step generation. The gap is most pronounced in the extremely low-step regime (4 steps), validating the crucial role of curvature-adaptive training when computational budget is highly constrained.

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

SOTA is likewise achieved across all step counts on SDXL, with 4-step FID being 1.97 lower than the strongest competitor PCM. Notably, performance consistently improves for CACFM as inference steps increase (e.g., 39.19 \(\rightarrow\) 36.96 \(\rightarrow\) 37.62 on FLUX), whereas baseline methods like Turbo often degenerate at higher steps. This indicates that RL-guided optimization genuinely "straightens" probability flow trajectories rather than merely overfitting to a specific training step count.

In human/subjective quality evaluations (HPSv2 / Aesthetic / PickScore), CACFM consistently ranks in the top 3 (mostly 1st) across all step configurations for both backbones. Its advantage is particularly outstanding in the Aesthetic metric (FLUX 4-step: 5.80 vs. Schnell: 5.52), indicating that the emergent curriculum learning focused on the refinement phase directly improves 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) causes the Avg Rank to drop from 1.00 to 3.22, and the Aesthetic score to plunge from 5.893 to 5.430, confirming 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 paper—the optimization difficulty of consistency distillation follows a U-shaped distribution rather than a bell curve. Blindly carrying over the sampling prior from Rectified Flow training proves counterproductive. - DMD Significantly Contributes to Aesthetic Quality: Without DMD, the Aesthetic score drops from 5.893 to 5.485, verifying the critical role of distribution-level alignment in preserving details. - Super-linear Gains of RL: Neither DMD alone (w/o RL) nor RL alone (w/o DMD) yields substantial improvements. However, combining them (full CACFM) produces a gain far exceeding the sum of their individual parts, demonstrating a strong synergy between geometric-aware scheduling and distribution matching.

Key Findings¶

Wall-clock Efficiency: CACFM incurs roughly 18% higher per-step overhead than PCM (due to the RL agent and discriminator forward pass). However, it consistently dominates within a fixed training budget: 24h of CACFM training (FID 40.82) already outperforms 36h of PCM training (FID 41.65). This indicates that the data efficiency improvements introduced by RL-guided curriculum learning far outweigh the minor computation overhead of the extra modules.
Emergent Curriculum Learning: The agent's phase selection frequency perfectly aligns 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 prioritizing Phase 3 (high-frequency detail refinement) in later stages, without any hand-crafted scheduling.
Zero-shot Step Generalization: The RL policy trained with \(M=4\) continuously yields quality improvements when inference steps are expanded to 8 or 16. This demonstrates that the optimization truly straightens the trajectory rather than merely overfitting to a specific number of training steps—a key advantage over pure distillation methods that often overfit to a fixed step count.

Highlights & Insights¶

Formulating "where to learn" as a learnable problem: While most acceleration techniques focus on "how to learn" (e.g., better design of losses and architectures), CACFM shifts focus to "where to learn." By leveraging an extremely lightweight RL agent (a 24-state Q-table with virtually zero computational overhead), it achieves a remarkable leap in training efficiency. This design philosophy can be translated to any training scenario with degrees of freedom in trajectory/stage selection, such as multi-stage VLM training or training stage allocation for RLHF reward models.
Ingenious baseline-advantage reward design: Because different sub-phases naturally possess varying loss scales, using the raw loss directly as a reward would result in a degenerate, greedy policy of "always selecting the phase with the highest loss." The EMA baseline eliminates this bias, steering the agent to learn the "marginal utility" rather than absolute difficulty. This serves as a highly reusable RL reward design trick.
Logit-Normal is counterproductive in distillation: This is a compelling, "counter-intuitive" finding. While the community routinely defaults to reusing the Logit-Normal sampling prior from Rectified Flow training, this work empirically demonstrates via oracle consistency error that the distillation difficulty distribution is U-shaped rather than bell-shaped. This insight itself serves as an important warning to research on consistency distillation.

Limitations & Future Work¶

Bound by the teacher's upper limit: The authors candidly acknowledge that the quality ceiling of CACFM is constrained by the teacher's vector field manifold and cannot exceed it—a shared limitation of all distillation approaches.
One-step generation remains challenging: In the extreme scenario of 1-step generation, highly curved topological flows cannot be perfectly approximated by a single linear step, still causing minor artifacts. The authors do not provide 1-step experimental results.
Granularity selection of M requires manual tuning: Appendix experiments indicate that \(M=4\) strikes the optimal balance (\(M=3\) is too coarse, while \(M=6\) creates a state space of 720 that hinders efficient exploration). However, this choice is derived from empirical search without theoretical guidelines, potentially requiring search reiteration for larger models or different tasks.
Comparison with more sophisticated RL algorithms remains unexplored: This work strictly employs tabular Q-learning, leaving the performance of deep RL methods like PPO unexplored for larger phase counts (\(M>6\)). As the phase resolution increases, the state-explosion issue of tabular methods will warrant new solutions such as neural network policies.

vs. PCM / TCD / Hyper-SD (Static Consistency Distillation): While these methods aim to improve the matching mechanism (e.g., better partition schemes and more stable training techniques), they all rely on static sampling priors. The core distinction of CACFM is that it treats sampling as a learning problem—the RL agent decides which segment to train on, rather than assuming a fixed difficulty distribution.
vs. DDPO / DPOK (RL for Alignment): These works utilize RL to optimize alignment goals for generated content (such as aesthetic scores or human preferences), whereas CACFM uses RL to optimize the training dynamics themselves (specifically sub-trajectory selection). While both sit inside the broader intersection of RL and generative models, they function at entirely different levels—one targets "what to generate" and the other targets "how to learn generation."
vs. Turbo / Lightning (Adversarial Distillation): These methods employ adversarial losses to enhance single-step or few-step generation. Although CACFM also integrates an adversarial consistency loss, its adversarial component plays an auxiliary role (regularization and stabilization of DMD self-distillation). The primary driver remains the RL geometric probe, as removing RL while keeping only the adversarial loss and DMD results in substantially inferior improvements.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ First to integrate RL into consistency distillation training scheduling. Discovering the U-shaped difficulty distribution and designing adaptive sampling based on it is highly novel—advancing from problem definition to solving strategies rather than presenting incremental modifications to existing frameworks.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Evaluated across two backbones (FLUX + SDXL), multiple step counts (4/8/16), and multiple metrics (FID + HPS + Aesthetic + PickScore). Includes wall-clock efficiency analysis, RL policy visualization (difficulty distributions and curriculum emergence heatmaps), and extensive comparisons of ablation and sampling strategies (Loss-Aware / MAB baseline). The appendix further supplements sensitivity analysis on M and Q-learning convergence analysis.
Writing Quality: ⭐⭐⭐⭐ The core insight (U-shaped difficulty vs. Logit-Normal bell shape) is consistently threaded throughout the text, and the policy visualization in Figure 6 is highly persuasive. The flow-adapted DMD score derivation is highly technical yet theoretically sound. A minor weakness is that FIDs of several baselines vary considerably across different tables (e.g., Turbo's 16-step FID reaches 77 on SDXL) without comprehensive discussion.
Value: ⭐⭐⭐⭐⭐ The concept of "adaptive training scheduling" is highly transferable—any learning process featuring discrete stage choices and dynamically shifting optimization difficulties (such as multi-task learning, curriculum learning, or multi-stage RLHF training) could potentially benefit. The method itself yields a tiny computational footprint, making it highly practical.