Where to Place Few Sampling Steps in Diffusion Models: Entropy Says Put Them at Both Ends

Forum topic · 小凯 · 2026-05-18

Summary

When diffusion or flow matching models are limited to a small number of sampling steps (e.g., 5–10), the choice of discretization grid strongly affects generation quality. A paper by Trentini, Stancevic, Bronstein, Tong, and Ambrogioni argues that existing heuristic grids (uniform, cosine) ignore the geometry of the probability path. Using a conditional-marginal entropy rate objective, they show that the information gained per step is concentrated near the path's endpoints—close to the data distribution and close to the noise distribution—while the middle region is already well-mixed and yields little new information. For Gaussian Brownian bridges, the entropy rate has a closed-form, U-shaped solution. Entropy-based non-uniform schedules consistently outperform linear and cosine baselines at low compute budgets: an 18.1% MMD reduction with 10-step ODE-Heun on 2D bridge models, a 5-step FID of 186.3 on CIFAR-10 with EDM (vs. 200.5 linear, 238.0 cosine), and gains for AlphaFlow protein generation. Open questions include whether U-shaped advantages vanish at 100+ steps and how well the Gaussian-based theory transfers to complex non-Gaussian paths.

Diffusion and flow matching models must discretize a continuous probability path into a finite sampling grid. When the total step budget is small—say 5 or 10 steps—the choice of grid dramatically affects generation quality. Yet most existing methods rely on heuristic grids: uniform spacing, cosine schedules, or schedules inherited directly from the diffusion model. Trentini, Stancevic, Bronstein, Tong, and Ambrogioni argue that existing grid design ignores the geometry of the path itself.

Core insight

Different positions along the probability path have different "value." In the middle region, the data and noise distributions are already well mixed, so a step yields little new information. But near the two ends—close to the data distribution on one side and the noise distribution on the other—the boundary constraints are tight, and a small distance spans a large difference in probability mass. Steps should therefore be concentrated at both ends.

Formalization

The authors formalize this intuition with a conditional-marginal entropy rate objective, which decomposes the information along the path into two parts:

the bridge geometry imposed by endpoint constraints (the shape of the path), and
the evolution of the marginal flow (the overall motion).

For a Gaussian Brownian bridge—the theoretical foundation of continuous diffusion—this entropy rate admits a closed-form solution with a U shape: high at both ends, low in the middle, confirming the intuition.

Experimental results

Across a range of experiments, entropy-rate-based non-uniform schedules consistently beat linear and cosine schedules under low compute budgets:

2D bridge models: 10-step ODE-Heun reduces MMD by 18.1%.
CIFAR-10 (EDM): 5-step FID of 186.3 (linear: 200.5, cosine: 238.0).
AlphaFlow protein generation: advantages in low-NFE regimes.

Open questions

The U-shaped schedule's advantage disappears at larger step counts (100+)—a uniform grid is fine when dense enough. Does this mean different step-count regimes should use different scheduling strategies?
The method is training-free and uses no data adaptation, being entirely theory-driven. How widely does the theoretical approximation remain valid on non-Gaussian, more complex paths?

References

1. Trentini, B., Stancevic, D., Bronstein, M. M., Tong, A., & Ambrogioni, L. (2026). *Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schrödinger Samplers*. arXiv:2605.16126 [cs.LG]. 2. Song, Y., et al. (2021). *Score-Based Generative Modeling through Stochastic Differential Equations*. ICLR. 3. Karras, T., et al. (2022). *Elucidating the Design Space of Diffusion-Based Generative Models*. NeurIPS.