Adaptive Discretization for Consistency Models

AI summary

Overview

• Research area: Generative modeling — specifically Consistency Models (CMs), a class of fast samplers derived from the probability flow ODE (PF-ODE) of Diffusion Models. arXiv:2510.17266v1 [cs.LG], 20 Oct 2025.
• Technical level: Intermediate to Advanced. The paper assumes familiarity with diffusion models, the PF-ODE, score estimation, consistency distillation vs. consistency training, FID, NFE, and basic constrained optimization (Lagrange multipliers, Gauss-Newton).
• Scope: A framework (ADCMs) that replaces hand-designed discretization schedules in CM training with an automatic, per-time-step discretization step derived from a constrained optimization problem balancing local and global consistency.
• Authors/affiliations: Jiayu Bai, Robert C. Qiu, Zenan Ling (corresponding, lingzenan@hust.edu.cn) — School of EIC, Huazhong University of Science and Technology; Zhanbo Feng, Zhijie Deng — School of Computer Science, Shanghai Jiao Tong University; Tianqi Hou — Huawei. Code: https://github.com/rainstonee/ADCM.

What This Paper Is About

Consistency Models generate samples in one step by mapping any point on a diffusion trajectory back to the data, and they are trained by forcing the model's outputs at two adjacent trajectory points to agree. Which points count as "adjacent" — the discretization step Δt — is chosen by hand in prior work, requiring repeated manual tuning for different noise schedules and datasets, and poorly chosen steps either make training ineffective (step too large) or unstable (step too small). This paper formulates the choice of Δt as an optimization problem solved automatically at each time t during training.

Key Contributions

1. A unified framework for CM discretization. The authors characterize discretization through two quantities: local consistency (the CM training objective on a local interval, which governs trainability) and global consistency (the denoising error of the training target, which governs stability). They turn the selection of Δt into a constrained optimization problem (Eq. 8) and show that prior methods are special cases of it (Remark 3.2).
2. ADCMs: an analytical, adaptive solution. The constrained problem is relaxed with a Lagrange multiplier λ, and the Gauss-Newton method plus a first-order Taylor expansion of the network output yields a closed-form discretization step (Eq. 10), computed via a Jacobian-vector product. λ is intended to be a constant independent of t, and is typically set to λ ≪ 1.
3. Practical training additions. An adaptive weighting function w(t) = 1/L_global to rebalance stability over time, together with a Pseudo-Huber distance metric, giving the overall loss in Eq. 11.
4. Empirical evidence of large efficiency gains. Experiments on unconditional CIFAR-10, class-conditional ImageNet 64×64, Flow Matching, and ImageNet 512×512 show large reductions in training budget with competitive one-step FID, plus roughly 4% additional training time cost under the same number of training epochs.

Main Findings

• CIFAR-10 data efficiency: ADCMs reach a 1-step FID of 3.16 with a training budget of 12.8M images. ECM — described as the most efficient prior CM — requires 51.2M images to reach a FID of 3.60. ADCMs reach a 1-step FID of 2.80 with 76.8M images, whereas iCT requires 409.6M images for comparable performance.
• ImageNet 64×64 data efficiency: ADCMs achieve a 1-step FID of 3.49 with 12.8M images at 2× model size, and 3.04 at 51.2M images. At 1× model size the progression is 5.12 → 4.65 → 4.23 for 12.8M → 25.6M → 51.2M images. The paper states ADCMs surpass iCT-deep (1638.4M images, 2× model, FID 3.25) using only 3% of its training budget.
• Table 1 headline comparison: ADCM (1 NFE) records FID 2.80 on CIFAR-10 and 3.04 on ImageNet 64×64, without additional training costs (entries marked * in the table indicate extra cost such as a larger model or auxiliary model).
• Computational overhead is small: ADCMs add about 4% time cost under the same training epochs, while converging faster and reaching better final performance than other CM approaches (Figure 3a, 3b).
• The learned schedule resembles empirical practice: ADCMs learn discretization trends similar to hand-designed ones without manual adjustment, using finer discretization at smaller t and coarser discretization at larger t, i.e., emphasizing time intervals closer to the data.
• Theory behind the trade-off (Remark 3.1): The global-consistency denoising error is an upper bound on the squared Wasserstein-2 distance between p_data and the distribution generated by the model at time t−Δt, and a lower bound on accumulated error from previous steps.
• Interpretation of the solution (Eq. 10): The optimal step is inversely proportional to the magnitude of the Jacobian, proportional to ||f_θ⁻(x_t) − x_0||₂ (an estimate of global consistency), and proportional to the linear correlation between v and f_θ⁻(x_t) − x_0.
• λ ablation: Figure 3(c) reports the effect of λ on training dynamics on unconditional CIFAR-10 as a trade-off between stability and final quality. The main text states that with small λ CMs converge quickly but final generation quality is relatively poor, whereas large λ makes CMs more unstable and difficult to converge; the figure caption states the reverse emphasis (large λ improves stability but hurts final performance, and too-small λ reduces stability and hinders convergence). The paper's default guidance is λ ≪ 1.
• Adaptability to other DM variants: On Flow Matching, all CMs were trained under a 12.8M image budget, and ADCMs outperform other CMs without manual adjustment (Table 5, results truncated in the provided content).
• Scalability: The paper reports an ImageNet 512×512 experiment using EDM2 as the base latent diffusion model with SD-VAE for encoding/decoding, compared against sCT and ECM; the specific numbers are in Table 5 and the visible text is cut off before reporting them.

Methodology in Plain English

The authors treat the discretization step Δt at each time t as a decision variable and ask what makes a step "good."

