Entropy-Controlled Flow Matching

Abstract: Modern vision generators transport a base distribution to data through time-indexed measures, implemented as deterministic flows (ODEs) or stochastic diffusions (SDEs). Despite strong empirical performance, standard flow-matching objectives do not directly control the information geometry of the trajectory, allowing low-entropy bottlenecks that can transiently deplete semantic modes. We propose Entropy-Controlled Flow Matching (ECFM): a constrained variational principle over continuity-equation paths enforcing a global entropy-rate budget d/dt H(mu_t) >= -lambda. ECFM is a convex optimization in Wasserstein space with a KKT/Pontryagin system, and admits a stochastic-control representation equivalent to a Schrodinger bridge with an explicit entropy multiplier. In the pure transport regime, ECFM recovers entropic OT geodesics and Gamma-converges to classical OT as lambda -> 0. We further obtain certificate-style mode-coverage and density-floor guarantees with Lipschitz stability, and construct near-optimal collapse counterexamples for unconstrained flow matching.

• The paper introduces Entropy-Controlled Flow Matching (ECFM) to rigorously control information geometry and prevent mode collapse in generative models.
• It formulates ECFM as a constrained dynamic optimal transport with an entropy-rate bound that enforces continuous, certificate-style non-collapse.
• Key theoretical results include existence, uniqueness, and stability guarantees, with compatibility across ODE- and SDE-based generative frameworks.

Entropy-Controlled Flow Matching: A Rigorous Constraint-Based Approach to Mode Coverage in Continuous Generative Transport

Introduction

This paper develops a mathematically rigorous framework for continuous generative modeling which directly controls trajectory-level information geometry via entropy budgets. The central contribution, Entropy-Controlled Flow Matching (ECFM), is formulated as a constrained dynamic optimal transport in Wasserstein space: trajectories mapping a source to target law (across time) are required to satisfy a lower bound on the entropy dissipation rate, thus controlling the evolution of compressibility and precluding temporal bottlenecks associated with mode collapse. ECFM generalizes and unifies continuity equation-based modeling (ODEs) and diffusion/Schrödinger bridge methods, inducing certificate-style non-collapse properties without the need for architectural or adversarial heuristics for mode preservation.

Motivation and Background

Bulk of modern deep generative models—particularly in vision—employ flow-based or diffusion-based architectures that (explicitly or implicitly) learn continuous trajectories from a simple base law to the data distribution. While score-based and flow-matching schemes (e.g., SDE-based generation, rectified flows) yield high-quality samples, the underlying regression-style objectives provide no inherent controls over the “information geometry” of the full generative path, allowing transient compressions which can deplete semantic mass and lead to structural mode collapse. Optimal transport (OT) theory, especially in its Benamou–Brenier dynamic form, addresses this with minimum-kinetic geodesics, but such paths remain highly brittle in high-dimensional or non-regular settings—often routing mass through low-density corridors or “bottlenecks.” Classic entropic regularization (e.g. Schrödinger bridges) regularizes the path-level geometry, but generative models rarely impose such control directly in their learning objectives.

ECFM integrates these insights, offering a variational constrained transport in which the entropy dissipation rate is upper-bounded throughout the trajectory. The key constraint: $\frac{d}{dt}\mathcal H(\mu_t) \geq -\lambda$ restricts average instantaneous compressibility, preventing arbitrary local collapse and ensuring sufficient dispersion along the path.

Formulation: ECFM as Constrained Optimal Control

Variational Formulation

ECFM is posed as: $$\min_{(\mu,v)\in\mathcal{A}_\lambda} \frac{1}{2}\int_0^T \int \|v(x,t) - u^\star(x,t)\|^2\,d\mu_t(x)\,dt$$ where $(\mu, v)$ solve the continuity equation (CE) linking endpoints, and $\mathcal{A}_\lambda$ enforces the global entropy-rate constraint. The reference drift $u^\star$ can arise from a teacher model, closed-form interpolation, or target field.

KKT/Dual Characterization

The ECFM minimizer satisfies a rigorous KKT system. The Lagrange multiplier $\eta(t)$ acts as a measure-valued pressure, dynamically activating to enforce the entropy constraint. When the unconstrained path is compressible beyond the permitted entropy decay, the dual variable injects a correction in the score-field direction, blocking low-entropy bottleneck formation. The dual admits a Hamiltonian (Pontryagin) structure and coincides in the stochastic setting with the entropy multiplier in Schrödinger bridge problems.

