Optimal Acceleration Transport

• OAT-FM is a framework that generalizes classical optimal transport by introducing acceleration-based dynamics, resulting in constant mass geodesics with minimal flow bending.
• The methodology employs a natural quadratic metric and Hamiltonian formulation to couple spatial transformations with mass variations, providing precise control in unbalanced transport scenarios.
• The approach significantly advances applications in quantum control, classical optimal control, and generative modeling by reducing transport costs and ensuring smoother interpolations.

Optimal Acceleration Transport (OAT-FM) is a framework that generalizes classical optimal transport (OT) by lifting the paradigm from velocity-based transport to acceleration-based transport. It provides a theoretical and algorithmic foundation for controlling and interpolating between probability distributions and physical states in both geometric, quantum, and generative modeling contexts. OAT-FM is characterized by formulating transport in the product space of samples and velocities, introducing a natural action functional based on acceleration, and yielding geodesics whose defining property is constant mass acceleration or, in a statistical transport setting, minimal aggregate bending of flows. The approach has significant implications for unbalanced OT, quantum control, and flow-matching generative models.

1. Geometric and Mathematical Formulation

OAT-FM is constructed by equipping the conical extension of the diffeomorphism group $\mathrm{Diff}(M)$ of a compact manifold $M$ (with reference volume $\mu$) with a natural quadratic metric. The extended space is $R = \mathrm{Diff}(M) \times \mathbb{R}_+$, representing both spatial transformations and variable total mass. The principal-bundle projection

$$\pi: \mathrm{Diff}(M) \times \mathbb{R}_+ \to \mathrm{Vol}, \quad \pi(\varphi, m) = m\,\varphi_*\mu$$

descends this geometry to the space of non-normalized volume forms. The metric, in coordinates $(\varphi, m)$ or $(\varphi, r)$ with $m = r^2$, is given by

$$\mathcal{G}_{(\varphi, m)}\big((\dot\varphi, \dot m), (\dot\varphi, \dot m)\big) = m\int_M |\dot\varphi(x)|^2\,\mu(dx) + \frac{\dot m^2}{m}$$

which underpins the sub-Riemannian geometry of unbalanced transport (Khesin et al., 2023).

The horizontal lift condition aligns tangent vectors to gradients ($v = \nabla\theta$) with mass change parameterized by $M$0, resulting in the unique decomposition for any tangent vector $M$1 at $M$2: $M$3 inducing the metric

$M$4

2. Hamiltonian and Dynamical Structure

The OAT-FM framework admits a Hamiltonian formulation on $M$5, with the dual variable $M$6: $M$7 The Hamiltonian equations,

$M$8

reflect a continuity equation with global mass source $M$9 and a Hamilton–Jacobi law for $\mu$0. Crucially, the equation for the total mass $\mu$1 yields

$\mu$2

i.e., the total transported mass evolves with constant acceleration along any OAT-FM geodesic, an essential feature absent in balanced OT (Khesin et al., 2023).

3. Cost Functional and Geodesics

OAT-FM defines the unbalanced-OT distance as the minimal dynamical action

$\mu$3

subject to the unbalanced continuity constraint

$\mu$4

Minimizers trace OAT-FM geodesics, which couple spatial transport with a uniform global mass acceleration. In the balanced OT (Wasserstein-2), the mass source $\mu$5 vanishes, and standard velocity-square minimization is recovered (Khesin et al., 2023).

In finite dimensions, the OAT-FM geometry realizes as a “cone-over-Wasserstein” structure on the manifold of Gaussian measures, parameterized by covariance $\mu$6 and mass $\mu$7, with induced metric and ODE geodesics sharing the constant $\mu$8 property.

4. OAT-FM in Quantum and Classical Control

OAT-FM extends to quantum transport where wavepackets in shallow or anharmonic traps are steered in acceleration to maximize fidelity even in dissipative, non-adiabatic regimes (Chakrabarti et al., 26 Jun 2025). The fundamental system considers a particle in a time-dependent potential minimum $\mu$9, coupled to a non-Markovian bath. The optimal control variable is the trap acceleration $R = \mathrm{Diff}(M) \times \mathbb{R}_+$0. The central objective is to maximize the Loschmidt-echo fidelity $R = \mathrm{Diff}(M) \times \mathbb{R}_+$1, with $R = \mathrm{Diff}(M) \times \mathbb{R}_+$2 integrating both non-adiabatic leakage and bath-induced dissipation, weighted by the entire trajectory of acceleration and velocity.

