Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schrödinger Samplers

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ödinger 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.

• The paper introduces a novel discretization strategy that leverages conditional-marginal entropy rates to optimally allocate steps in bridge samplers.
• It employs Hutchinson’s stochastic trace estimator for efficient divergence estimation, resulting in a reusable, training-free grid with endpoint concentration.
• Empirical evaluations on synthetic tasks and high-dimensional settings like CIFAR-10 and protein generation demonstrate notable performance gains over heuristic schedules.

Entropic Discretization: Conditional--Marginal Scheduling for Flow and Schrödinger Bridge Samplers

Overview

The paper "Entropy Across the Bridge: Conditional-Marginal Discretization for Flow and Schrödinger Samplers" (2605.16126) presents a theoretically grounded method for constructing inference-time discretization schedules for flow-based and Schrödinger bridge generative models. The approach explicitly leverages the conditional and marginal entropy rate differences along the generative trajectory, introducing a conditional--marginal (cond--marg) entropy-rate signal as a bridge-aware criterion for discretization. This signal is used to derive a training-free grid allocation scheme, with strong empirical performance in low function evaluation (low-NFE) regimes across synthetic and high-dimensional domains (CIFAR-10 images, AlphaFlow protein generation).

Theoretical Framework: Entropy Rate as a Discretization Signal

The crux of the approach is the recognition that for bridge-based samplers, the conventional use of heuristic time discretization grids (linear, cosine, sigmoid) fails to exploit path-dependent information geometry. In the diffusion context, conditional entropy has guided time scheduling [stancevic2025entropic], but such one-endpoint analyses are insufficient for two-endpoint generative bridges.

The paper formally derives the conditional--marginal entropy identity: $$\frac{d}{dt} H(Z|X_t) = E_{Z,X_t|Z}[\mathrm{div}~v_t(X_t|Z)] - E_{X_t}[\mathrm{div}~\bar{v}_t(X_t)]$$ where $Z$ is the bridge condition (e.g., endpoint pair), $X_t$ is the state at time $t$, $v_t$ is the conditional vector field, and $\bar{v}_t$ is its marginalization. For Brownian bridges, the conditional divergence is closed-form, U-shaped, and singular at the endpoints, motivating nonuniform, boundary-heavy discretization.

Figure 1: Entropy profiles for EDM (left—endpoint concentrated) and AlphaFlow (right—boundary-heavy as predicted by the Brownian-bridge computation), demonstrating the cond--marg rate's configurational distinction.

The method's key point is that the correct discretization signal in bridge models is the contrast between conditional and marginal volume-change rates—not merely the marginal, as for classical diffusions.

Computational Implementation: Estimation and Grid Construction

Practically, the cond--marg rate is estimated using Hutchinson’s stochastic trace estimator for the divergence terms, avoiding the need for full Jacobians in high dimensions. The resulting entropy-rate profile is regularized (e.g., log-transform) for numerical stability and its cumulative distribution is inverted to produce the discretization grid. This grid is constructed once as a calibration step and then reused for generating samples at fixed NFE.

Figure 2: Allocation of time steps by the AlphaFlow cond--marg scheduler; higher grid density at boundaries demonstrates rate-driven node concentration.

Figure 3: Geometry diagnostics for AlphaFlow grids: the log1p-transformed cond--marg grid remains close to linear but retains endpoint bias.

This approach separates three axes that are often confounded: (1) the entropy profile, (2) the numerical discretization, and (3) the resulting sample quality metric.

Empirical Results: Synthetic Probes, Image and Protein Generation

Synthetic 2D Transports

Experiments on low-dimensional synthetic bridges confirm the predicted U-shaped profile and demonstrate the schedule’s sensitivity to endpoint geometry. The largest gains occur at small NFE, where step misallocation most severely degrades the trajectory.

Figure 4: Synthetic (toy) controlled transport scenarios and associated entropy-induced grids, displaying how different geometric endpoints require different step allocations.

Figure 5: Scenario-dependent performance on toy bridges: schedule benefit is endpoint-geometry contingent.

High-Dimensional Generative Models

EDM Image Generation (CIFAR-10)

In the pretrained EDM (Elucidated Diffusion Models) [karras2022edm] setting, the log-tempered cond--marg scheduler achieves the strongest reported five-step FID ($186.3 \pm 4.0$), outperforming linear ($200.5 \pm 2.9$) and cosine ($238.0 \pm 5.3$) schedules at low NFE. With increasing NFE, performance differences converge, consistent with the reduction of integration error.

AlphaFlow Protein Generation

For AlphaFlow [jing2024alphafoldmeetsflowmatching], the cond--marg scheduler generates grids highly concentrated at endpoints (approx. 31.1% of mass in [0,0.1] and [0.9,1]), in line with bridge theory. On CAMEO22 and ATLAS benchmarks, cond--marg scheduling significantly enhances endpoint confidence (plDDT) on medium and large proteins at 5 evaluation steps ($97.62 \pm 0.22$ and $Z$0, leading all baselines), and remains competitive for small proteins. At higher NFE and with raw entropy allocation (unregularized), performance degrades, emphasizing the importance of tempering entropy signals.

Figure 6: Trajectory pLDDT endpoint confidence as a function of dataset size for AlphaFlow, demonstrating cond--marg schedule impact.

Comparative Analysis and Diagnostic Tools

The approach is benchmarked against standard grids (linear, cosine, sigmoid, power) and recent learned or optimized step schedulers. The paper demonstrates that heuristic schedules may incidentally approximate bridge entropy profiles, but only the entropy-driven method explicitly measures and adapts to path-dependent complexity. The boundary concentration ratio (BCR) is introduced as a scalar diagnostic of grid endpoint bias, disentangling geometric allocation from sample quality metrics.

Figure 7: EDM trajectory at 25 steps; as discretizations become fine, schedule choices converge, matching FID results.

Implications, Limitations, and Theoretical Significance

The principal implication is that conditional--marginal entropy rate provides a model-agnostic, diagnostic, and training-free allocation rule for time discretization in generative samplers. This is particularly effective for high-dimensional problems and scientific domains—e.g., proteins, molecules—where endpoint geometry and manifold structure introduce local stiffness and integration challenges missed by generic grids.

The method is constrained by the need for access to both conditional and marginal fields; some pretrained models may only expose one. Hutchinson estimation introduces additional variance. The U-shaped closed-form profile is specific to Gaussian Brownian bridges, and for non-Gaussian or learned bridges, the estimation must be used empirically.

Critical for use in practice, entropy-derived schedules can be combined with existing fast ODE solvers and schedulers, and do not interfere with distillation or model weight changes. The entropy approach thus composes with, rather than replaces, current inference-time acceleration techniques.

Future Perspectives

Directions for future research include aligning model training objectives with the cond--marg rate used at inference, developing schedule adaptivity within few-step distilled samplers, and further biological interpretation—leveraging measured entropy profiles as signals for structural or functional constraints in generated proteins or molecules. The adaptation of the method to non-bridge models (e.g., EDM) also indicates its broader applicability as a general discretization diagnostic.

Conclusion

This work introduces and empirically validates a bridge-aware entropy-rate discretization for flow and Schrödinger bridge generative models. By formalizing and exploiting the conditional--marginal rate, the schedule systematically addresses numerical integration allocation in low-NFE regimes, outperforming or matching heuristic baselines in both synthetic and practical high-dimensional generative settings. The theoretical derivation, estimator design, and empirical validations together provide a robust framework for principled grid construction in generative sampling, with substantial implications for scientific machine learning applications.