Theoretical Results

Existence, Uniqueness, and Convexity

• Existence is established via direct method in the space of measure-valued trajectories with finite kinetic energy, leveraging convexity and lower semicontinuity.
• Uniqueness holds strictly in the Schrödinger regime (KL-form) and for pure transport (entropic OT geodesics).
• The induced objective is convex in path law, strict for relevant regular data.

Relation to Schrödinger Bridges and Entropic OT

• For a Brownian reference path, ECFM is equivalent to minimization of KL divergence to the reference among all path measures matching endpoint distributions, under the entropy-rate feasibility constraint.
• In the pure transport regime, ECFM coincides with (and thus selects) the unique entropic OT geodesic. As $\lambda \to 0$, Gamma-convergence to classical OT in the Benamou–Brenier form is established.

Practical Enforcement (Primal–Dual Algorithm)

• Enforcement uses a time discretization of the entropy-rate constraint, with primal–dual updates (augmented Lagrangian/FISTA) and empirical estimation of $\dot{\mathcal H}$ via batch divergence (ODE) or Fokker–Planck/Fisher identity (SDE), with stepsizes and dual variable updates to maintain feasibility.
• Feasibility of the entropy budget is certified via lower confidence bounds across the time grid, supporting statistical control over enforcement.

Mode Coverage, Density Floors, and Stability

Quantitative Anti-Collapse Guarantees

• For any chosen semantic mode partitioning $\{A_k\}$, if each endpoint law puts nontrivial mass on all modes, then all intermediate-time marginals $\mu_t$ retain at least $\beta_k$ mass in each mode, where $\beta_k$ is a function of the endpoint mass, entropy budget, and time horizon.
• In addition, interior “density floors” on compact mode cores are established, which propagate even under nontrivial perturbations.

Stability Properties

• Perturbations in endpoint data, in reference drifts, and in the learned velocity field propagate in a controlled Lipschitz manner to the marginals and coverage/density floors.
• The mode-floor guarantees thus provide explicit quantitative robustness to deployment shifts and optimization noise.

Necessity: Collapse Channels Without Entropy Control

• ECFM is provably necessary for certificate-level non-collapse: flow matching without an entropy constraint admits sequences of near-optimal paths exhibiting bottleneck-induced mode depletion, with entropy dropping to $-\infty$ and arbitrarily low mass in at least one mode at some $t$.

Connections and Implications for Vision Generative Models

• ECFM is compatible with both ODE- and SDE-style schemes, integrating directly with existing flow matching, rectified, or score-based generative paradigms.
• It enables provable, certificate-style anti-collapse properties that are model-agnostic and do not depend on empirical SOTA comparisons.
• Compatibility with KL-based (Schrödinger bridge/diffusion) regularization includes strict convexity, uniqueness, and pathwise stability.

Asymptotics and Limit Connections

• Under vanishing entropy budget ($\lambda \to 0$), ECFM solutions Gamma-converge to classical OT geodesics.
• This connects the framework to classical displacement interpolation and recovers strict coverage even in sharp, measure-concentrating limits.

Conclusion

ECFM introduces an explicit, mathematically tractable mechanism for controlling information geometry in continuous-time generative modeling. By enforcing, certifying, and adaptively regularizing entropy dissipation along generative trajectories, ECFM rules out bottleneck channel mode collapse at the constraint level, as opposed to post hoc empirical or adversarial remedies. The duality with Schrödinger bridge models, together with formal guarantees on mode coverage, density, and stability, position ECFM as a foundational tool for robust and certifiable generative modeling in high-dimensional settings.

Implications and Future Directions

The theoretical guarantees for mode coverage and Lipschitz stability open the prospect of deployment-grade certification for generative models, where empirical recall metrics can be replaced (or accompanied) by explicit statistical certificates. As the formulation generalizes beyond vision (to any generative transport with meaningful entropy geometry), ECFM also sets the groundwork for stochastic control–optimal transport hybrid architectures in continuous, high-dimensional modeling. Potential directions include further regularity analysis, practical large-scale optimization schemes with tighter statistical control, and integration with recent neural Schrödinger and diffusion-bridge matchings, as well as applications for adversarial robustness and interpretability.

---

Reference:

"Entropy-Controlled Flow Matching" (2602.22265)