Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schr\"odinger Samplers

2026-05-15

AI Summary

This work addresses the critical dependence of sample quality in flow models and Schrödinger bridge samplers on time discretization under limited inference budgets, where existing approaches typically rely on heuristic or uniform grids. By decoupling the conditional geometry from the marginal evolution of bridge processes through an information-theoretic lens, we propose a novel objective based on the conditional–marginal entropy rate, derive its closed-form expression, and leverage it to construct a training-free non-uniform sampling schedule. Utilizing analytical entropy-rate solutions for Gaussian Brownian bridges, a Heun solver, and entropy estimation techniques, our method achieves substantial performance gains at low NFE: reducing MMD by 18.1% (ODE) and 22.7% (SDE) on 2D tasks, attaining a 5-step FID of 186.3 on CIFAR-10, and outperforming baselines in protein generation on both CAMEO22 and ATLAS benchmarks.

Abstract

For a fixed flow-based generative model under a small inference budget, sample quality can depend strongly on where the sampler spends its few function evaluations. Flow matching and Schr\"odinger bridges define probability paths, yet their inference grids are usually heuristic or inherited from one-endpoint diffusion. We derive a conditional-marginal entropy-rate objective for bridge-aware discretization, separating endpoint-conditioned bridge geometry from marginal flow evolution, and use it to build a training-free entropic inference-time scheduler from first principles. For Gaussian Brownian bridges this rate is closed-form and U-shaped, motivating boundary-heavy nonuniform grids. On trained two-dimensional bridge/flow models, the estimated profile recovers the predicted shape and improves 10-step ODE-Heun MMD over linear by 18.1%, with a paired 22.7% SDE-Heun improvement in the same low-NFE sweep. On EDM/CIFAR-10, the entropic time-discretization gives the best tested five-step FID (186.3 \pm 4.0 versus 200.5 \pm 2.9 for linear and 238.0 \pm 5.3 for cosine). On AlphaFlow protein generation, entropic conditional-marginal (cond-marg) scheduling shows advantage in low-NFE regimes on both CAMEO22 and ATLAS benchmarks. These results support entropy-rate scheduling as a practical low-budget allocation signal for high-dimensional bridge and flow samplers.

Problem

• flow-based generative models
• Schrödinger bridges
• time discretization
• entropy rate
• low-NFE sampling

Innovation

• entropy-rate scheduling
• conditional-marginal discretization
• Schrödinger bridge
• flow matching
• nonuniform time grid