The variational problem leads to an integro-differential Euler–Lagrange equation for $R = \mathrm{Diff}(M) \times \mathbb{R}_+$3 or $R = \mathrm{Diff}(M) \times \mathbb{R}_+$4: $R = \mathrm{Diff}(M) \times \mathbb{R}_+$5 This approach generically outperforms shortcut-to-adiabaticity (STA) strategies, especially in regimes where transport occurs faster than bath excitation propagation (“supersonic” transfer).

The frequency modulation (FM) extension introduces a phase $R = \mathrm{Diff}(M) \times \mathbb{R}_+$6 in the EL kernel, allowing for simultaneous amplitude and frequency modulation of the bath response, and achieving improved transport performance for quantum systems (Chakrabarti et al., 26 Jun 2025).

5. OAT-FM in Flow Matching and Generative Modeling

OAT-FM forms the theoretical foundation for improved flow matching (FM) in generative modeling (Yue et al., 29 Sep 2025). Standard FM learns a velocity field $R = \mathrm{Diff}(M) \times \mathbb{R}_+$7 driving a distribution from noise $R = \mathrm{Diff}(M) \times \mathbb{R}_+$8 to data $R = \mathrm{Diff}(M) \times \mathbb{R}_+$9 under a continuity equation. OT-based FM (OT-CFM) corresponds to Benamou–Brenier transport, regressing to constant velocity interpolations.

OAT-FM generalizes FM by working in the joint $$\pi: \mathrm{Diff}(M) \times \mathbb{R}_+ \to \mathrm{Vol}, \quad \pi(\varphi, m) = m\,\varphi_*\mu$$0 space, introducing second-order Vlasov dynamics, and minimizing the action functional of acceleration squared: $$\pi: \mathrm{Diff}(M) \times \mathbb{R}_+ \to \mathrm{Vol}, \quad \pi(\varphi, m) = m\,\varphi_*\mu$$1 subject to

$$\pi: \mathrm{Diff}(M) \times \mathbb{R}_+ \to \mathrm{Vol}, \quad \pi(\varphi, m) = m\,\varphi_*\mu$$2

with boundary conditions $$\pi: \mathrm{Diff}(M) \times \mathbb{R}_+ \to \mathrm{Vol}, \quad \pi(\varphi, m) = m\,\varphi_*\mu$$3.

The OAT-FM loss function includes carefully weighted proxies for acceleration, enforcing both velocity alignment and minimal bending (straightness) of characteristic trajectories. The two-phase paradigm refines a base FM model via OAT-FM, using mini-batch Sinkhorn OT for endpoint coupling and gradient steps to minimize the second-order action. This reduces transport cost (e.g., Wasserstein-2 loss), straightens flows, and consistently improves generative model quality on a suite of tasks, including CIFAR-10 and ImageNet (Yue et al., 29 Sep 2025).

6. Comparative Properties and Limits

OAT-FM admits a precise relation to classical OT. In balanced OT, the objective enforces constant velocity along geodesics; OAT-FM generalizes this to geodesics of constant acceleration. The extra scalar degree of freedom, associated with mass creation/destruction or velocity endpoint alignment, produces non-trivial geodesic curves that can straighten flows and reduce transport costs beyond OT (Khesin et al., 2023, Yue et al., 29 Sep 2025).

Theoretical analysis establishes that OAT-FM’s proxy loss lower-bounds the true second-order action, and that straightness of the resulting cubic interpolants is both necessary and sufficient for global optimality (i.e., minimal action). In the practical regime, OAT-FM yields monotonic improvements when combined with FM and OT-CFM, controlling both endpoint velocities and alignment of the flow (Yue et al., 29 Sep 2025).

7. Applications and Implementation

OAT-FM is broadly applicable:

• Quantum Transport: Protects quantum wavepacket transport against non-adiabatic and dissipative losses in engineered potentials (e.g., impurity transport in Bose–Einstein condensates), outperforming STA/CDF in finite-resource and fast regimes (Chakrabarti et al., 26 Jun 2025).
• Optimal Control: Defines time-optimal transport protocols for classical oscillators, yielding bang–bang, multi-switch acceleration laws and accommodating frequency modulation for minimum-time transfer (Hegerfeldt, 2023).
• Generative Modeling: Refines generative models (diffusion, FM, OT-CFM, EDM) via a lightweight plug-in procedure, improving sample quality, reducing path energy, and regularizing flow straightness, as demonstrated empirically on low-dimensional benchmarks and large-scale image datasets (Yue et al., 29 Sep 2025).

Implementation leverages standard OT solvers (Sinkhorn), velocity parameterization, and acceleration proxy losses. Best practices recommend small entropic regularization for exact OT coupling, large mini-batch sizes to stabilize Sinkhorn, and EMA endpoints for gradient propagation.

OAT-FM thus integrates geometric, control-theoretic, and data-driven methodologies, enabling principled acceleration-based transport and interpolation across diverse domains.