Title: HOTA: Hamiltonian framework for Optimal Transport Advection
Venue: NeurIPS 2025
Accepted: no
Keywords: Optimal transport

optimal control

generalized Schrödinger bridge

diffusion models
Abstract: Optimal transport (OT) has become a natural framework for guiding the probability flows. Yet, the majority of recent generative models assume trivial geometry (e.g.,  Euclidean) and rely on strong density-estimation assumptions, yielding trajectories that do not respect the true principles of optimality in the underlying manifold. We present Hamiltonian Optimal Transport Advection (HOTA), a Hamilton–Jacobi–Bellman based method that tackles the dual dynamical OT problem explicitly through Kantorovich potentials, enabling efficient and scalable trajectory optimization.  Our approach effectively evades the need for explicit density modeling, performing even when the cost functionals are non-smooth. Empirically, HOTA outperforms all baselines in standard benchmarks, as well as in custom datasets with non-differentiable costs, both in terms of feasibility and optimality.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a method to learn trajectories between a source and a target distribution under a Lagrangian cost which might not be differentiable. The authors propose to rewrite the dynamic optimal transport problem with such cost using a formulation based on the maximization of a potential function with respect to a value function, which is approximated using a neural network, and under the constraing that it must satisfies the HJB equation. The method is shown to work better than several baselines on different benchmarks, and is shown to be scalable to high-dimensional data, and ablation studies are performed for the different choices of the algorithm.\n\nSTRENGTHS AND WEAKNESSES:\nLearning trajectories between two distribution with respect to a general cost is an important problem which received lots of attention recently. This work tackles this task successfully using a composite approach combining optimal control and RL. Moreover, the method is shown to work on several datasets and compared to different baselines. Nonetheless, I think some elements of the framework could be clarified, and experiments could be improved to real problems\n\n**Strengths**:\n- A method allowing to compute trajectories between two distributions while being optimal relative to pretty general costs.\n- The method is based on an original formulation of the problem, and combines optimal control and RL elements.\n- The method is compared to baselines on standard datasets. \n- Ablations on some choices of the algorithms are provided.\n\n**Weaknesses**:\n- I found that Section 3, where the main theoretical results to understand the method are described, was not very clear.\n- In the description of the method, there are some elements that are not clear. In particular, a target model $\\bar{s}^k$ appears in Section 4 as well as \"EMA\" parameters. But they are not introduced before.\n- The experiments are on toy problems.\n\nQUESTIONS:\nIn Section 3, equation (2) is referred to as the dynamic cost of an OT problem. But this OT problem is not introduced before, and the cost is not in a standard form. This is a bit confusing. Is it the cost of the weak OT problem? Same question for equations (3) and (4). It seems to correspond to the weak OT problem. If it is the case, I do not think Theorem 5.10 of (Villani, 2008) is the right reference for the dual given in equation (3).\n\nThe Hamilton-Jacobi-Bellman equation is not introduced before equation (7), and I do not find it very clear why the function $s(t,x)$ solves it?\n\nWhy OT trajectories can be generated by equation (12)?\n\nWhat are EMA parameters, and target models in Section 4?\n\nThe experiments focus mainly on toy datasets. Even the experiment in higher dimension go only to dimension 700 on data simulated on a sphere. Would it be possible to also do experiments on real data such as images? Are there real case scenario where there would be an interest to use non differentiable potential?\n\nSince the method can be used with general costs, it could also be a nice addition to perform metric learning as in [1,2,3].\n\n\nTypos:\n- In equation (6), therer are 2 parameters call $t$.\n- Line 115: \"and as a sequence\": consequence?\n- Line 116-117: \"Based on dynamic programming approach, define the value function. For\"\n- Equation (11): miss a point\n\n\n[1] Scarvelis, C., & Solomon, J. (2023, May). Riemannian metric learning via optimal transport. In International Conference on Learning Representations.\n\n[2] Pooladian, A. A., Domingo-Enrich, C., Chen, R. T., & Amos, B. (2024). Neural optimal transport with lagrangian costs. arXiv preprint arXiv:2406.00288.\n\n[3] Howard, S., Deligiannidis, G., Rebeschini, P., & Thornton, J. (2024). Differentiable cost-parameterized monge map estimators. arXiv preprint arXiv:2406.08399.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose a method to solve the Generalized Schrödinger Bridge (GSB) problem. Given a source distribution $\\alpha$, a target distribution $\\beta$, and a prior process, the Schrödinger Bridge (SB) problem seek the optimal dirft $v_t(x)$ that takes $\\alpha$ to $\\beta$, where in the dynamical formulation the cost is $||v_t(x)||^2$. GSB also allows potential $U_t(x)$ term in the cost. In contrast to [1], this work do not make any restricting assumption on the potential or on the marginal of the solution conditioned the edges (i.e., t=0 and t=1). This allows them find the optimal solution for larger class of problems but comes in price of higher computation cost. The paper present the dual formulation of the GSB in Theorem 1 and propose to optimize cost Eq. (8) by alleviating the constraint to soft constraint Eq. (9).  Algorithm for the training is presented and the method is validated on a number of small scale problems.\n\n[1] Liu, Guan-Horng, et al. \"Generalized Schr\\\" odinger Bridge Matching.\" arXiv preprint arXiv:2310.02233 (2023).\n\nSTRENGTHS AND WEAKNESSES:\nI will begin by noting that I was not eligible for biding on papers as a reviewer. Although the topic of the paper lies outside my primary research domain, I have made a sincere effort to evaluate it fairly and thoughtfully.\n\n> Strengths:\n1. The general formulation of the proposed method allows to solve a large class of problems.\n\n> Weaknesses:\n1. Though no reference was given, Theorem 1 is provided in the preliminaries section and hence I assume it was proved in previous works. Thus the main contribution of the paper is simply writing the estimator of the terms in the theorem and alleviating the constraint to a soft constraint. I find this contribution of low novelty, since no additional lemmas or theorems were need to write the provided estimator. \n\n2. It seems the provided estimator is not scalable:\n - It requires simulation of the process during training. The author is proposing to use a buffer which potentially would solve this problem, but it is not clear if this is the case and no discussion on the matter is done by the authors.\n- The network parametrize the value function $s(x,t)$, hence the optimization will require second order derivative which tend to be untractable in larger scale.\n- The problem of the high dimension seems relatively easy, i.e., only transporting the mean of the distribution on the sphere.\n\nQUESTIONS:\n1. Is the N-dimensional experiments  has same target and source distribution and just requires to transport the mean of the distribution between the poles? Additionally, why no other baseline is reported for this experiment?\n\n2. Could you please provide a discussion/evaluation of computational cost or training times of your method vs [1]?\n\n[1] Liu, Guan-Horng, et al. \"Generalized Schr\\\" odinger Bridge Matching.\" arXiv preprint arXiv:2310.02233 (2023).\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors formulate and prove equivalence of a dual formulation of the generalized Schrodinger bridge problem. This formulation allows one to avoid explicit representation of the evolving densities and to generate trajectories directly with an Euler-Maruyama approach. Also, it can handle non-smooth potentials and non-differentiable cost functions.\n\nThey then design a method for solution via network parameterization of the value function and an RL approach. Their loss function combines potential terms with a Hamiltonian loss. The network is an MLP with Fourier feature encoding of the time component (with a frequency cap). Empirically, they show that their method outperforms existing methods in terms of matching the target distribution and path optimality.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The problem at hand is an important one and they come up with a formulation that handles non-smooth potentials and non-differentiable cost functions. \n2. The method outperforms the SOTA on some benchmark datasets.\n3. Pseudocode is clearly provided as well as experimental details, so reproduction of their method should be achievable.\n\nWeaknesses:\n1. It might be nice to have a comparison with competing methods on computational cost / training time and the rate of convergence of the results. \n2. Perhaps, the method could be tested and compared on more datasets. There is only a single high-dimensional example, and I could envision more challenging examples in the planar setting (perhaps a not so \"baby\" maze).\n\nQUESTIONS:\n1. Is there any potential for getting stuck in local minima by the initiation of the approximate flow region as linear interpolation? I.e., is it possible for the trajectories to be attracted to a higher cost path that happens to be closer to this initialization, rather than finding a more distant lower cost path?\n2. Can you be more specific on the potential and smoothness parameter for the Sphere dataset? This seemed a bit vague unless I missed something.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for solving the Schrodinger Bridge problem with dynamical state costs by minimizing (i) the violation of the Hamiltion-Jacobi equation along a trajectory and (ii) relevant boundary terms.   The authors propose several techniques for stabilizing optimization dynamics, such as separating the time- and space- derivative regression targets using an EMA network and scaling the contribution of gradients arising from (i) and (ii).   These choices are ablated on several two-dimensional datasets.\n\nSTRENGTHS AND WEAKNESSES:\n** Strengths ** \nThe authors provide a nice analysis of the c-duality associated with Lagrangian OT costs.   Several clever tricks for optimization are presented, and it is interesting that the method works as is for non-differentiable costs.\n\n** Weaknesses ** \nThe distinction between unregularized OT and regularized OT is blurred in this work, where we invoke the unregularized dual potential perspective for optimizing over SDEs.    It seems this could be addressed by reparameterization $u_t = v_t - \\frac{\\sigma_t^2}{2} \\nabla \\log p_t$ (e.g. [1] Ex. 4.4) and adding an additional term into the cost $U(x)$ involving $p_t(x)$.   However, this also suggests the proposed method may be equally valid for deterministic transport maps.\n\nExperiments are fairly limited, focused on small two-dimensional examples.   I did not fully understand the Spheres dataset example and more detail on the dataset would be appreciated.   Are there complications arising from working on a manifold?  \n\nUltimately, I am unsure of the significance of this work due to limited experimental evaluation.\n\n[1] Neklyudov et. al 2024 \"Wasserstein Lagrangian Flows\"\n\nQUESTIONS:\nI would be interested in an ablation which removes the EMA splitting of the objective, with all terms involving $s_\\theta$ optimized together.\n\nI would also be interested in acceleration and replay buffer ablations for the BabyMaze dataset.   In principle, it seems odd to introduce (i) straight-path inductive bias and consider (ii) fixed linear interpolation sampling (buffer-free), since the main feature and challenge of problems with state-cost is that the straight paths are *not* the solution.     I'd like to see the authors discuss these questions in the context of Table 2 / Figure 3 ablations, perhaps recommending that an expectation of \"reasonably straight paths\" may encourage high $\\lambda_a$, for example.\n\n\n\n\n\n**Minor Comments**\n\nPlaces where notation might be cleaned up? \n\n- lines 89-92 ($\\alpha = \\delta_x$?   where does $\\beta$ appear?)\n\n- $\\mu_{v}(\\cdot | x)$ could be written as a conditional distribution / map induced by the SDE with drift $v$?    $\\mathbb{E}_{y \\sim \\mu(x)}[c(x,\\mu)]$ is confusing notation, and  $c(x,\\mu)$ in Eq. 2 depends on the entire trajectory, not just the static map.   \n\n- Eq 6 $\\inf$ over $x_t$ appears to be incorrect\n\n- different notation should be chosen for the loss-gradient scaling term $\\alpha$ vs. initial marginal distribution $\\alpha$\n\nWhat is NLSB as a baseline?    I thought it might be Koshizuka & Sato 2024, \"Neural Lagrangian Schrödinger Bridge\" which might merit citation.   Pooladian et. al 2024 solves the deterministic problem.\n\n\nIn Table 2, what does \"gray values correspond to the method's divergence\" mean?   It seems these (lower) values are not to be trusted?\n\nLIMITATIONS:\nThe authors discuss sensitivity to the Fourier feature encoding of time, and provide ablation results suggesting the necessity of design choices such as adaptive gradient scaling and replay buffers based on simulated trajectories.      \n\nHowever, compared to diffusion-based models or GSBM, which parameterize the drift function directly as a neural network, HJB-based methods require a scalar neural network $s(t,x) \\in \\mathbb{R}$ such that we can take derivatives $\\partial_t s, \\nabla_x s, \\Delta_x s$, adding computational cost.    \n\nWhat parameterization was used for the scalar function?   This may also constitute an important design choice.   See discussion in Salimans & Ho 2021 \"Should EBMs Model the Energy or the Score?\", Neklyudov et. al 2023 \"Action Matching\" and 2024 \"Wasserstein Lagrangian Flows\", Thornton et. al 2025 \"Composition and Control with Distilled Energy Diffusion Models and SMC\".\n\nI cannot find the theoretical guarantees claimed in line 48, although none immediately come to mind as feasible and relevant.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-05"}