Diffusion models and flow-matching models generate images by progressively denoising a noisy sample over a series of timesteps. How that time is allocated across the denoising process, the schedule, has a direct impact on sample quality. Standard schedules are chosen for convenience: cosine curves for DDPM, linear interpolation for flow matching. These choices are model-agnostic, meaning they ignore how well the model actually predicts noise or velocity at different noise levels. A model might struggle in high-noise regimes and sail through low-noise ones, but the standard schedule allocates equal time to both. This paper proposes a way to fix that by combining a model-dependent prediction risk with a transport-motivated kinetic objective, yielding schedules that allocate more time where the model needs it most.
Why Model-Agnostic Schedules Leave Performance on the Table
The standard approach to schedule design for diffusion and flow-matching models minimizes a coefficient-path kinetic action, the integral of the squared speed of the signal and noise coefficients along their trajectory. This objective is motivated by optimal transport: minimizing transport cost across a probability path. The result is clean closed-form schedules like cosine for VP/DDPM and the Cond-OT linear parameterization for flow matching. These are strong baselines, but they are model-agnostic. They optimize the geometry of the coefficient path without regard to how well the trained model actually performs at each point along that path.
The fiberwise prediction risk introduced in this paper captures something different. At a fixed time t and state x along the probability path, the true signal/noise decomposition and the predictor's decomposition both sum to x, but they differ in how they split it. All compatible decompositions at a given (t, x) form an affine fiber, a flat surface of possible splits. The fiberwise risk measures the optimal-transport cost between the true decomposition's distribution and the predictor's deterministic decomposition within this fiber. This is base-preserving: it compares decompositions at a fixed state without changing the marginal distribution of the noisy samples. The result is a risk profile R(t) that tells you how badly the model's predictions disagree with reality at each noise level.
The Closed-Form Allocation Rule
The key insight is that on a fixed coefficient curve, combining the kinetic action with a budget on the integrated fiberwise risk yields a closed-form optimal time allocation. The allocation density w*(tau) is proportional to L'(tau) divided by the square root of eta plus lambda times R(tau), where L'(tau) is the coefficient-space speed, R(tau) is the risk profile in kinetic reference coordinates, and lambda is a Lagrange multiplier trading off kinetic cost against risk. In the kinetic reference coordinate where L'(tau) is constant, higher risk means less schedule time. The model spends less time at noise levels where its predictions are poor and more time where they are good.
This construction works for any linear prediction target. For noise prediction, the risk reduces to s_t^2 times the noise-prediction MSE. For v-prediction or velocity prediction, the same fiberwise risk applies with a different coordinate expression. The risk profile can be estimated from a single early checkpoint, making the schedule construction one-shot: estimate risk, compute allocation, retrain with the new schedule.
38.6% FID Reduction on CIFAR-10
The results are consistent across a wide range of settings. On CIFAR-10 with DDPM epsilon-prediction, the model-aware schedule reduces FID by 16.4% relative to the baseline with DPM++3M at 16 function evaluations (NFE), from 9.64 to 8.06. At 32 NFE, the reduction is 11.8% (7.69 to 6.79). At 64 NFE, it is 9.3% (6.83 to 6.19). DDIM at 16 NFE improves by 11.4%, and DPM++2M at 16 NFE by 15.4%. On ImageNet-64, DPM++3M at 16 NFE improves from 25.01 to 23.26 (7.0% reduction), with similar gains at higher NFE.
Flow matching shows even larger improvements. On CIFAR-10, midpoint integration at 16 NFE drops FID from 7.62 to 4.68, a 38.6% relative reduction. At 32 NFE, the reduction is 14.4% (5.52 to 4.72). Heun3 at 15 NFE improves by 32.2% (7.19 to 4.88). On ImageNet-64, midpoint at 16 NFE improves from 41.52 to 30.11 (27.5%), and Heun3 at 15 NFE from 41.38 to 29.04 (29.8%). Euler integration also improves, from 40.29 to 36.50 at 16 NFE (9.4%) and from 37.33 to 30.34 at 64 NFE (18.7%).
The improvements transfer across prediction targets, architectures, training configurations, and risk-estimation checkpoints. V-prediction DDPM improves by 7.4% to 14.7% across all six sampling configurations. U-ViT-S/2, a Vision Transformer backbone, improves by 12.2% at 16 NFE without architecture-specific retuning. The lambda values (220 for DDPM, 450 for flow matching) are selected once via coarse sweeps and held fixed across all subsequent experiments.
Empirical Universality in the Risk Profiles
The most striking finding is that independently trained models across different settings produce nearly identical risk profiles when expressed in kinetic reference coordinates and normalized to unit area. DDPM epsilon-prediction, DDPM v-prediction, and OT-CFM risk profiles align closely. Their allocation deformations, the changes in schedule-time allocation relative to the baseline, also align. More time is allocated near the endpoints (clean data and pure noise) than in the middle.
This alignment is not assumed by the construction. The kinetic baselines determine the reference coordinates independently, and the risk profiles are estimated from separate models. The fact that they agree after normalization suggests something deeper about the relationship between coefficient curves, training dynamics, and fiberwise risk. The authors characterize this as empirical universality across the evaluated models and settings.
The shared shape is well-approximated by sin^{5/4}(pi*tau) after unit-area normalization. Using this frozen analytic shape, the authors construct a template allocation density proportional to (1 + 4*sin^{5/4}(pi*tau))^{-1/2}. This template retains 98.1% of the DDPM epsilon-prediction gains, 98.0% of the DDPM v-prediction gains, and 75.8% of the flow-matching gains on CIFAR-10. On ImageNet-64, the template slightly outperforms the model-aware DDPM schedule (109.7% gain recovery) and recovers 76.2% of the flow-matching gain. The template requires no risk estimation or model-specific fitting.
What the Universality Means and What It Does Not
The authors are careful about scope. The shared profile shape and cross-system agreement are empirically observed, not theoretically derived. An explanation would require understanding how the coefficient curve and training dynamics jointly shape fiberwise risk, which remains an open problem. The pretrained-checkpoint diagnostics on DiT-XL/2 and InstaFlow 2-RF models extend the normalized-risk agreement to larger conditional latent models, but end-to-end evaluation with retraining remains future work. The analytic template distills the model-aware construction rather than replacing it: the full construction identifies the shared shape, maps it to an allocation, and retains additional gains, particularly in flow matching.
The finite-step integration results add nuance. Model-aware schedules improve FID with higher-order integrators (midpoint, Heun3, DPM++ variants) but can worsen FID with Euler integration on CIFAR-10 flow matching. On ImageNet-64, even Euler improves. The degradation is setting-specific rather than systematic, tied to discretization effects at low NFE rather than a fundamental incompatibility.
An empirical Bayes decomposition attributes 93.7% to 96.3% of the integrated fiberwise risk to predictor-dependent excess risk, the component that varies with model quality. The small remaining Bayes-risk component is negatively correlated with total risk, suggesting it represents an irreducible floor that becomes more prominent as the model improves. This decomposition supports the interpretation that the risk profile is primarily driven by model performance, not by the fixed geometry of the coefficient curve.
What This Means for practitioners
If you train diffusion or flow-matching models, the practical takeaway is straightforward. The frozen analytic template can be used as a drop-in schedule without any risk estimation or model-specific fitting, and it retains most of the model-aware improvement. For maximum gains, estimate the risk profile from an early checkpoint and compute the allocation. The one-shot construction means you do not need to iterate: estimate once, train with the resulting schedule, and the risk profile does not need to be re-estimated after the model improves. The lambda parameter trades off kinetic cost against risk, and the coarse sweep that sets it transfers across targets, couplings, datasets, and architectures.
The universality finding is the most intellectually interesting result. If the normalized risk shape is truly universal across model families and settings, it suggests that the allocation problem has a common structure that is independent of specific architectural choices. The analytic template is a first approximation to this structure. Whether it holds for text, audio, video, or 3D generation models remains to be tested, but the pretrained-checkpoint diagnostics on DiT and 2-RF models suggest it extends beyond the unconditional image models used for controlled comparison.