• Trainability argument (local consistency): CMs are trained to match outputs at time t and t−Δt. If Δt is tiny, this objective is easy to satisfy. So minimizing the local objective pushes Δt toward zero.
• Stability argument (global consistency): The target output at t−Δt is not a perfect denoiser of x_{t−Δt} back to x_0. The larger that denoising error, the more the training target misleads the model, causing instability. Constraining this error pushes Δt toward larger values.
• Combining them: Because the two forces point in opposite directions, they write a constrained problem (minimize the local objective subject to the global error staying under a threshold δ) and relax it with a Lagrange multiplier λ, giving a single objective: minimize L_local + λ·L_global over Δt. λ ≪ 1 by default, since trainability is treated as more important than stability.
• Getting a closed form: The model output at t−Δt is approximated with a first-order Taylor expansion using v, the directional derivative computed by a Jacobian-vector product. This turns the problem into a least-squares problem with the closed-form solution in Eq. 10, which can be estimated from a single mini-batch — the paper notes only the trend of the step over t is needed, not precise values.
• Training loop (Algorithm 1): Starting at t = T, the method repeatedly computes Δt* and steps t ← t − Δt* until t ≤ ε, producing a time grid T. The network is then trained on that grid, and the grid is refreshed after every m = 25,000 parameter updates, alternating between re-solving for the discretization and updating θ.
• Loss refinements: An adaptive weight w(t) = 1/L_global downweights times where the denoising error (and hence instability) is large, and the Pseudo-Huber metric is used instead of squared L2 to reduce training variance.
• Experimental setup: Unconditional CIFAR-10 initialized from a pretrained EDM diffusion model; class-conditional ImageNet 64×64 initialized from EDM2. Unless stated otherwise, experiments use VE SDE settings, FID for quality, and NFE for speed.

Why This Matters

• For research: It reframes CM discretization from a schedule-design heuristic into an optimization problem with a closed-form, per-time-step solution, and shows that diffusion models, discrete CMs, and continuous CMs sit at extremes of the same framework (λ → ∞, empirical intermediate values, and λ = 0 respectively). This gives a principled language for comparing and hybridizing prior discretization schemes.
• Practical consequence: Training CMs to strong one-step quality with far fewer training images (for example, a 1-step FID of 3.16 at 12.8M images on CIFAR-10 versus 3.60 at 51.2M for ECM) directly reduces compute and data cost, which matters most for groups without large-scale training budgets.

Real-world applications (derived from the paper's context of fast image, video, audio and 3D generation):

• Interactive image and video generation where a single forward pass (1 NFE) is needed for latency-sensitive interfaces.
• Content-creation and design tooling that needs diffusion-quality output on modest hardware.
• Audio and 3D content generation, which the paper lists among diffusion application areas and which share the same sampling-cost bottleneck.
• Lower-cost model retraining and fine-tuning pipelines for new datasets or noise schedules, since the discretization adapts without manual retuning.

Industry relevance: The efficiency claims target the two dominant costs of generative model deployment — training budget (measured here in millions of training images) and inference latency (measured in NFE). A roughly 4% training-time overhead for substantially faster convergence and better one-step quality is a favorable trade for production training pipelines, and the reported ability to carry over to Flow Matching and to ImageNet 512×512 latent diffusion suggests applicability beyond a single architecture.

Future Directions

• Report and extend the high-resolution results: The ImageNet 512×512 comparison against sCT and ECM (Table 5) is only partially visible in the provided content; extending this scaling study to higher resolutions and video is a natural next step.
• Is λ really constant? The paper chooses a single λ independent of t and notes the trade-off it controls; whether a time-dependent λ could improve the stability/performance frontier is left open.
• Unify the theoretical picture on stability: The paper's Remark 3.2 maps continuous CMs to λ = 0 and DMs to λ → ∞, but instability in continuous CMs under general noise schedules remains, in the paper's framing, unresolved for the general setting.
• Approximation fidelity: The closed-form solution relies on a first-order Taylor expansion and single-mini-batch estimation of the expectation; the sensitivity of ADCMs to that approximation quality, and to the manually set bound δ, are open questions.
• Beyond the tested families: The paper claims strong adaptability to advanced DM variants; systematic evaluation across more noise schedules, architectures, and modalities would test that claim.

Target Audience

• Most valuable for: Researchers and engineers working on fast sampling, distillation, or from-scratch training of diffusion and consistency models, and anyone who needs to train these models under tight compute budgets.
• Also useful for: Graduate students studying the theory of PF-ODE-based generative models and constrained optimization applied to deep learning, since the paper makes the trainability/stability trade-off explicit and solvable in closed form.
• Prerequisites: Comfort with diffusion model objectives, the probability flow ODE, score estimation, Jacobian-vector products, FID and NFE metrics, and Lagrange multiplier reasoning.

Authors’ abstract

Consistency Models (CMs) have shown promise for efficient one-step generation. However, most existing CMs rely on manually designed discretization schemes, which can cause repeated adjustments for different noise schedules and datasets. To address this, we propose a unified framework for the automatic and adaptive discretization of CMs, formulating it as an optimization problem with respect to the discretization step. Concretely, during the consistency training process, we propose using local consistency as the optimization objective to ensure trainability by avoiding excessive discretization, and taking global consistency as a constraint to ensure stability by controlling the denoising error in the training target. We establish the trade-off between local and global consistency with a Lagrange multiplier. Building on this framework, we achieve adaptive discretization for CMs using the Gauss-Newton method. We refer to our approach as ADCMs. Experiments demonstrate that ADCMs significantly improve the training efficiency of CMs, achieving superior generative performance with minimal training overhead on both CIFAR-10 and ImageNet. Moreover, ADCMs exhibit strong adaptability to more advanced DM variants. Code is available at https://github.com/rainstonee/ADCM.