Underdamped Stochastic Optimal Transport with Landauer boundary conditions: machine learning and analytical solution
Authors: Julia Sanders, Dario Lucente, Marco Baldovin, Paolo Muratore-Ginanneschi
Organizations: Department of Mathematics and Statistics, University of Helsinki, Helsinki, 00014 Finland · Department of Physics, Sapienza University of Rome, Rome, 00185, Italy · Institute for Complex Systems (CNR), Rome, 00185 Italy
Finding optimal transport protocols for underdamped stochastic processes is a difficult yet important problem in control theory. In this paper we develop a direct method, based on machine learning, to find the optimal solution of the problem in the presence of non-linear boundary conditions, including those modelling the Landauer bit-erasure problem. The dynamics is steered from a single-peak to a double-peak distribution in position space, by only acting on a controllable external potential. The problem is highly non-linear, and analytical solutions are not available in general. We benchmark our numerical results with a perturbative expansion, which is expected to hold in a suitable overdamped limit.
Figures & tables
Figure 1: Visual summary of the multiscale perturbative scheme adopted to solve the problem in the ε≪1 limit.
Figure 2: Solution of the optimal problem with Gaussian boundary conditions: Ui(q)=q2/2 , Uf(q)=2q2 . Panels (a) - (c) show second order cumulants, while the parameter k(t) of the linear drift [see Eq. ( 17 )] is plotted in (d). Solid lines represent the neural network solution, see Sec. IV . Dashed lines are the analytical benchmark, see Sec. III.1 , obtained by collocation methods. Two values of the weight are considered: M=1 (blue) and M=200 (orange). The solution of the problem with fixed boundary conditions (“full bridge” case, see Ref. [ 50 ] ) is shown in green. We initialize the neural network as −∂U=−q and then select the minimal-cost realization out of 2000 epochs. We set ε=1 .
Figure 3: Solution of the optimal problem with non-Gaussian boundary conditions in the perturbative regime ε=0.2 , for M=8 . Panels (a)-(d) show the evolution of the joint and marginal pdf at 4 different times. Neural network solution (blue) and perturbative prediction (red) are compared. Panel (e) shows the drift evolution. Here Ui=(q−1)4 and Uf=(q2−1)2 . The network is initialized at epoch 0 with the perturbative analytical solution. We train the network for 100 epochs.
Figure 4: Solution of the optimal problem with non-Gaussian boundary conditions in the in the perturbative regime, for large M . Same as Fig. 3 , with M=200 . The network is initialized with the same drift used in Fig. 3 (i.e., the analytical solution with M=8 and ε=0.2 ).
Figure 5: Probability density function at the final time T , for different values of ε and M , computed with the machine learning strategy in the presence of non-Gaussian boundary conditions. Blue lines represent the solution at epoch 100, orange lines the solution at epoch 0 (no training). The green shaded area represents the target distribution. Details as in Figs. 3 and 4 .
Figure 6: Time-dependent variance, for different values of ε and M , computed with the machine learning strategy in the presence of non-Gaussian boundary conditions. Blue lines represent the solution at epoch 100, orange lines the solution at epoch 0 (no training). The green dashed line represents free diffusion. Details as in Figs. 3 and 4 .
Controlling the evolution of a many-body stochastic system from a disordered reference state to a structured target ensemble, characterized empirically through samples, arises naturally in non-equilibrium statistical mechanics and stochastic control. The natural relaxation of such a system - driven by diffusion - runs from the structured target toward the disordered reference. The natural question is then: what is the minimum-work stochastic process that reverses this relaxation, given a pathwise cost functional combining spatial penalties and control effort? Computing this optimal process requires knowledge of trajectories that already sample the target ensemble - precisely the object one is trying to construct. We resolve this by establishing a time-reversal duality: the value function governing the hard backward dynamics satisfies an equivalent forward-in-time HJB equation, whose solution can be read off directly from the tractable forward relaxation trajectories. Via the Cole-Hopf transformation and its associated Feynman-Kac representation, this forward potential is computed as a path-space free energy averaged over these forward trajectories - the same relaxation paths that are easy to simulate - without any backward simulation or knowledge of the target beyond samples. The resulting framework provides a physically interpretable description of stochastic transport in terms of path-space free energy, risk-sensitive control, and spatial cost geometry. We illustrate the theory with numerical examples that visualize the learned value function and the induced controlled diffusions, demonstrating how spatial cost fields shape transport geometry analogously to Fermat's Principle in inhomogeneous media. Our results establish a unifying connection between stochastic optimal control, Schrödinger bridge theory, and non-equilibrium statistical mechanics.
Haiqian Yang, Vishaal Krishnan, Sumit Sinha +1
Department of Mechanical Engineering, MIT, Cambridge, MA 02139 · SEAS, Harvard University, Cambridge, MA 02138 · SEAS, Physics and OEB, Harvard University, Cambridge, MA 02138
We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.
Yesom Park, Eric Gelphman, Stanley Osher +1
Department of Mathematics, University of California, Los Angeles · Department of Applied Mathematics and Statistics, Colorado School of Mines
This paper shows that the semi-dual formulation of the optimal transport problem has a degenerate saddle-point structure, and that its numerical solution is equivalent to solving a constrained optimization problem. We derive necessary and sufficient conditions for the convergence of Monge maps without requiring optimality of the dual potential. This analysis helps explain why, in practice, numerical algorithms often require more iterations to update the transport map than the potential.
Anton Selitskiy, David Millard
Dept. of Electrical and Computer Engineering Stony Brook University Stony Brook, NY 11790, USA · Dept. of Electrical and Computer Engineering University of Rochester Rochester, NY 14627, USA