Neural Hamiltonian Diffusions for Modeling Structured Geometric Dynamics
Authors: Sungwoo Park
Venue: NeurIPS 2025 poster ()
Keywords: geometric deep learning
Riemannian diffusion
Hamiltonian diffusion
Primary area: probabilistic_methods

Abstract: We introduce Neural Hamiltonian Diffusion, a unified framework for learning stochastic Hamiltonian dynamics on differentiable manifolds. While Hamiltonian Neural Networks (HNNs) model conservative systems in flat Euclidean space, they fail to account for geometric structure and intrinsic stochasticity. Conversely, diffusion models on Riemannian manifolds offer geometry-aware stochastic modeling but lack physical inductive biases. Our method parameterizes a Hamiltonian with a neural network and defines its dynamics as a stochastic differential equation on a (pseudo-)Riemannian manifold equipped with a Poisson structure. This enables physically consistent modeling of dynamics on curved, periodic, or causally structured spaces. We demonstrate that the proposed geometric dynamics generalizes existing approaches and applies to systems ranging from molecular dynamics to relativistic n-body problems.

Neural Hamiltonian Diffusions for Modeling Structured Geometric Dynamics
Sungwoo Park
Department of Computer Science and Engineering
Korea University
sungwoo_ wungwoo.parkkkorea.ac.. ac.
Abstract
We propose Neural Hamiltonian Diffussion (NHD), a unified framework for learning stochastic Hamiltonian dynamics differentiable manifolds. Unlike conventional Hamiltonian Neural Networks (HNNs), which assume noise-free dynamics in flat Euclidean spaces, approach models stochastic differential equations (SDEs) curved manifolds endowed with both Riemannian metric and Poisson structure. Specifically, we parameterize neural Hamiltonian and define the dynamics via a Stratonovich SDE whose drift is the Poisson vector field lifted horizontally to the orthonormal frame bundle. This construction ensures coordinate-invariant, gauge- consistent dynamics across (pseudo- Miindd--)Rmmaaiaan manifolds, enabling physically plausible modeling systems with geometric constraints, periodicity, or relativis- tic structure. We establish generalization guarantees under curvature-dependee complexity and demonstrate applications across diverse scientific domains, includ- ing toroidal molecular dynamics, quantum spin systems, and relativistic n-body problems in Schwarzschild spacetime.
Introduction
Modeling physical dynamics from data is fundamental challenge in machine learning, with applica- tions ranging from molecular simulations and protein folding [Karplus and McCammon, 2002 Noć
et al. 2020 to planetary motion and gravitational systems with curved spacetime [Rein and Liu 2019
Pretorius 2005]. A central goal this context to learn a dynamical model that not only predicts future states accurately but also reflects the underlying physical laws such symplectic structure, and geometric invariance [Hairer et al., 2006]. However, many existing learning- basedd approaches focus approximating state transitions Euclidean spaces without explicitly encoding the physics or geometry of the system. In practice, physical systems often evolve on non- oon-Huclidean domains. In molecular dynamics, for example, internal coordinates such as dihedral angles naturally reside toroidal or pseudo-Rie Riemannian manifolds [Zhou et al.) 2020 Townsend et al.) [021]. Likewise, in modeling N N-body interactions near massive celestial bodies, the trajectories of particles evolve in strongly curved spacetimes, where the underlying geometry plays crucial role in determining the causal and dynamical structure [Rezzolla and Zanotti 2013). Quantum spin systems modeled on com- pact Lie groups also exhibit inherently curved dynamics due to their non-Euclidean group geometry [Sakurai and Napolitano, 2017]. These systems are inherently geometric, and their governing laws are often described by Hamiltonian mechanics on manifolds equipped with symplectic structures.
Recent advances in Hamiltonian learning including Hamiltonian Neural Networks (HNNs) Grey-
danus et al. 2019 and their variants [Zhong et al.
2020
Chen et al.) 2021 Cranmer et al. 2020
Wang al. 2023 Dierkes al. 2023 Khoo et al.) 2023 have shown that incorporating symplectic
structure can significantly improve generalization and long-term stability. Complementary approaches
such Symplectic ODE-Nets [Zhong al. 2020), Symplectic Recurrent Neural Networks (Chen
et al. 2021 and Symplectic Transformers Finzi et al. 2020 embed symplectic constraints directly
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

into the learning architecture. However, these models typically operate in flat Euclidean phase space and do not generalize to curved configuration spaces or non-canonical geometries. Orthogonally, growing line of research on modeling learnable stochastic dynamics on Riemannian manifolds [De Bortoli et al., 2022, Huang et al., 2022 Park et al., 2022 Mathieu et al., 2023] has enabled geometry-aware stochastic modeling. Yet, these approaches are not physically grounded: they do not enforce Hamiltonian or symplectic structure to preserve physical fidelity. In this context, our contri- bution lies in unifying these previously disconnected pillars. We propose a stochastic Hamiltonian modeling framework that incorporates both the geometric complexity of differentiable manifolds and the structural constraints of Hamiltonian mechanics. We extend Hamiltonian learning to curved, periodic, and causally structured domains with stochasticity. We highlight the following contributions:
Neural Hamiltonian Diffusions on Curved Manifolds. We propose a novel framework that unifies stochastic diffusion processes and Hamiltonian mechanies on general curved spaces. By incorporating gauge consistency into the modeling, we ensure that the learned dynamies remain physically meaningful and independent of local coordinate choices. Our approach respects both the symplectic structure and the intrinsic geometry of the system, enabling faithful simulation of stochastic physical processes beyond flat spaces.
Freometry-Consisee Learning via Frame Bundle Lifts. Instead of learning dynamics directly the base manifold, we lift the formulation to the frame bundle to handle curvature and coordinate- dependence explicitly. This allows the model to represent vector fields in a unified way across locally varying orthonormal frames, ensuring compatibility between overlapping charts. This geometric design provides consistency across varying local frames and improves the physical reliability of the learned vector fields.
Theoretical Guarantees and Empirical Superiority. We establish theoretical generalization bounds that link curvature, network capacity, and frame symmetry, and prove that gauge consistency intrinsically reduces worst-case deviations. Our method achieves superior performance compared to existing approaches across various structured geometric systems, demonstrating the practical benefits of incorporating geometric and physical consistency into learning.
2 Neural Hamiltonian Diffusion
Hamiltonian Dynamics. Let m, == (q:, P:) E R2d denote the canonical position and momentum coordinates in phase space, and let H C (R)d be smooth Hamiltonian function. We briefly recall the canonical formulation of Hamiltonian dynamics in Euclidean phase space, where the system evolves according to smooth Hamiltonian function H via the associated Poisson bracket structure:
d q dt Pt
0 - 0]
= (m, H} == JVH(q, Pr), where J==
The operator {f.g} == V JVg defines the canonical Poisson bracket for any smooth functions f.g E C* (R24), and {:, } denotes the Hamiltonian vector field applied to observables. This formulation is analytically tractable and serves as the foundation of classical conservative dynamics.
However, such Euclidean and deterministic formulations may face fundamental limitations when applied to structured data domain. First, they assume globally flat phase space, making them ill- suited for modeling systems with curved or topologically structured configuration spaces such as those encountered in relativistic, periodic, or molecular settings. Second, real-world physical dynamics are often inherently stochastic, due latent variables, thermal fluctuations, or observational uncertainty, none of which are reflected in the deterministic formulation. To overcome these limitations, we move beyond the classical regime and explore a generalized class of Hamiltonian systems that operate over differentiable manifolds and evolve according to stochastic dynamics. Specifically, we adopt the framework of Hamiltonian diffussion Bismut [198T]. which preserves the structural fidelity of Hamiltonian flows while incorporating both the intrinsic geometry of the underlying space and the probabilistic nature of physical systems. This generalized formulation enables the modeling of structured, curved, and noisy dynamics in principled and physically consistent manner.
Hamiltonian Diffusion on Manifold Throughout, we work on 2d-dimensional symplectic man- ifold (M,w) w) equipped with a Poisson structure i.e., f, } with local coordinates m := (q. p) E
For detailed explanation of the background, we refer the reader to the Appendix A

---

(()
S(U+))
p 1M-UU
M TQ
Hu
U
Hor
U.h
Figure I: Horizontally-Iin -lftted Hamiltonian Diffusion and Gauge Equivariance on Frame Bundle. (Left) A red stochastic trajectory X, evolves on M = '" Q. The learned horizontal vector field G transports orthonormal frame U, to U, that spans T.M. S its connection-drivee rotation is highlighted by the blue arrow (U h)., realizing the symplectic structure
-bundle. (Right) We visualize specific types equivariant Hamiltonian vector fields along with their fiber rotations under O(d). The preserved structure under transformations highlights the gauge equivariance property of our model.
principal
T*Q M, where the configuration space is set to Riemannian manifolds (2.g). Now, we give formal definition of Hamiltonian diffusion on manifolds:
Definition 2.1. Let (S2, = =.F.FFFTT be the augmented probability space with 2d-dimensional Brownian motion B, Given the standard non-degenerate symplectic 2-form w = Ei= dg' A dpi, define a M-valued semi-martingale that solves the system of stochastic differential equations:
dX. (m,
dB., 44,Ho] w -HHg Ha H(Ta()), 7p(())) (T* θ), (2)
where Tq Q Q and Tp T - T are canonical projections onto configuration manifold and the fiber, and and à denote the interior product and the exterior derivative on M, respectively.
A neural Hamiltonian diffusion (NHD) refers to stochastic process X, evolving manifold under the Hamiltonian flow, where the Hamiltonian Hg is modeled by neural networks with parameters 0 E θ to reflect the induced physical structure of the system. As can be seen, the Poisson bracket formulation naturally leads to geometry-awarr dynamics, ensuring that the induced flow respects the curvature and structure of the underlying manifold. To be more specific, the vector field on the manifold takes the form
G Ga) V V HoH.o...). [---(0)
(m. Ho}
Jo Ho
where (q.p) T and G-)(g) is the inverse of the Riemannian metric the configuration manifold Q. The resulting vector field retains the skew-symmetric structure of Hamiltonian flows while encoding local geometric information through the metric G == [gt]] and can be viewed as geometry-awarg generalization of Euclidean Hamiltonian vector fields in Eq
In deterministic Hamiltonian systems in Eq ① energy conservation is encoded by the identity H = 0, which holds along every trajectory, ensuring exact invariance of the Hamiltonian over time. This reflects the fact that the energy function remains constant along deterministic flows. In contrast, stochastic Hamiltonian framework characterizes energy conservation through the stationarity of the equilibrium distribution, given by Lax O, where π e -x),,,, Rather than preserving energy along individual sample paths, the Hamiltonian in this case governs the long-term statistical behavior of the system via the generator Co. Table summarizes the distinction between these two paradigms.
Horizontal Lift of Hamiltonian Diffusion. One major difficulty in the simulation of Eq. 2 lies in the absence of canonical coordinates, as well as the lack of principled method to define stochasticity on manifolds. Recently, (De Bortoli et al. 2022] suggested geodesic random walks (GRWs) which hamness the property of extrinsic geometry by using Riemannian exponential maps. Yet, there are open questions to respect Hamiltonian and symplectic structures by using GRWs. In contrast to extrinsic approaches, we formulate the intrinsic geometry by lifting the process to the frame bundle (Hsu

---

2002], where geometry consistent noise can be defined, naturally allowing for a principled realization of stochastic Hamiltonian dynamics on manifolds Lázaro-Cami and Ortegal [2008]
Formally, we introduce a horizontal lift (i.e., U²) of the diffusion process (i.e. X₂) to the frame bundle (M), where the stochastic dynamics admit local frame coordinates adapted to the manifold:
Proposition 2.2 (Horizontal Hamiltonian Diffusion). Let U, E ((M) be the horizontal lift of the diffiusion process X TUU)), where и : ((M) - M is the canonical projection and m denotes local coordinate function on M. The lifted process U, evolves according to the Stratonovich SDE:
dU, == Gloor (U,) o dB,, T(U,) X,, E
(4)
GB
(U+) = (1m, Ho (X), (Lm
((m, ((x))))(x))))) (TM)) Ves
(5)
EHy: Horizontal
=Go
= E Vu Vertical Fiber
") is the indx-llwered connection tensor (i.e., Christoffel symbol), and (E) E M(d?, 1) is the vectorized local orthonormal frame. V and Ve denote the vectorized gradients with respect to the configuration point and the frame coordinates, respectively.
where [rM M(2d x d
Here, the horizontal process U, E (M) augments the manifold trajectory with a local orthonormal frame, enabling noise to be defined in the tangent space and lifted via the horizontal distribution. In the formulation to ensure horizontal transport, we remove the vertical fiber part using the connection one-form uB The resulting vector field G (U.) E HU, is obtained by projecting Go onto the horizontal distribution Cf (Id - w²) Go. which ensures that the Stratonovich increment lies in Hu = ker w. Figure [ [cchmatically illustrates this construction. Note that we lift the base metric g to Sasaki-type metric gMM to induce geometric structure on the total space M = T Q. The resulting connection tensor D-M serves as a geometric object defined on the total manifold of trajectory X₁:
Definition 2.3 (Lifted Metric on the Total Space). Lel Q be a configuration manifold equipped with Riemannian metric with its total space M := T Q. We define Sasaki-type metric on M as g.d = n @ g where rag is a pull-back metric with respect प.. Then, the connection lensor associated with gM take the following block tensor form: form:
T
1/2 (0 + Banggt)
M be =
jk
Va,b,cE (L...., 2d].
(6)
Oaxd
Having the complete definition of Hamiltonian diffusion on M by using horizontal lifts in Proposi- tion we now shift our focus from defining the model learning it from physical data.
Learning Geometric Dynamics. Main feature of modeling physical systems is that the system is required to reconstruct both positional and complementary information such velocity and momentum. This additional physical information is represented as trajectories evolving in the second- order tangent bundle of the phase space, formally modeled % (qt. Pt, q p.) E T(T* Q). Given observational trajectory sampled from a data path distribution f/2 ~ PP..daa on physical data space, we seek to recover the underlying Hamiltonian Hg that approximately generates the observed dynamics. The objective is to align the empirical drift of the process with the vector field induced by the learned Hamiltonian. To this end, we define the training loss over path distribution:
Definition 2.4 (Hamiltonian Learning). Let it Pt, qt. (t. E T(T Q) be a sample trajectory
on the physical data space, and let 7t ~ Pr,ant denote the corresponding data distribution. Given a parameterized Hamiltonian Hg T' + IR, the geometric learning objective is defined as follows:
[| p.I (m, Ho} r(U.)) gM
C(0) := E.-Ujo.T..
+ dellu Ta o (U,))o) Uo
(7)
where I1-llgm is Riemannian norm, and dq denotes a distance function on configuration space.
2 Abusing notation, we denote both the connection one-form and the symplectic form by w for simplicity (with o is its null-back), the meaning will be unambiguous from context.
iv

---

Eq. (7) defines a hybrid loss that combines vector field alignment tern and a trajectory reconstruction term. The first term ensures that the learned Hamiltonian vector field aligns with the empirical time derivatives of the state to reflect the underlying physical dynamics. The second term penalizes deviations between the predicted and observed positions on the configuration manifold. The entire objective is conditioned on fixed initial state Uo x ([qo: Pol), reflecting the initial value nature of Hamiltonian systems where the full trajectory is determined by the initial condition. Together, these objectives encourage the model to learn Hamiltonian function that faithfully captures both the evolution of the system and its observed behavior.
Designing Neural Hamiltonian Functions. To model physically consistent dynamics in curved or structured spaces, we aim to construct Hamiltonian functions that reflect the underlying geometry of the configuration space. As motivated from conventional force-field modeling Salomon- Ferrer et al. [2013], we begin by formulating N-body interaction systems on (pseudo-JRiemannial Riemannian manifolds parameterized by neural networks.
N
Hg
E
'p' +
W
We
+
Vo
(8)
2
q
Kinetic Energy
Single purticle Potential
(iii) Pairwise Interaction
Eq. (8) defines parameterized Hamiltonian function for system of interacting particles, where the joint state (q = {q..... ,9N), ,, pd == (PI..... }N E M evolves on the cotangent space of the joint configuration manifold (MM (g The proposed Hamiltonian consists of three distinct components: (i) the first kinetic energy term represents the kinetic energy of each particle and incorporates geometry-awareness awareness by using the local inverse Riemannian metric, (ii) the second potential term captures single-particle effects through neural potentials that depend on individual states, including temporal, environmental, or local structural influences, and (iii) the third term accounts for pairwise interactions that model spatial dependencies and mutual influence across particles, which are essential for capturing correlated behavior on the manifold. Having established the parameterized Hamiltonian function Hg, we now derive the associated Poisson bracket {-, Ho expressed explicitly in local coordinates on .M as follows:
The
Vool' q²) 1"
aw
aw
m Ho} == J, V Hg = Jg G G-'p" +
QN).
0p 0q
0q
In what follows, we impose a structural constraint on the neural Hamiltonian Hg that respects the gauge symmetry of the underlying frame bundle, leading the formulation of gauge equivariance (MM) to ensure the efficient learning of the parameterized Poisson bracket {m, Ho).
Gauge Equivariance on Frame Bundle. While the choice of frame can be arbitrary, the proposed horizontal lift canonically projects dynamics from the base manifold to its frame bundle a manner that is independent of specific frame parametrization. In this context, gauge equivariance is an essential property for ensuring that learned dynamics remain consistent across arbitrary frame choices and preserve geometric coherence during transport [ohen al. [2019]. This principle also formally be realized through the geometry of the orthonormal frame bundle in our framework.
Let U E O(M) and h E ((d) be orthonormal frame and the rotation defined by (U) := U h, then the orthonormal frame bundle ((M) admits a natural right action of the structure group O(d) = {h GL(d, R) hTT Id}. This action describes local gauge transformations within each fiber, and naturally lifts geometric quantities from base manifold M into an equivariant bundle.
Definition 2.5 (Gauge Equivariance on the Frame Bundle). Let f O(M) + V be map into representation space V R + R x Rd corresponding to position and momentum components in the cotangent bundle. We say that f is gauge equivariant if. for every h E O(d),
p(h) = Ta 0 h
f(U h) = p(h))(()))
E GL(2d),
(9)
where U E (MM) is an orthonormal frame, and U h denotes the right action of h on U.
Figure provides an illustrative example of gauge equivariance defined over the orthonormal frame bundle. The equivariance property reflects that under frame transformation, local momentum vectors

---

E Fd transform covariantly as p + hp, while position vectors E R remain invariant as base coordinates. In the context of Hamiltonian learning, the goal is then to construct the neural Hamiltonian Hg such that the resulting geometry-inuuced drift term C is gauge equivariant under the frame transformation U - U h, which holds if the following condition is satisfied:
{m, Hop ( (U. ()) - p(h Hm, (a)))(U)),
vh E O(d).
(10)
To realize equivariant learning within frame bundle coordinates, we propose the Frame Equivariant Transformer U-Ne This architecture integrates canonicalization by transforming coordinates into the local orthonormal frame via U performing invariant computations, and reconstructing outputs via Up, thereby ensuring gauge-consistent predictions. In Appendix, we provide the algorithm and pseudo-code for sampling frame- -quiiarruivari and neural network architectures.
3 Theoretical Analysis
In this section, we present two theoretical results: a generalization bound linking curvature and model capacity, and a deviation bound showing how gauge equivariance improves stability across frames.
Uniform-in- in-timme Generalization of Hamiltonian Learning. With the objective function posed earlier, a natural question arises: If the model achieves near-perfect trajectory matching, why is Hamiltonian learning still necessary? This question is fundamental, as trajectory matching alone does not guarantee physically meaningful generalization. To analyze this rigorously, we consider the neural network 8*, which exactly reproduces the physical trajectories.
(C1)
0r = arg min de(ăt. x (U+lo)),
θ = Bw (0", R)-
(11)
o5e
Here, the radius R of the ball θ reflects the capacity of neural networks. Unfortunately, although the neural network 8" perfectly fits the physical trajectory in the training phase, it fails to capture the holistic physical information such as velocity and momentum. Proposition [3.1 demonstrates that proposed geometric Hamiltonian learning significantly improves generalization.
Proposition 3.1 (Informal), Let (0 == Law (re (0) be an associated probability measure of model trajectory, and assume the condition (C1) in Eq. I holds, Under the mild regularity conditions of Hamiltonian function, the learned model distribution fails to remain close to the data distribution uniformly over time with high probability:
P sup sup W (P,((), S S exp (-2 11/22 (log R))/ [E(O.T] 0EB
where W = W Q) denotes the squared Wasserstein distance on the physical data space T(T Q), and Ω is a constant depending on geometric and model-specific quantities.
Generalization in geometric Hamiltonian learning hinges on two key factors: the curvature of the configuration manifold (i.e., I) and the network capacity (i.e., Rt.. High curvature (IIT')oo) intensifies stochastic distortion, while large R increases variance. This induces a rrade-ff--eprressvvv models capture complex geometry but risk overfitting under curvature. Trajectory matching alone is insuf- ficient, often neglecting velocity and momentum structure. Our method resolves this by enforcing physically consistent dynamics beyond position-level fitting.
Gauge Equivariance Ensures Smaller Deviations. In the second theoretical finding, Proposition shows that enforcing gauge equivariance not only yields uniformly smaller worst-case Wasserstein deviations between the learned trajectories and the reference geodesic across all admissible frames, but also tightens the resulting generalization bounds by eliminating spurious frame-dependent variance.
Proposition 3.2 (Informal). Let < : [0, T] T be reference physical data represented as geodesic. For any Hamiltonian Hç define the frame-rotated trajectory by X₂ (h) ;= (qt, or (U,-)))) for h E O(d). Then, for arbitrary h' O(d), there exists constants K. C 0 such that the following inequality holds:
(KP 8x(0)) = sup (1F) S sup W(P:(h), < Ce hEO(d)
(12)
heO(d)
where (e) ~ (h9(h) is generated by a gauge-equivariant Hamiltonian function Hg.
AAppendix D contains in-detailed information of model architecture.
vi

---

OFF
True
HDM (Ours)
Figure 2: Visualization of Tody Quantum Spin Dynamies via Hopf Projection. Each subplot shows the spin trajectory of single body on the 3-sphere S projected orthogonal complex planes: (orange) and 22 = z + iw (magenta). (Left) Ground-truth trajectories reveal nonlinear yet phase-coherent dynamics. (Right) Our model (HDM) accurately reproduces the spin geometry across badies.
4 Related Works
Hamiltonian Neural Networks. Hamiltonian Neural Networks (HNNs) [Greydanus et al., 2019
introduced the idea of learning scalar energy function Hq: p) whose gradients define conservative dynamics. Several extensions have since been proposed to improve generality, structure preservation, application-specifie modeling. [Cranmer et al.. 2020) proposed learning Lagrangian dynamics as an alternative to the Hamiltonian formulation. [Chen et al. [2021] introduced symplectic recurrence for better long-term stability. [Wang al. 2023] incorporated symplectic constraints into the learning process. (Simiao et al. 2023 adapted HNNs to rigid-body dynamics with energy-aware formula- tions. [Dirrees al. 2023] focused on automatic symmetry detection and exploitation. (Khoo et al.. 2023 proposed modeling separable Hamiltonians to reflect physical modularity. Unlike prior work constrained Euclidean domains, our model learns Hamiltonian dynamics directly on manifolds, enabling faithful modeling of geometry-aware physical systems.
Neural Diffusion Manifolds. Recent work has extended neural stochastic modeling non- Euclidean spaces, particularly Riemannian manifolds, by incorporating geometric structure into 2022 proposed Riemannian score-based generative modeling on smooth manifolds, generalizing Langevin dynamics and score matching to curved spaces. (Huang et al., 2022 developed Riemannian diffusion models by extending continuous- time stochastic differential equations (SDEs) to arbitrary manifolds. (Park et al.) [022] introduced
diffusion score-based generative models. (De Bortoli
Riemannian Neural SDEs, enabling stochastic representation learning directly on manifolds using
intrinsic geometry. (Lou et al. (003) addressed the scalability of Riemannian diffusion models
for soth--imensionall and complex manifold settings. (Fishman and Cunningham 2023) tackled constrained diffusion in non-Fuclidean domains by incorporating boundary-awar mechanisms. These works lay the foundation for stochastic modeling on manifolds. Building on this line of research, our approach extends geometric diffusion models with a Hamiltonian perspective, enabling structured modeling of physical dynamics on curved spaces.
5 Experiments
Problem Formulation. In this section, validate our proposed framework across three distinct physical scenarios that reflect a diverse range of geometric structures: (i) an interacting spin system evolving on the compact Lie group manifold SU(2) ~ $3, (ii) relativistic N N-body dynamics formu- lated on Lorentzian spacetimes such as the Schwarzschild manifolds and (iii) molecular dynamics of protein backbones represented high nigh-dimensional toroidal configuration spaces TV Each setting highlights a unique combination of curvature, topology, and physical constraints, allowing us to the generality and fidelity of neural Hamiltonian diffusion non-Fuclidean domains. We compare our method with recent state-of-the art methodologies in geometric sequential modeling including GeoTTDM Han al. 2024 Eq Xu al. [2023] EGNN Satorras al.) (2021] SE-3
transformer Fuchs et al. [220 Hamiltonian learning based such as HNN Greydanus et al. [2019]
SympHNN David and Méhats 2023 Noether vaan van der Ouderaa al. 2024 are also considered.
In all scenarios, we formulate physical dynamics prediction as a sequence modeling problem on non-Euclidean manifolds. Let {nx }/_ denote a trajectory of geometric states " E T(T* Q) sampled from a Hamiltonian system. During training, each model is conditioned on a single initial state Yi and trained to autoregressively predict the subsequent states {T At test time, the predicted state re-fed into the model to generate the next one, allowing the model to learn long-range extrapolation dynamics at each step. This setup reflects realistic forecasting settings where long-term evolution must be inferred from geometric observations. A comprehensive summary of the experimental setups is included in the Appendix.
Tag
vii

---

Model
AD-3
2AA
4AA
Spin
Schwarzschild
HNN
1.13/0.777
0.612 / 0.859
2 1.0
0.141 70.275
0.106/ / 002218
Noether
0.554/ 0.644
0.580 / 0.723
1.0
0.077 / 00162
0.053 00116
SympHNN
0.596/ / 0.717
0.519/ 0.736
2 1.0
0.083 / 0.124
0.063 / 0.175
SE
0.312/ 0.513
0.445 0.596
0.649 / 0.830
0.384 / 0.665
0.338 00.437
EGNN
0.251 / 0.501
0.367 / 0.405
0.417 / 0.474
0.182 / 0.242
0.155 0.226
EqMotion
0.081 / 0.117
0.062/ / 0.152
0.131 / 0.174
0.090 / 00100
0.081 / 0.144
GeoTDM
0.045 / 0.102
0.079 / 0.145
0.093 / 0.179
0.037 / 0.085
0.026 0. 0.046
NHD (Ours)
0.023/ / 0.084
0.055 / 0.
0.117
0.019 / 00.097 0.012 / 0.035
Table Comparison of toroidal protein trajectory prediction and curved-space N-body dynamics. The first three columms (AD-3, 2AA, 4AA) report ADE/FDE on protein torsion angle trajectories. The last two columns (Spin and Schwarzschild) report ADE from 2V-body simulations of spin-based and Schwarzschild-metrr particle systems with N = 3 and N particles. The first and second best is highlighted with bold and blue.
Molecular Dynamics of Protein Backbones. In modeling Hamiltonian formulations of protein molecular dynamics, we are motivated by classical force fields used in molecular model- ing Cornell et al. [1995]. Maier et al. (2015). Tian et al. [019]] which incorporate structured physical interactions such as bond stretching, angle bending, torsional rotations, and non-bonded forces. We reinterpret these force fields as learnable neural po- tentials while preserving underlying geometric and physical consistency. The configuration space is set to nigh-dimensional torus Q == Nuugl where Nangle denotes the number of torsional degrees of freedom (e.g.. backbone dihedral angles Φ, vi, w and side-chain angles Xt). To evaluate the proposed framework, we perform experiments on three representative peptide systems of increasing complexity: AD-3 Alanine dipeptide, which exhibits simple dynamics on T 2AA dipeptides, where T arises from torsions, and 4AA tetrapeptides. which form structured dynamics on T due to multiple interacting torsional modes. To extract geometrie Hamiltonian states, we post-process time-aligned atomic trajectories from the Timewarp Klein et al. [2023] to compute angles (6, 16 w, Xi) as generalized coordinates, and approximate their corresponding momenta by estimating the reduced moment of inertia associated with each torsional mode. Time derivatives are computed via finite differences across consecutive frames. We evaluate trajectory quality using standard geodesic metrics on the torus manifold, including average displacement error (ADE) and final displacement error (FDE), where distances are measured along Nad::: = $1 ... x S², As summarized in Table our method outperforms existing benchmarks by a significant margin across all evaluated metrics.
dynatecc
Model presictio
Figure 3: Spatiotemporal Ramachan- dran Map. Torsional state evolution over time compared between true and predicted trajectories.
backbone and occasional side-chain
Quantum Spin System. In quantum physics, a spin system refers a collection of particles, each possessing an intrinsic angular momentum (spin) that interacts with neighboring spins according to specified coupling rules. Mathematically. spin states are modeled as unit vectors on sphere or, as elements of compact Lie groups. We model the dynamics of mutually interacting quantum spins on the unit 3-sphere $ c Rd,, where each spin represented as unit quaternion that evolves under rigid-bode dynamics. The system is equipped with anisotropic inertia and pairwise coupling Hamiltonians, giving rise to nonlinear, geometry-come constrained motion. The total Hamiltonian of the system takes the following form: H(q", .p = = (p. IIp.) - Eiss des (fq, ) where wi E R is the body angular velocity of the i-th spin, I E 3R3x3 is the moment of inertia tensor, and dej is the coupling strength promoting alignment between spins q; and q. The inner product q,) = rızy + yıy; + 222y + wayr measures the similarity of unit quaternions on $3 The time evolution is governed by the Hamiltonian equations şi f[((w)) )q; and W, ~',, where ((w) encodes angular velocity and т; is the coupling torque promoting spin alignment. The induced dynamies are thus constrained to a Riemannian manifold, specifically the 3-sphere endowed with its canonical metric. We visualize the resulting trajectories using Hopf projection :: + S² in Figure 2 where each spin is mapped to complex plane components (21, 22) E C² with 21 z + iy. F2 + iw. |2 + |22 2 1. Both the qualitative trajectories in Figure and the quantitative metrics in Table demonstrate that our dynamics accurately capture the underlying spin system evolution.
viii

---

True
HDM (Ours)
HNN
RM
GM
G=c=1
Figure 5: Three-body Trajectories in the Spacetime of Schwarzschild Black Holes. Left: The ground-truth simulation obtained by numerically integrating the exact relativistic dynamics, shows three mutually interacting bodies (labels 1-3) spiraling toward the event horizon (f.e., Schwarzschild Radius 2M). Center: The proposed method accurately captures the relativistic deflection and inward inspiral of all three trajectories, remaining faithful to the ground-truth. Right: Existing Euclidean HNN trained without explicit geometric conditioning yields inconsistent trajectories that indicate an incorrect physical regime.
Relativistic Particle Dynamics. In the last exper- iment, we consider the dynamics N interacting bodies in the eurved spacetime surrounding com- pact astrophysical objects such as Schwarzschild black hole. In formulation, the background force field is derived from general relativity, encapsulat- ing the relativistic geometry of spacetime. Mean- while, the interaction forces between bodies fol- low classical modeling assumptions, e.g., pairwise Newtonian-lic potentials. This setup allows us to generalize classical N N-body systems to curved spacetime environments beyond the fat spaces [Satoraas et al. [202T]. The Hamiltonian consists of a kinetic term defined via the inverse Schwarzschild metric and classical pairwise potential: H(q ,p*)
-15
Figure Comparison total Hamiltonian H(t) and cumulative relative drift E.
across models.
GM² N En.= 9l" (q')p"p Eiss v - 1e The geodesic structure
of the spacetime introduces oun-Euclidean curvature effects the momentum transport, while inter-body forces remain Newtonian-like. We implement symplectic leapfrog integrator adapted relativistic Hamiltonian flow and simulate multi-body systems initialized near stable orbital radii. The results clearly indicate that Euclidean methods struggle to model particle behavior in curved geometry. In contrast, our proposed HDM achieves superior reconstruction accuracy and substantially lower energy drift in Table T reflecting improved alignment with the intrinsic geometry of the system.
Ablation Study. We assess the numerical stability and seal- ability of our model via two ablation criteria: (i) energy conser- vation, and (ii) robustness across varying system sizes. Figure shows that our method yields significantly lower energy drift compared to Euclidean baselines (i.e., SympHNN), indicating better consistency with the underlying geometric structure. In addition, Table 2 reports how performance degrades as the number of spin particles increases. While both models exhibit reduced accuracy for larger N, the proposed gauge-equivariant model remains consistently more stable. For instance, while both models experience increasing eиror as N grows, the non-equivariant variant exhibits a sharp deterioration at N = 10, with over sevenfold increase ADE relative N = 3. In contrast, the gauge-equivarian model maintains more gradual degradation, reflecting improved scalability under growing system complexity.
N
G-equiv
Non-equiv
0.019
0.024
(+24%)
0.097
0.103
(+6%)
10
0.120
0.173
(+44%)
20
0.148
0.169 (+14%)
Table 2: Performance degradation as the number of spin particles increases.
Conclusion
This work presented Neural Hamiltonian Diffuusion (NHD), a unified framework that integrates geometry-awar diffusion processes with pruuture-preservvngg Hamiltonian learning. We formulated a diffusion process lifted the frame bundle and constructed neural Hamiltonian vector fields that are equivariant under frame transformations. We provided theoretical results characterizing the generalization properties of the proposed method, including uniform-in-time bounds and frame- wise deviation under gauge transformations. Experiments results across diverse scientific domains demonstrated that our NHDD consistently improves physical fidelity and predictive stability compared to Euclidean or non- Hom-amiltonnan baselines.
ix

---

Acknowledgments
This work was supported by ICT Creative Consilience Program through the Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) 2025-RS-2020
References
Jean-Michel Bismut. Mécanique aléatoire, volume 866 of lecture notes in mathematics, 1981.
Zhengdao Chen, Jian Zhang, Martin Arjovsky, Léon Bottou, and Joan Bruna. Symplectic recurrent neural networks. arXiv preprint arXiv:2010.07003, 2021.
Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral cnn. In International conference on Machine learning, pages 1321-1330. PMLR, 2019.
Wendy D Cornell, Piotr Cieplak, Christopher Bayly, Ian Gould, Kenneth M Merz, David M Ferguson, David C Spellmeyer, Thomas Fox, James W Caldwell, and Peter A Kollman. A second generation force field for the simulation of proteins, nucleic acids, and organic molecules. Journal of the American Chemical Society, 117(19)5517--519, 1995.
Miles Cranmer, Alvaro Sanchez- Sanchez-Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Lagrangian neural networks. International Conference on Learning Representations (ICLR), 2020.
Marco David and Florian Méhats. Symplectic learning for hamiltonian neural networks. Journal of Computational Physics, 494: 194:112499 12495, 2023.
Valentin De Bortoli, James Thortoo,, Jeremy Heng, Alexandre Bouchard-Côté, Paul Vicol, and Yee Whye Teh. Riemannian score-based generative modeling, arXiv preprint arXiv:2202.02763, 2022.
Eva Dierkes, Christian Offen, Sina Ober- Blöbaum, and Kathrin Flaßkamp. Hamiltonian neural networks with automatic symmetry detection. arXiv preprint arXiv:2301.072 000928, 2023.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing hamiltonian mechanics with learned symplectic structure. In International Conference on Machine Learning (ICML), pages 3146-3154. PMLR, 2020.
Nic Fishman and John P. Cunningham. Diffusion models for constrained domains. arXiv preprint arX1v:2304.05364, 2023.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformmrs: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 33: 1970-1981, 2020.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, volume 32, 2019.
Ernst Hairer, Christian Lubich, and Gerhard Wanner. Geometric numerical integration: structure- preserving algorithms for ordinary differential equations, volume 31. Springer Science & Business Media, 2006.
Jiaqi Han, Minkai Xu, Aaron Lou, Haotian Ye, and Stefano Ermon. Geometric trajectory diffusion models. arXiv preprint arXiv:2410.1 13027, 2024.
Elton P Hsu. Stochastic analysis on manifolds. Number 38. American Mathematical Soc., 2002.
Chin-Wer Wei Huang, Milad Aghajohari, Joey Bose, Prakash Panangaden, and Aaron Courville. Rieman- nian diffusion models. In Advances Neural Information Processing Systems, 2022.
Martin Karplus and Andrew McCammon. Molecular dynamics simulations of biomolecules. Nature Structural Biology, 9:646-652, 2002.
X

---

Zi- Yu Khoo, Dawen Wu, Jonathan Sze Choong Low, and Stéphane Bressan. Separable hamiltonian neural networks. arXiv preprint arX1v:2309.010669 2023.
Leon Klein, Andrew Foong, Tor Fjelde, Bruno Mlodozeniec, Mare Brockschmidr, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. Advances in Neural Information Processing Systems, 36: 52863-52883, 2023.
Joan-Andreu Lázaro-Camí and Juan-Pablo Ortega. Stochastic hamiltonian dynamical systems. Reports on Mathematical Physics, 61(1):65-12,, 2008.
Aaron Lou, Minkai Xu, Adam Farris, and Stefano Ermon. Scaling riemannian diffusion models. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
James A Maier, Carlos Martinez, Krishna Kasavajhala, Lauren Wickstrom, Kevin E Hauser, and Carlos Simmerling. ff14sb: improving the accuracy of protein side chain and backbone parameters from f199sb. Journal of chemical theory and computation, 11(8):3696-377 3713,, 2015.
Pascal Massart. About the constants in talagrand's 's concentration inequalities for empirical processes. The Annals of Probability, 28(2):863-88) 884, 2000.
Emile Mathieu, Vincent Dutordoir, Michael John Hutchinson, Valentin De Bortoli, Yee Whye Teh, and Richard Turner. Geometric neural diffusion processes. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
Frank Noé, Alexandre Tkatchenko, Klaus-Robert Müller, and Cecilia Clementi. Machine learning for molecular simulation. Annual Review of Physical Chemistry, 71:361-390, 2020.
Bernt Øksendal and Bernt Øksendal. Stochastic differential equations. Springer, 2003.
Sung Woo Park, Seongmin Kim, Junhyun Lee, Seonguk Joo, Jaehyung Choi, and Eunho Yang. Riemannian neural sde: Learning stochastic representations on manifolds. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Frans Pretorius. Evolution of binary black hole spacetimes. Physical Review Letters, 921121, 2005.
Нanno Rein and Shang-Fei Liu. Rebound: an open-source multi-purpose n-body code for collisional dynamics. Monthly Notices of the Royal Astronomical Society, 485(4):5490-5511 2019.
Luciano Rezzolla and Olindo Zanotti. Relativistic hydrodynamics. Oxford University Press, 2013.
J.J. Sakurai and Jim Napolitano. Modern quantum mechanics. Cambridge University Press, 2017.
Romelia Salomon-Ferrer, David A Case, and Ross C Walker. An overview of the amber biomolecular simulation package. Wiley Interdisciplinary Reviews: Computational Molecular Science, 3(2): 198-2107 2013.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. (n) equivariant graph neural networks. In International conference on machine learning, pages 9323-9332. PMLR, 2021.
Fan Simiao, Linlin Huo, Fei Wang, Lei Zheng, and Shuang Yang. Hamiltonian neural network 6-dof rogyd-body dynamic modeling based energy variation estimation. International Journal of Intelligent Systems, 2023. doi: 10.1155/2023/8882781.
Michel Talagrand. The generic chaining: upper and lower bounds of stochastic processes. Springer Science & Business Media, 2005.
Cong Tian, Krishna Kasavajhala, Kaërl Belfon, Lauren Raguette, Hao Huang, Andre Migues, Jessica Bickel, Yumin Wang, Jennifer Pincay, Qi Wu, et al. ffl 9sb: amino-acid-specinf protein backbone parameters trained against quantum mechanies energy surfaces solution. Journal of chemical theory and computation, 16(1):52--552, 2019.
xi

---

James Townsend, Eric Vogeley, Christoph Wehmeyer, and Frank Noé. Geometric constraints in molecular simulations differential geometry. Journal of Chemical Physics, 154(23)233108, 2021.
Tycho F. A. van der Ouderaa, Mark van der Wilk, and Pim De Haan. Noether's razor: Learning conserved quantities. In The eighth Annual Conference on Neural Information Processing Systems, 2024. URL
Yiming Wang, Yujie Sun, Yiqun Liu, and Hao Xu. Symplectic learning for hamiltonian neural networks. Journal of Computational Physics, 484: 12060, 2023.
Jon Wellner et al. Weak convergence and empirical processes: with applications to statistics. Springer Science & Business Media, 2013.
Chenxin Xu, Robby Tan, Yuhong Tan, Siheng Chen, Yu Guang Wang, Xinchao Wang, and Yanfeng Wang. Eqmotion: Equivariant multi-agent motion prediction with invariant interaction reasoning. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 1410-1420, 2023.
Yicheng Zhong, Debasmit Dey, Rudrasis Chakraborty, and Vikas Singh. Symplectic ode-net: Learning hamiltonian dynamics with control. International Conference on Learning Representations (ICLR), 2020.
Yao Zhou, James Zou, Will Grathwohl, and Frank Noé. Torchmd: A deep learning framework for molecular simulations. arXiv preprint arXiv:2012 12 12106, 2020.
xii

---

A Backgrounds
A.1 Stochastic Riemannian Geometry, Hamiltonian Dynamies
Let M be a smooth d-dimensional Riemannian manifold with metric The tangent bundle T M the disjoint union of all tangent spaces T&M for x E M. The orthonormal frame bundle (MM) is principal -nn)-bundle over M, where each point U ((M) is an ordered orthonormal basis (e1,...,Ca Cn) of T at some x E M. The canonical projection π O(M) - M maps a frame
to its base point. This allows lifting trajectories from M to (MM) in geometrically structured
way. connection on (MM) decomposes the tangent space ((M) into vertical and horizontal components:
(MO(M) = Hy + Vu-
The connection form w e (()))) is Lie algebra-values I-form satisfying:
* w(A")= = A for all A E sa(n), where A is the fundamental vertical vector field, Ryw Adg- under the right action Rg of E O()..
vector field X on M lifts horizontally ((M) if (X) O. This horizontal lift allows stochastic processes on M to be lifted into ((M) while preserving the connection structure.
Let X, M be semimartingale the frame bundle such that π U,) = X, We say U, the horizontal lift of X, it satisfies the
smooth manifold M. Let U, E ((M) denote a process
following condition from Hsu 20022:
dU, Σ H.(U.) o dB 1=1
Un =
where {H₂} is the canonical horizomtal vector field associated with the standard basis vectors E R", and B, an -aaluel Brownian motion. This construction allows noise to be defined canonically on R", lifted horizontally via H, and transported on the manifold through U, The projected process X, = (U, then inherits stochastic dynamics that respect the geometry induced by the connection.
An orthonormal frame U E ((M) at x xUU is represented as an isometry U R." + T, i.e. for standard basis vector e; R" Ue: = E T.M. Thus, the process U, evolves according to the geometry of M with stochastic noise applied in the frame coordinates and mapped into the tangent bundle via horizontal transport. This perspective ensures that noise is not only intrinsic the manifold but also compatible with the Levi-Civita connection and geometric constraints of M.
Hamiltonian Vector Fields Manifolds. Let H T Q R be Hamiltonian. The Hamiltonian vector field Xи is defined implicitly via the symplectic form w on T Q, i.e., 'X, = dH. In local Darboux coordinates (q', Ps), Xи takes the standard form:
OH a OH a
X Σ
Op, δqe
Dq' Öp.
This is the generator of deterministic Hamiltonian evolution and provides the basis for its stochastic extension. For smooth functions f.g T + R. the Poisson bracket is defined as:
Of 0g Of 0g Oie Op. 3p. 0q
U.g) = : = Σ
Horizontal Lift and Connection form. In Section 2 and Figure 1 described the evolution of lifted Hamiltonian dynamics on the orthonormal frame bundle ((M) The key construction relies decomposing a lifted vector field into horizontal and vertical components with respect to the principal connection. Let Go (U,) Tu, (M) denote the full lifted vector field constructed from the Hamiltonian flow {m. H. This field is not guaranteed to lie in the horizontal distribution Hu, kerwv, and must be projected to ensure that resulting SDE respects the manifold connection structure. The canonical projection is defined via the connection 1-form E (O()) which satisfies:
wu, (Vu.) so(d),
Wu, WU, (Hu.) 0.
xiii

---

The vertical component is extracted using and the projection onto the horizontal space is given by:
o (U.) Go( - wu. (Go(U.))",
(13)
where wu, (Go)* E Vu, denotes the vertical lift of the Lie algebra element associated to the vertical component, and maps Lie algebra elements to fundamental vector fields. This ensures that the resulting direction Gffor E HU, lies entirely in the horizontal distribution, satisfying the condition (G
In local coordinates used in Eq. (5). the vertical component explicitly appears as the second term involving the connection tensor [ IMP and the vectorized frame vec(E). The subtraction in Eq. (5) therefore realizes the above projection in local form, decomposing the lifted vector field into:
Go G G.ver
(14)
This decomposition plays crucial role in ensuring that the Stratonovich increment dU, Gg (U,) dBB evolves along direction consistent with the geometry of T*Q. This geomet- ric consistency essential for transporting noise on the manifold without introducing spurious curvature-induced distortions, and forms the foundation of the gauge-consistent stochastic Hamilto- nian dynamics defined in this paper. Note that the proof of Proposition 2.2 will deliver the detailed calculation of deriving the vanishing connection 1 1-form to show the validity.
Infinitesimal Generator and Fokker-Planck Equation. The stochastic process X; E TM governed by the Stratonovich SDE proposed in main manuscript
dX.- = {m, (Xa) dBi,
defines diffusion process on the symplectic manifold M with Hamiltonian vector field Ho). The corresponding infinitesimal generator Lg acts on smooth test functions f E - as:
2d Σ
Lof(x)
m, HoY' (xr) a (1m, (x).0.((.)))
2
where the index runs over local coordinates (q'..... g pt pa) on T* and (m, Ho}' denotes the i-th component of the Hamiltonian vector field. Let (x()) E Dens(M) denote the time-dependent probability density function of X, Then, the evolution of Pr is governed by the Fokker-Planck equation associated with the generator Le:
(p.(x) Cop((x). Òt
(15)
where Cj is the formal adjoint of La in the L²(M) sense. Explicitly, using integration by parts, this yields:
0pr(z) Øt Σ a. ((m, Ho)" (x:) 0: (Im, Hoy (x) (x) -a(a)))
where the density Pe allows the Radon-Nikodym derivative with respect probability measure P, by the formula dP,- = p:dx with Lebesgue measure dx. This formulation describes how the probability mass of the stochastic process spreads over the symplectic manifold under the influence of the geometry-awar Hamiltonian noise.
xiv

---

Table 3: Comparison between Euclidean HNNS and Neural Hamiltonian Diffusions (Ours).
Comparison Item
Euclidean HNNs
Neural Hamiltonian Diffusions (Ours)
Space
Red (flat)
General manifold T*Q
Structure
Fixed symplectic form J
Poisson structure from lifted geometry
Noise
None (deterministic)
Intrinsic stochasticity (via horizontal lift)
Energy Conservation
Pathwise 11
Statistical: Lœι 0
B Experimental Setup
B.1 Relativistic Dynamics
We consider Hamiltonian system defined full spacetime phase space (M,g), where M is Lorentzian manifold with Schwarzschild metric g. The canonical geodesic Hamiltonian governing N interacting particles takes the form:
H)x,
2M ged (2x)) Pup,
1
Σ
GM
(16)
VIE. 2,
where (z) is the inverse Schwarzschild metric and Pj is the four- nou-mmonntum conjugate to the spacetime coordinate x² = (t,, ri. 02 0.). The second term encodes softened pairwise gravitational interactions.
The inverse Schwarzschild metric in spherical coordinates is given by:
2GM
2GM
0
geyd (r, 0) =
(17)
0
0
0
7 sin
The dynamics follow the relativistic Hamilton equations:
daff
OH
dp
aH
(18)
dt
Öp
dt Ox.
We summarize the physical-to-code variable mapping as:
(r,0,
Position in Schwarzschild coordinates,
0H
Relativistic velocity (spatial),
Op
(Po, P1, P1.P2,P3) Ps)
Canonical four-momentum,
0H
VV
Force from geometry and pairwise interaction.
0q
Each particle E (1....,M, N) is initialized with a four-position q: = (14, ri, 0g db) drawn from:
ri ~ 0.800,0.011),
0 ~ (=T/2,00000)
(0) | 1.37), 6i ~ No (0 (..001) 0.01%),
centered around stable orbital radius with small angular spread.
Initial four- momentum is sampled induce slightly perturbed circular orbits:
2GM/ra 0.1
with Eg ~((,0.001)
Pt
XV

---

Table Initial state (reprodueible random seed 42).
Body
æ:
2
w
we
wy
Wx
1
0.496
0.647
0.576
-0.011
0.228
-0.238
0.040
2
0.261
0.400
0.812
-0.314
-0.141
0.136
-0.113
3
-0.183
0.861
0.318
0.351
-0.119
-0.438
0.176
Frame Bundle Structure and Spatial Diffusiin. Let (M,g) be a Seeudo-Riemannian mani- fold with signature (-, +y +, +). The pseudo-orthonormall frame bundle O (M) principal (0(1,))-undlk A point in this bundle is written as:
U, (xt
et a): where (2) et e. MJab.
and 17 = diag(- 1, 1, 1) is the Minkowski metric. The frame e" forms a local orthonormal basis of M. To respect relativistic causality, we restrict stochastic diffusion to spatial directions
= 1, Let B, = (B("). B3(2), BB(3) be Brownian noise on the spatial frame. The spatially- restricted horizontal SDE on the frame bundle is then:
du, Hor (U.)
(U,)) = X,, U. GLEOOLE (
(19)
dB.)
No perturbation is applied to the temporal component (a 0), and the drift term is entirely de- terministic, maintaining consistency with the Lorentzian structure. The simulation integrates this Hamiltonian system using symplectic leapfrog method, with symbolic metric evaluation via SymPy.
The background mass M = 1.0 determines the Schwarzschild radius r. = 2GM = 2.0.
B.2
Spin Dynamics
Our goal in the expereiment is to model the time-evolution of three nutuully-interactinn ggid-body spins living on the unit 3-sphere 3 CR,, simulate three dynamical regimes, and render the results through the Hopf fibration. Let q (3r, y, 2, w) E 3 be unit quaternion representing rigid rotation.
We split q into the complex pair (a z2) CC,
+ ey, 22 + iw, that |21 + |22|2 1. We treat each spin as a point mass with principal inertia (2, 1, 0.5) and equip the system with the pair-exchange Hamiltonian defined as follows:
J
H H(g (x) Σ
2 w Iwi
( 4z),
ifj
where wi E R is the spatial angular velocity, I= diag (2, 1. 1.0.). and J > promotes aligmment. Then, the Hamilton equations read d: (l(ws) qs W; I-1 Tr, with I(w) and 7, = = EJE
4))1:3 the coupling torque. To simulate the dynamics, we employ an explicit Euler step of size At = 0.02s and renormalize 9; to unit length after each step avoid drift off $3. All runs start from the randomized seed quaternion/velocity ensemble (Table for seed 42). Each trajectory contains T 300 frames. While unit quaternion (21, 22)
C C² projects via the Hopf projection 1 (21, z2) (2)(21))), (2(2)), |21|2 |22/?) E S? we display each trajectory in polar arg magnitude coordinates of (2): (0c, rk) (arg 2k, |=x)), k = 1,2.
B.3 Toroidal Protein Sequence
We convert the raw Cartesian trajectory contained in traj-arrays hat-aarays..pp.. npz into generalized coor- dinates (0,0, P. i)) the dihedral-torus to support Hamiltonian learning. The nnp file provides positions E IR xx velocities Vt. and time stamps t E Dihedral angles 0, = (d,, th, wi) e (- T are computed from atomic coordinates using standard torsion definitions applied to atom quadruplets A A A C (1...., NJJ, which are in turn extracted from the molecular topology file staj-state0 pdb via MDTRAJ. This topology file also provides element-wise atomic masses {m, for moment of inertia computation.
We first estimate a uniform time step Al mean; (4++1 - (x) in femtoseconds from the raw time array. Angular velocities for each torsion angle 0 E (0, 1) w are then computed by finite differencing:
xvi

---

Variable
Symbol
Shape
Angular velocities
0
(T - 1) x
Angular momenta
P
(T - 1) x 3
Momentum derivatives
P
(T 2) x
Table 5: Dimensions of the augmented physical variables.
= (0++1-00)) followed by periodic unwrapping using 0g wrap(--.s)) (02) where the wrap function applies asin2(sin,, cos-) elementwise. The resulting angular velocity matrix O ( is stored under the key torsion, dots. To compute the scalar moment of inertia 1 for given torsion time lg, approximate the rotation axis the normalized vector (Xaz Xay //1 and use only the two terminal atoms to define transverse distances. Specifically,
4 = ag r (xoo) + mas r (Xas),
(r( ||(r xa1) [[r x el
This produces single scalar inertia for each dihedral and time step. Using this, we compute the conjugate angular momentum as P& I 0 and aggregate the result matrix p E RTT-1)x3 stored as torsion, torsion_momentuee _momentum. To obtain generalized forces, we compute the time derivatives of angular momentum by finite differencing:
Pk+1
Pk
Pk (r. (Pt))
At
where we again apply periodic unwrapping. This finally yields the data P E (TT-2)x3, Table summarizes the shape of the augmented tensors.
B.4 Experimental Details
All experiments were conducted on single NVIDIA RTX 5090 GPU using Python 3.11 and Py Torch 2.1.0 with CUDA 11. support. The proposed framework is evaluated on pre-processed protein trajectory dataset embedded curved configuration space. We use an 80%/20% temporal split for training and testing. Each sub-trajectory consists of 0.87 frames: the initial frame to serves as the input, and the following 0.87 frames (f..0.T)) are used for supervision. The model input and target sequences include generalized coordinates, velocities, momenta, and their time derivatives:
(g. q. P, P) TRTx
where 7' denotes the total number of time steps in each sequence.
The neural architecture in Algorithm T is Hamiltonian vector fields on curved manifolds by incorporating symmetry-ppesservng inductive biases. The input to the network is concatenation of configuration and momentum coordinates (9. p) E 2d transformed into canonical local gauge frame using Cholesky-basee projection with gauge matrix G Rddd This projected input is passed through a linear embedding layer and fused with temporal encoding via sinusoidal or MLP-based TimeEmbedding.
Transformer UNet designed to model
The model consists of an encoder-decoder Transformer with L 16 residual blocks, each using multi-head self-attention, GELU activations, and layer normalization. Skip connections and projection layers link encoder and decoder stages. Predictions for time derivatives frame and mapped back to the global frame using a Cholesky-batt based projection, ensuring gauge equivariance. The network uses hidden size of 128 and contains approximately 3M parameters. Training is performed for 10 epochs using the Adam optimizer with learning rate 10-1 and batch size 128. The loss combines local alignment term and long-range reconstruction objective.
generated in local gauge
((0) B P..
[9,, P. 1
1
wrap q 11
A = 0.1.
Nur
Lalign
Lyconn
To simulate future trajectories, we employ geometry-aware simulator that integrates stochastic Hamiltonian dynamics via Stratonovich SDEs on the cotangent bundle. This simulator leverages
xvii

---

Algorithm GAUGE EQUIVARIANT TRANSFORMER UNET
1: Input:
19 1 N. 99 .
EddN, time I, bundle metric G
2: Output:
3: // Canonical gauge transform
4: Qem + GG o., Tcaa + [ean
5: // Encoder-decoder trunk ( (shared context)
6: h Linear (FFem) + TimeEmbed(()
7: < Unsqueeze 1)
8: for i
to do h + EncBlock, (h), enc; +-
9: end for
10: for i = 1 to L do Concat(,, encz-s41); ht DecProj h DecBlock.(A)
11: end for
ot GlobalPool(h)
|| context vector shared by all potentials
13: / Shared trunk Ta and two heads Hap Hgune
14: for i = 1 to N do
b branch shared weights
15: d: Tdlg',p'\,c)
16: E.p Hbp(
// head
17: end for
pairs (i,j), <
pairwise branch // same trunk
19: Pij To(l9' 9') c)
20:
E < Hauir (d:s)
// head 2
21: end for
22: KE 1 G p
// kinetic
23: SP Σ, Fij.
Pt+ Eies F.j
24: He = KE + SP + PL
25: 9 V xx Ho. Tºd = v He
t N
26: ④
// back global frame
N :N
27: return .P'
gauge-equivariant drift fields and lifts the dynamics into a frame bundle, where the noise transported horizontally. Chart transitions are handled using manifold-aware update rule. The numerical integrator uses fixed step size Al 2.0 fs and isotropic Gaussian noise of magnitude 10-2.AI Performance is evaluated using two common metrics: the Average Displacement Error (ADE) and the Final Displacement Error (FDE), defined as
ADE
1
Σ 1mm.. en.e n.t
FDE
1 N
Σ
An To,TTr
N
N
Random seeds for torch, numpy, and random are fixed to 42 for reproducibilityy The codebase and configuration files will be made publicly available at htjhttps://githui com Anonymous /HDM
Simulation of Neural Hamiltonian Diffusion. Algorithm outlines the simulation process of our Hamiltonian Diffusion Model (HDM) based on Proposition 2.2. Starting from an initial state (go, Po), we construct a lifted representation Un on the frame bundle using the inverse metric and its Cholesky decomposition. At each step, the model predicts a gauge-equivariant Hamiltonian drift, and isotropic noise is projected onto the horizontal space to ensure geometric consistency. The state is updated via Stratonovich integrator that respects the manifold structure, and chart transitions are handled as needed. The algorithm outputs both the lifted trajectory on the frame bundle and its projection onto the base manifold, enabling structured simulation over eurved geometric spaces.
Initialisation of the Lifted State Ug = (zo. Eo) E O(M). The horizontal SDE of Proposition2.2 requires an initial condition on the orthonormal frame burulle. This amounts to choosing (i) a base point æn E M-which fixes the particle's is initial configuration-nnt (ii) an orthonormal frame Eo E SO(T.,.M) that serves as the local gauge which all subsequent tangent-space computations expressed.
Xviii

---

Algorithm 2 SIMULATE NEURAL HAMILTONIAN DIFFUSION(J1- Po, lo, nyy Ha)
Require: Initial state (go.Po) TQQ start-end times (lo. T). Stratonovich steps ny, parameterized Hamilto-
nian Hp
Ensure: Lifted trajectory (UrhtoseTT and its projection (X, -(U))}
1: // Imitial frame lift
2: 9: 9:(0u, Pu)
// inverse bundle metric
3: Po + gh)
local orthonormal frame If coordinates on O Q)
4: Uo f- [9o. Po. vec(vo))
5: Co 11
initial chart index
6: I! Stratonovich SDE integration
7: for l = =0
do
8: t to
At T-to "f
9: Tt.Pt, U. + unpack(Uk)
// Horizontal Hamiltonian drift
1g: wp] + FODDELFORWARR((..P.... U.)
Gh (Uk) + (Wy. Wpp
If {m, Ho} part
13:
// Horizontal diffusion term
14:
He + Hirame (qt, U.)
I/ horizontal projector
15:
E E~ NO, I2á). sto V 715
16:
// Stratonovich increment
17:
UR+1 &- UR G (Uk) At + sto
18:
I/ (optionnal Chart update frame bundle CR+1, UK+1 <- CHARTUPDATE(UL C+)
20:
end for return (Uk) (Xx = x(.))
Base point Tp- In practice vo is dictated by the task: for trajectory prediction one sets 20p = Edata (the observed configuration at time 0); for sampling controlled experiments may draw 25 from a prescribed distribution on M (c.g. the uniform measure on S²²).
2. Orthonormal frame Eg. Given a coordinate chart (a) ..... rd)) around zo with metric matrix (xro), one constructs Eg by orthonormalising the coordinate basis via the Gram-Schmidt (or Cholesky) procedure:
Eg = [c1 ... ea). (e,, ()) = -5g-
For the sphere example (d = 2) with (x, up) coordinates one obtains
1
0
el = 0x
e2 sin XO
Opp
Eo
0
1/ sin Xo
Phase-space variables. the model evolves on T'M one also specifies the initial momentum Po E T M. Typical choices are (a) the empirical momentum if one starts from real data, or (b) draw from the canonical Gibbs distribution po (N,0,, g (xxo), which is consistent with the kinetic term 2pp g -1pp in the Hamiltonian.
Vectorised form for the SDE. For implementation the frame is flattened, vec Eo) E Rdd and concatenated with (go, Po) to produce the full initial vector fed into Algorithm ??. Because Eg is already orthonormal, the integrator starts in the horizontal sub-bunn sub-bundle, and the structural proper- ties guaranteed by Proposition2.22 are preserved from the first step onward without any corrective projection.
xix

---

C
Lemmas
Lemma C.1. Ler T T be the connection (i.e., Christoffel symbols) of a smooth configura- tion manifold (2,9). and let R "ikk be the components of the Riemann curvature tensor. Let us
denote ||T||> == z maxi,jj [1 3x (2)), and || supzeQ maxi,j.J (x Then the following sup-norm inequality holds:
2 & 2 + 2
|||| =
||||| = 0.
Proof of Lemma C.7 Recall that, in local coordinate chart on a smooth Riemannian manifold Q2 g). the components of the Riemann curvature tensor are
R' jkl = 0,1 0,1% Fimm J
1 I'm jk-
where r are the Christoffel symbols. Define the sup-norms
|| = sup max #c
(x)(, 13
|0 = sup max Os (x))(())):
and
| == sup max
R' (x)((
ZE R.j.R.)
For the derivative part of R jkt, we have
| 3ry &
|0.1 +
|orj S 2 bl..
For other quadratic parts, we control the terms by showing that iim gkl
T
km m
<
km 37
Mi |
2 |
Applying above results and then we take the maximum over all indices at each point to have
R (x(())
< 2| + || VrE
Finally, this gives
R |x S 2 0 + 2 1
If || 30 = O (so every 13 vanishes identically), then 0&T; = as well, and R = 0 by (*).
Consequently || = O. This concludes the proof.
Lemma C.2. Let (M,9) be Riemannian manifold. Fix smooth reference curve y :0, + M and define
Fi(y) == exp (y) E M, J= F(.:).
where X, is solution the proposed Stratonovich Hamiltonian Diffusion Model (HDM), Assume that X (7(t)) = for all t E [0,T), Then, to first order in the stochastic differential, we have
Dexp (dX,) J., {m, Ho] ()) dBB.
(20)
Proof. We start by recalling the behavior of the exponential map and its derivative. For v E TTg()) M, define the geodesic Yo (s) := exPxce (sm). Given E T (e) f, the differential of the exponential map satisfies
D Dexp (v)[E) (1),
(21)
where Jg denotes the unique Jacobi field along Yo satisfying the initial conditions
(e) O,
V v.. (e() =
At the origin 0, the exponential map behaves simply:
D Dexp (0) ld.
D2 ex 0.
XX

---

This tells that near O, the differential D (t) (v) is close to the identity up to second-ore order terms. Thus, for y close to (()) (writing =xxpx(1) (v)), the differential the inverse exponential map satisfies
7o) (D -dl
(22)
Next, we apply Taylor expansion to the Hamiltonian vector field X around (t)). Since Jr exp, (Xe) represents small deviation, we have
Xи (X,) X (7(t)) + (x)))
(23)
By the assumption (z)((()) = 0, this simplifies to
XH (X+) (10))
Thus, the stochastic differential dX, is given by
dX, J. X (+(1)) dB, + O(1 dB,
Now, we apply the differential of the inverse exponential map to both sides. Using 22 we find
D Dexp (dX.) [ld (0(++( Nw (7(1) dB, + ((II).1))) dB.] = xn o dB. + ((II.1I)) o dB,
Finally, since J. Bxl) under small-noise scaling, the remainder term (1I dB, becomes negligible compared to dB, in the Stratonovich limit. Thus, we conclude the desired first-order approximation:
D Dexp 7()) ( (dXi) fm Ho Hol (r(()) dBi.
Lemma C.3. Let (2,g) be a d-dimensional Riemannian manifold. Assume that the scalar curvature Scal(x) satisfies
sup [Scal(r)] <
(24)
ZEQ
for some constant S 2 0. Then the operator sup-norm of the Riemann curvature tensor satisfies the estimate
28
| < d(d 1)
(25)
In particular, the curvature-induced deviation term (B) in the stochastic Jacobi analysis can be uniformly bounded in terms of the scalar curvature bound S.
Proof. Recall that for any point E Q and any orthonormal basis ferld of T, the scalar curvature is given by
Scal(x)
2K
(26)
i<
where K(e; e;) denotes the sectional curvature of the 2-plane spanned by e, and ey:
K(e,, (()),,c))= g(R(ei, ej)e,, ei)
Ae,l? -((((e,,e))),.c))
(27)
There are exactly (9) d(d-1) independent pairs (i, and each sectional curvature contributes linearly to the scalar curvature. Therefore, taking absolute values and using the triangle inequality, we obtain
( > |Scal())],
2 max K ((e,, [c
(28)
xxi

---

which rearranges to
[x)
max ((e,,, Kg
(29)
d(d 1)
Since by definition
|||
sup # Q
sup ||
| u, u) and [[ [K v)
up to constants depending on the wedge norm u A all = 1 for orthonormal pairs, we obtain
2S
[|]
S sup max |K (ei, v d(d - 1)
(30)
=EQ (c)
where the factor 2 arises from symmetrization conventions in the definition of the Riemann tensor versus the sectional curvature. Thus, the curvature tensor's sup-norm is explicitly controlled by the scalar curvature bound S.
Lemma C.A. Oksendal and Øksendal 22003) Let (Be) be standard Brownian motion. Then, the mean-squared expectation of Stratonovich SDEs can be calculated as follows:
E $I[
f E| ds
1 -
Z, oB.
dB,
E
VZ Zz ds
+
Zg
dB,:
VZ. Zg ds
(31)
xxii

---

D Proofs
This section serves to rigorously formalize all theoretical results that were informally stated in the main text. The goal is provide complete proofs that fill in the technical gaps and support the conceptual developments discussed earlier.
D.1 Proof of Proposition 2.2
Proposition D.1 (Horizontal Hamiltonian Diffusion). Let U, E O(M) be the horizontal lift of the diffiusion process X (U,), where π O(M) + M is the canonical projection and m denotes local coordinate function on M. The lifted process U, evolves according to the Stratonovich SDE;
dU, == =G (U,) dB,, TUU)) Xx, U. E,EO(M),
(32)
G (U.) = ((m,
V. - [[ (1m, Ho),H))))))))) vec)) Ve,
(33)
where [laI M(2d x d is the inde--lowered connection tensor (i.e., Christoffel symbol), and (E) E M(d 11 is the vectorized local orthonormal frame. Va and Ve denote the vectorized gradients with respect to the configuration point and the frame coordinates, respectively.
Proof. We begin by deriving the horizontal lift of stochastic Hamiltonian dynamics in local coor- dinates. This involves expressing the dynamics {X,] on the cotangent bundle '" Q under local trivialization of the frame bundle OM, equipped with moving frame E, For explicit analytical and numerical handling, we represent the Stratonovich SDEs in terms of Euclidean coordinates via the horizontal lift operator. In this coordinate system, the horizontal lift of the Stratonovich-tyee stochastic Hamiltonian dynamics is written as:
HHg Opi
(X,)
72
r (9:l)) eaa eq.h)
(aHL
(Xe)e dB Opp)
j=1
(34)
HHg 0q'
HHe Oqy
p:
(X,)
r (9)) ct
(X))e as
dBc
The above system describes the horizontal lift of Hamiltonian dynamics where the curvature of the base manifold Q encoded by the Christoffel symbols I and the stochastic transport along local orthonormal frames E, jointly modulate the diffusion. The geometric structure is embedded via the lifted noise term on T Q. ensuring that Brownian motion remains horizontal with respect to the Levi-Civita connection.
To enable stochastic calculus, we now transform the Stratonovich integrals in (34) into Itô form. This allows the introduction of correction terms due to the nonlinear dependence the coefficients on the stochastic process. The position dynamics in Itô form become:
HH, 0p:
OHe Op.
d
ri jk
a.k
6 a.
dB,
j=1
H
To rk Op,
et
a.k
HHo Opr
e, a.
r
OE k Op,
(HH
c a,
(35)
Op,0p;
Opr
a=1
Ti a. Ho Öp,üpr
(HH, 3pr
dea
E + tr
aik aak
V dt.
Op:
xxiii

---

Similarly, the momentum equation is converted as:
HHg dqp
HHH
p
E
dB.
jk
==
k=1
Hg .q' Oq)
or rk Öqg
HHH
OE. k Oq''
OHg
Fe
r
F
Öqr
nk
3qr
Ho dq 0ür
HHg
OE,
1 E.
Et
& + r
Vi dl.
dq.-
0qy
(36)
Here, the additional drift induced by curvature and moving frames is absorbed into the auxiliary term VJ, defined as:
((x)= (x()=(115))
155d
HHa Dpy
(X)
-1s
OHp Oqpid d (x)
Σ
(g)
(
(x) E. (37)
o.,r.k
In Euclidean coordinates with Cartesian frames, all connection coefficients vanish (T 0), and the orthonormal frame F,, becomes static. As a result, the geometric correction term V also disappears, recovering the standard stochastic Hamiltonian flow.
To summarize and simplify the geometric formulation, we now express the entire Hamiltonian diffusion in matrix notation. Let us define the following tensorial representations:
{m, Ho} } R 1xx2d
Id @ {m, Ho} Rdx2d
I 2dxx vec(E,) d
Let a r vec((F)) F.) denote the geometric distortion vector. Then, the matrix form of the Hamiltonian diffusion reads:
(Ito) dx = (m, HX)()) dB, +
Σ j=1
D2 (m, Ho} a²) a dt,
(38)
where a denotes the j-th column of and (m, Ho] 2dxx2dx2d is rank-3 tensor containing Hessian of Hamiltonian. This reformulation makes explicit the second-order geometry- cor- rection arising from horizontal noise transport in local coordinates. With the form of Stratonovich's diffusion, one can recover the original definition used in Eq (5) as follows:
(Stratonovich)
dx,
[m. (X.) - (Id @ {m, Ho) Td veo(Ex) (Ft))
dB,
(39)
The first part of the proof is complete by rewriting the above dynamics presented Stratonovich SDEs.
As next step, we aim to establish the theoretical validity of our geometric construction by verifying whether the proposed vector fields Getrr indeed lie in the horizontal distribution of the orthonormal frame bundle ((M), where M : 7 @ is the cotangent bundle equipped with a Sasaki-type metric.
We introduce canonical coordinates on M as
x² ..),
E (I,...,244 2d}.
(40)
With block index conventions where q- indicss are i, E (1....,d d) and fibre indices are 7: d d i, the Sasaki-type metric on M is given by
g
(41)
where g is the base Riemannian metric on Q. The Christoffel symbols of the Levi-Civita connection on M, denoted I fr. have the following block structure (with 0 0/0q'):
r jkl
i, f j. k.
1/2
(0,gth +
a i, B = j, Tr.
Ma By
=
(42)
a = B = 7 = k.
0,
otherwise.
xxiv

---

Let U Ez E) (M)) be point on the orthonormal frame bundle, where E = -EE a) E R2dx2d
an orthonormal frame z, We define the vectorized frame by
vec ()) E IR (2d))
(43)
which stacks the columns of E. We also define the index-lowered connection tensor
πT E (22))x(24)
(44)
via the transformation
(ooM) == T E a.
(45)
Then for any R2dd, we have the key identity:
mat (Ita & v [ (v)E,
(46)
where (M) = I' u7, 07, and mat(-)) reshapes a vector of length (2d)² into 2d x 2d matrix.
We now set
0 v: (m, (z) (0, -0,Hg) E Red
(47)
which defines the base vector field associated with the Hamiltonian dynamics. Then the lifted horizontal vector field on the frame bundle is given by
G = vr [12a @ o "] [ [ 0g.
(48)
To confirm horizontality, we define the temporal derivative of the frame:
E == - [1 @ v (PcMMMM
(49)
which is equivalent to
E (vcc(TM(
(50)
The Levi-Civine connection 1-form evaluated vector field V is given by
wa (V) E (E%, + I o oE,.)
(51)
Substituting the expression for És, b. we obtain
En, an 0 E
(52)
which implies
wO = Ea (-r
By o E E8 + + M 0 = ())==
= 0.
(53)
Therefore, the connection 1-form vanishes:
(G 0,
(54)
which confirms that the vector field lies in the horizontal distribution:
GHor HO(M).
(55)
In conclusion, the proposed Stratonovich SDE
dU,- (U,) dB,
(56)
is the horizontal lift of the Hamiltonian diffusion process the cotangent bundle T* Q, ensuring geometric consistency with the underlying connection on M.
XXV

---

D.2 Proof of Proposition 3.1
Proposition D.2 (Time-uniform Generalization Bound of Hamiltonian Diffusion). Let P. (.())
Law(r)(0)) be an associated probability measure of model trajectory, and assume that the proposed neural networks lie in Sobolev ball i.e., |0,, S R, v0 E G and first and second derivatives of Hamiltonian are Lipschitzian.
P
sup sup W (P.(),), Pr,da) [E
S exp (-2 61/2 |T R/ (1og))
(57)
While the first bound captures the uniform deviation of the learned trajectory distribution from the target data measure across time, we next provide concentration result that controls the deviation between the empirical Wasserstein distance and its population expectation.
P
1 E sup sup W
8y Pduna
13S (7(0).²)
[ejo,]] AEG
x exp
(S2--2/2
3/2
R
"log R)
(58)
n1/1
where W = Q) stands for the squared Wasserstein distance on physical data space T(T Q), Ω Qa Amary LH, LvH, d, s), 2s is a constant depending on metric tensor g and the smoothness, Lipschitz constant of Hamiltonian.
Proof. While the proposed stochastic system is semi-martinge the chain rule with respect to Poisson bracket (i.e., Eq. (2.8) Lázaro-Came and Ortega (0088 direct gives the following result:
d Σg
Th T (U,) mi o T(U)) =
[ Hx)(X.) dB'
j=1 0
Om TH
Om 2He Opk Oqk.
(X.) o dB:
(59)
Ogn Opm
j=1
HH! Opk
10" aHd j=1 Op:
t
(X.) o d = Σ
X, o dB
k-=1
where 5% denotes the Kronecker delta. This reveals that the stochastic evolution of the velocity field depends explicitly on the derivatives with respect to momentum coordinates in Eq. 59 highlighting the necessity of incorporating additional physical information. For further discussion, we first give Sasaki-type fiber metric (i.e., norm) on TM = T(I" Q). Then, the squared distance between y and y² on tangent bundle T'M can be naturally defined as follows:
(1 (7,) = |1(x' p. (d) ppppppp00 P2..
11
(60)
= 1|q q² 11 + llp'
p p21. 11
q 11a0
+ 11 -11??-
where 9- (a 8 9 (a, B) is a dual metric configuration space. Following by the definition of the norm tangent bundle of total manifold (T) the discrepancy between model trajectory 7(0 and data trajectory 7 can be calculated as follows:
(70)))
((()((), Pa (0), (6), (9)) (âu. p á P
[[OO,, (0), qr (0), 0a) - (0,, p.. á. 0z) 1
(61)
< [(0,,. P. (0), d (0), 0.)) (0,, (00). p. (00), 0z) (
+ (O.,, (0o), q (00). 0.)) (Oa, p.. é 0.) 11
Let assume that the test neural network 0 satisñies perfectly matches particle trajectories almost surely i.e., :q o TT( (U)) (0) q.. and assume both mapping vpH(q. : and V.Hlpp :;0) is
XXvi

---

an injective mapping for each fixed and p. Consider another neural network dg both matches both particle trajectory and their corresponding momentum behavior e.g., 8n arg ming (
If there exists inverse Lipschitz constant I of second mapping, then the assumptions leads the second and third inequality in Eq. (62):
W Q(T Q) (1a(()()),
Ed @) (7((),
(62)
$ (l) (60) + Blpp(0) p(
$ (1 + 1 ()) - (000))l
Given the fact that velocity field lies the tangent configuration space dt TQ = RR taking supremum with respect neural networks in both side of inequality in Eq. 62 gives
sup W (Law(>(0), Law (f)) S (1 T sup TE le
(1) qe (60)
DEE
(63)
S (1+ + LH Xmax (g0) sup E [1l á (0))
For readability, we simplify the notation as W: -Wrr-QQ Note that the expectation in this context is taken from My, E PTT Q) for each E [0. In first inequality, normalize the Riemannian inner product with the Euclidean correspondence by using the property: || ( for
any vectors v Tp, Q. Next, our goal to obtain the following type of decomposition
sup [l)c(0) - (G)) [2:] S f. sup ( 1 OE
(64)
of
where the time-dependent constant j == f(t, r, ar, Ho) depends on the connection form and their derivative Ör, and the Lipschitz constants of higher-order derivatives Lip(O'Ho), VI 2. To this end, we first define four auxiliary processes as follows:
D.(0) = Op.Ho.
Aza(9))
(65)
Byaa( (0) = (0p, Ho) Z. ===1 == (g (4),
J.a
where the mean-squared norm of processes || || | 2 D / is bounded above with some constants Ca, CB, Cz, CD, Cvz. Having the definition in hands, the proposed velocity vector fields for arbitrary network can be simplified with the following form:
(0) / D2(0)=
Σ (4, J.0
dB..
(66)
With the definition ōq. (0) q (0o) for deviation between two velocity vector fields, direct calculation leads to have norm- corm-equared expectation as follows:
l 8D. - A)ya(ga) (s) dB,
E [1 11 = E
3.02
(67)
11
22
11
2
$ 2E
6D ds
+ 2E
Σ
Ajaa (90) (B) odB,
3.0
=2
Following by the conversion of Stratonovich SDE into Ito's SDE in Lemma C.4 the second term of right-hand in last inequality can be upper-bounded with the following form:
E 11.. Z, dB.
2
< <
1+ 1 1 f.zz+t8c +1CCx2
(68)
xxvii

---

This bound reflects the second-moment moment structure of the Stratonovich integral, where the dominant contribution arises from the squared noise norm 02 and the correction terms involve both the norm and gradient of the stochastic vector field 74.8
|| CA S Vd |
|
00 sup sup ja (9)
(69)
ge g j..
| g CB - Lu 1 Bull || = C2 d' CCCAB - / mmx
(70)
LH 10 0oll-
(71)
This shows that the magnitude of the noise vector Ze grows quadratically with the model distance ||0 Boll, and is modulated by the geometric curvature ||T|| and the Hamiltonian smoothness Lи-
VZ. Σ (A),((g)) (B) + Ag,a(a)) ( J,
(72)
This decomposition separates the gradient of the stochastic vector field Zr into two terms: one involving the spatial derivative the geometry-aware coefficient A3 and the other involving the gradient of the perturbation jB both of which are influenced by the manifold structure and the Hamiltonian model.
[A(gg) = OVA Vallo
|F
sup sup jark (g)l² q j.ak
(73)
| = CVB S (LH + L Eg | Boll.
(74)
Since the orthonormal frame is locally updated by parallel transport, one can obtain
Ogg Ea (qa) = (a) IV
(75)
max
Therefore, the spatial variation of each frame component E entirely determined by the Christoffel symbol and remains uniformly bounded under smooth parallel transport.s
V = Czz - (C + CACUB)
(76)
s 5 50 LL + [[ + LH [+] -
Hence, the total spatial gradient 74 scales linearly with the parameter deviation ||0 0oll. and is tightly controlled by the geometry through ||T'|ooo ||01]] and the Lipschitz constants LH, L
E
ds
S f(,,r,r E [1000
(77)
For I. T. V Lv f is defined as follows:
Kt, ar) = LV L' (1 (1- 6 [|0T + [[]] +||[) [|T)))
(78)
Now, our goal to simplify the inequality, making f related to the curvature skewness. Using Lemma C the curvature tensor provides a lower bound on |[ which allows to climinate the explicit derivative dependence and reparameterize Ri in terms of |[
and |[T]
|o S || + 1|
|| 2 2 ||
||
(79)
This inequality gives a curvature-dependen upper bound on the Riemann tensor norm terms of the supremum of the partial derivatives and ||[],.. which allows to replace || by [[]] in subsequent expressions. By combining the previous inequality with the triangle incquality, we obtain
uniform bound || + ao (1 + |[T) that is linear in || and ||] facilitating simplification of higher-order terms.
|| + [| + |[[])=) < || - || + |1 + [T]
1
2 R| + || (80)
Thus, the curvature dependence be simplified to function of R2] and ||||| only.
(| + [1] ))] < 2
2 R 2 (81)
xxviii

---

Squaring both sides, we derive bound for the squared norm (||r + |[] (1 + ([))) which ensures that second-order curvature contributions can be expressed as function of || and ||
alone.
(1+ | (5 (1 211
||
(82)
As a result, the original function i(,,T, R) can now be written in terms of |||] only, up to multiplicative constant, removing explicit curvature dependence from the generalization bound.
We now consider the of neural networks 0o constrained within metric ball defined by a Sobolev- type functional distance. Let θ denote the set of such neural networks whose Sobolev norm and supremum norm are simultaneously bounded by constant O. We then define the associated function class with respect to the -o-no and consider a probability space (O, E Pn) supported on θ
θ = 10 E w n 1 (mg); [10l Sr, 100 S RJ.
(83)
F = F(0) = |00g ( 0,00
(84)
For the random variable ((w) Ene let define auxiliary processes as
1 Tt Σ F(O)) EF(O.)
3 sup sup F(O:), x sup sup FEFE
(85)
(E(O.T)FFEP
n
For the metric dyt(t, F), (s, G)) = [(())) - G( ((0,))lx>(r)) where s I t E [0, the e-covering number on product space [0, 7] x Fi can be interpreted as a product of two sub-coverings in separate spaces:
(Cc, [0, T x Fé, d=) - N(E,, 1l, ( [0, I.
(86)
This shows that the metric entropy can decomposed as summation of two sub-terms:
log (e, [0, T] x Fá. dy) S log (E1. Fb, [1-112,) + log (52, [0,T].1
(87)
where the algebraic constraint on E and 52 is given as
of + (eB,)² (5VI
Bg (0.
(88)
As next step. we derive the upper bound of expectation for the variable 3. assuming that F(,)) has controlled by Gaussian-like long-tail property. Proposition 1.2.1 [Tlagrand (2005). Specifically, we apply Dudley's entropy integral bound to have the following result:
E. [3] = Es
sup sup F(0) [E(o.T]
FEFS
(e)
lim sup y-000
NogN(e,10,T] x F7, d
(9)
lim sup
Log.N |1 + log (E2, |0, T). Dde
1-000
(89)
(9)
lim sup
VlogN θ 1, log ((e,, 00, T). Ddde
(e)
lim sup p-000
Nog.N(eBB, BW l1:)) + log No Be, [0,
D)de
+ log 1v + - d>s,
EB
where the expectation Eg is taken with the probability measure P.. The third inequality naturally follows from the embedding Wa. ب 1,2 and the fact that s<(a) ||0 - all2,2 1-Lipschitz for
some (r where consider the function composition Fg &p Bw..arr) (r). In the fourth
xxix

---

inequality, we follow the entropy number in metric ball with radius r in Sobolev space W ct al. [2013). and taking supremum under the constraint ||- Il. I R.
Wellner
Since the final expression is non-integrable in general case, we only provide their approximation bound with Taylor expansion in the case when the first term in square root nominates the other term.
[3 < (2r) R2- d/2s
1
Rd/26+1
21-d/2syd/ d/2s log(R)
d/2s+1] OR 2),
(90)
d/2s
where d > 2s, and we simply set the variable 3g 0.5. After optimizing with respect to the radius the right-hand side, we finally have the time-uniform upper-bound of empirical estimates.
E OR log R))
(9))
Let us now turn our attention to the empirical concentration behavior of the Wasserstein distance, initiating our analysis with the standard symmetrization lemma.
1 Σ a F(O)) 8
Eo [x] < 2 = 2ar
sup sup (E(O.T] FEFS
Diam ())
lim sup
log,Nee, [0, x Fo, dz) de
(92)
Vn
Ju
7-+00
12
Diam( (
limsup 1-000
VlggN(s1, Fū. 11-111>) + log ((2,, [0, T) Ddde
Vn Ju
where the empirical version of metrics can be rewritten as following form:
Fl IFFO)
F
- G(O))-
(93)
Here, ag ~ U[{+}}] and R,, are denote both Rademacher variables and their corresponding empirical Rademacher complexity. As with the similar calculation conducted in Eq. (89), the expectation of X admits an upper bound that involves both the function class radius R and sample complexity n.
(x) (RA log R. n
(94)
The result was obtained by using the identical metric entropy calculated in Eq 89 to derive the upper bound. Note that the additional assumption of Gaussian-like property is not considered here as opposite to first inequality in Eq. (89). Combining the result in Eq. (63). Eq. (64) and the definition
of auxiliary processes Eq. 85). we obtain two inequalities
R11) = sup sup W (La((>())) P Faam))
21 3,
(95)
telo.]] Bee
n Σ (,,)) Panaa)
- EOW (1(0).)) < 21 x.
RR2) = sup sup (E(o,T] T lEe
W
(96)
21 = (1 + LH 2 f.
(97)
Next, we show the exponential probability inequality associated with (I2()t), En,, Pp,)) by introducing the classical result from Theorem 3 Massart (2000. Here, the probability space is considered as one specific choice of generic metric space.
3 1 x
Eg 33 E. [3C]
54
5
216
P
< (1 61+
Ra
S -
(98)
5
where 2 = sUP Var( (()) is maximal variance of function class, and two positive constants
>> 0 are arbitrary. Collecting the result obtained from Eq.
889
(95)
(66).
988. we have
(R1) (R2)
E. (1+ oV
54
216
P
<
+
Rx
S e
(99)
21
E. [3]
2
55
XXX

---

Our goal is to find the optimal constant D such that the quadratic expression in right-hand side serves an upper bound for the left-hand side polynomial of linear order in Eq. TTOp.
A - Bvz Cx S 2Dz²
(100)
Here, the constants A, B, and C capture the contributions from the empirical moments, noise level, and smoothness complexity, and are given by:
A = a +) = Eg 3) [S [2]]
216
(1++),
B 05 C
(101)
2 58
By minimizing the right-hand side with respect to the free variable æ, we obtaín the optimal constant D that balances the quadratic and linear terms:
D CE + -E-A
E -3 + v9B² 32AC
(102)
2
4C
Substituting the identity Dx = ô into the exponential tail inequality, we arrive the following probabilistic bound for (R.)::
6 P₂ (R1) S 0] S exp (
(103)
Observing that V9B² + 32AC 2 /32AC and setting = we obtain the following upper bound for Da, 2) by simplifying the denominator expression:
C² +
BC
208849 3 4573/22
D(a, 2) S 2
+
(104)
3/2
100a
103/2 5/
By substituting the upper bound of D(a, 21 obtained in Eq. (104) into the general exponential inequality in Eq. (103), we derive the following explicit bound on the probability of the event (R1), which reflects the asymptotic decay behavior in terms of and the structural parameters:
P [R1) S
(-v 10R (log R)3//$
1) (105)
$ exp
V 208849R2//2 R)1/1 18536.41 30g
We now further simplify the expression by isolating the leading-order terms and observing that the dominant contribution arises from the linear dependence on R! / R) 1/4 in the denominator. This leads to a more interpretable asymptotic bound expressed in terms of the supremum norm of the function class.
P (RI) S 5) exp (-2 8 1 R (log R1)/4).
(106)
Similarly, we consider the case where the effective complexity a scales with the number of samples as a R /lg R/n. This reflects regime where the resolution increases with sample size, leading to the following generalization bound:
Py [R2) < 5]
10R3/4 3/ (log 3)3/8
208849v R([log R2/14 18536 41 /3n=t/4 ) (107)
1
2 exp
-V
1/4
As before, we simplify the expression by extracting the leading dependence on Rt log R, and n to arrive at an asympiotic bound that reveals the effect of sample size scaling on the generalization rate:
Pa (R2) <
exp (-2 8./8 1//2/[111/22
3/2
R R2 R) R)1/4
(108)
1/4
where ? = M(o Amax, LH, L s), d > 2s. This final bound highlights that under the sample- Rv/log R/n, the generalization error decays exponentially in the effective resolution scale It and the logarithmic complexity log Rt with an additional improvement in rate proportional 1/4, The result reveals how incorporating geometric inductive bias and adaptive complexity can yield sharper generalization guarantees in high-dimensional structured models. 0
size-aware complexity scaling
XXXi

---

D.3 Proof of Proposition 3.2
Proposition D.3 (Worst-case geodesie deviation under frame rotations), Let [0.T] - be
reference physical data represented geodesic. For any Hamiltonian Hg define the rotate trajectory (h) by Xz(h) == (qt, hp.) for h E O(d). Write Kg == (Le + || Co) where Le and Cg are the Lipschitz bounds of (m, Ha} and {m, Ho). D is the diameter of some compact domain K., and R is the Riemann tensor. Then, for every & E [O,
sup d' (1), hEO(d)
sup ((x(()) s dP(Xo. c hEO(d)
where X is generated by la-gagge-eqivariin Hamiltonian H, satisfying the same bound with smaller constants LØ < Lg. C$ < Cg. Thus gauge equivariance minimizes the worst-case geodesic deviation over all frame actions h.
Proof. Let assume that the physical data trajectory y [0, Q forms geodesic, which satisfies the vanishing connection:
(()) = 0.
(109)
This identity ensures that the acceleration of y with respect to the connection vanishes, meaning that 7 locally minimizes path length and follows the intrinsic geometry of the manifold. In order to quantify small deviations from the reference geodesic 7, we define a deviation vector field Je as the logarithmic map from yt) to nearby perturbed point X₁:
Jr = exp,7)) (X,) log (X,) E (T(()
(110)
This construction allows to express perturbations within common tangent space at 7(), facilitat- ing differential analysis. To derive the stochastic differential equation governing J.. we apply the chain rule for Stratonovich differentials adapted to manifold settings. This yields
dF (t, dt)t + DF X [X 2
1 Σ R
dJ, dt
F(t,
(111)
(X. + DF X.) [X.]
Ja + J. (m, Ha (T()) dB,
where F(t,x) = exp (x) and {E&} denotes an orthonormal basis the tangent space. Expanding each term individually, we observe that the time derivative of the logarithmic map corresponds to the covariant derivative along the base curve 7(), giving
dFF 0 di at exp F.
(112)
Additionally, the differentials involving dX, are computed via the pullback under Ft, while second- order corrections involving V F are responsible for curvature effects, although these higher-order terms will vanish to leading order under our assumptions.
To measure the growth of deviations quantitatively, we introduce an energy functional E()) defined by the Riemannian norm of the differential of Jy:
1 dF(()) z (dJ,, dJ...
(113)
Here, g denotes the Riemannian metric lifted to the tangent bundle, such as the Sasaki metric if necessary. For notational clarity, we introduce two auxiliary processes:
K, V L, = J., (m, (+)()
(114)
The term K2 represents the covariant derivative of J along the trajectory yt), while L, captures how the Hamiltonian vector field varies along the perturbation direction Jr. Following by standard estimation of SDEs (i.e., dt² = 0, dt o dB, = 0, (dB.)² = dt), we have
dF(()) = 5 (ddt,, dJ.)) = z (Ke, Kr, (dt)² + (K,, I.) a dB, + 20(T,, (dB.)) 3 (V,, (m, (t)( fm, H,)((()))) dt _ I (m, (m,
(115)
xxxii

---

Finally, differentiating once more with respect to time and expanding the covariant derivatives using standard curvature identities yields
dÉ(()= dt (29 (m, Ho)()))), V., (m, (m, Ho)(((())))
d
9 V i(0)) J., (m. (..mm.HH)(r(tt))
(116)
9 J. V 5()) (m. (o)(())) (i)( Ho), J {m. Ho)(()))))
where R denotes the Riemannian curvature tensor. This equation connects the second derivative of the energy with the curvature of the manifold and the structure of the Hamiltonian flow.
While the lifted Riemannian metric g is compatible with the Levi-Civita connection (i.e., Vg = 0), we can relate the time derivative of the metric pairing along the trajectory as follows:
d di (U,V) - (V+00) V) + U. (V(())
(117)
where U= - (m, Ho} denotes the Hamiltonian vector field evaluated along the curve (4). This relation reflects the fundamental property of metric compatibility and provides way to track how the energy associated with the Hamiltonian flow evolves along the trajectory. To simplify further, we notice that the covariant derivative of vector field composed with >(t) can be expanded by the product rule of covariant derivatives along curves:
V
(() J., W = Va W ((1(t)
(118)
where W = {m, (r))), and R denotes the Riemannian curvature tensor. This decomposition separates the effects directional covariant changes along the flow from the intrinsic curvature- induced distortions arising from the manifold's geometry.
Given the above expansion, we can derive an upper bound for the second derivative of the energy functional Ett in terms of the norms of relevant geometric quantities:
dÉ((t= z
J.) V V 100 (m, Ho)(() + (4) Halll, (B)
J (m, H+)((
(A)
=dE(1)
(119)
Here, term (A) corresponds to the covariant second derivative contribution, while term (B) captures the effect of curvature-induced deviations along the geodesic trajectory. These two contributions
govern the overall behavior of the energy growth along the stochastic Hamiltonian flow.
Applying Grönwall's inequality to this differential inequality, we obtain an exponential upper bound on the evolution of the energy deviation:
d² (X,, (t)) dE(tt) < dE(()) (A)).-
(120)
where (A) and (B) represent the supremum bounds of the two terms over the time interval of interest.
This estimate provides key control over the divergence between the stochastic trajectory X, and the reference geodesic (t) under Hamiltonian diffusion dynamics. We now aim to show that imposing an equivariance constraint the Hamiltonian function intrinsically reduces the upper bound on the growth geodesic distance. Specifically, the two key terms in the upper bound (92) can be estimated separately as follows:
(A) = vV,
(1))()))) & | Il ( (m, HoHlop ||lllo (1(())
(121)
BB) [R()()), J.) Ho)ll, < [| [| [[)( [m,, Ho
(122)
where || denotes the sup-norm of the Riemannian curvature tensor over the manifold. Next, consider the action of the orthogonal group O(d) on the configuration and momentum coordinates (g. p). For any h Odd), the equivariant Hamiltonian satisfies
Ho(g. p, U h) Ho(g, p(h '1p,.')pp U)
Wh e O(d).
(123)
xxxiii

---

h GL((2) acts on (g.p) coordinates. Differentiating both sides with respect
where p(h)
0
la]
to (q.p) and using the chain rule, we obtain
p(h V (1.p) (g. p. U h) phh V (() U),
(124)
where we have used the fact that p(h) is orthogonal, and that differentiation of p(h introduces right multiplication by p(h). Differentiating once more yields the second derivative relation:
V (q.0)) (g, p, U- h) = p(h) V..ri (g p(h ")p, p)p(h")
(125)
Since p(h) is orthogonal, it preserves the operator norm of tensors. Therefore, for any matrix A, |||| for all A.
||(h) ||||| (pph)))p(k I
(126)
This property ensures that the norm of the differential operators remains invariant under frame transformations. Using these observations, can describe the differential operators acting the Hamiltonian as
- Vim, Ho] (a.)) Ha}, V (mm Ho) = Via {m. Ho(g, U)}.
(127)
From the orthogonality and isometry of p(h), follows that
sup kEO(d)
(o U U-h]] keO(d)
sup
Ha p(1 "),, U)l
sup
IV (a UU|
(128)
heO(d)
(V U)Il.
and similarly for the second derivative,
sup
IV (a... U b)ll = IV (g. U)]..
(129)
keO(d)
Let & s (dd) be the gauge subgroup that leaves the Hamiltonian Ha invariant, that is,
Ha(g.p, U He p(h U), Vh E 4, U O(d).
where p: 0 > G2() is the canonical block embedding p(h) = diag h. 1a) acting on (q.p), and the right action U- is free and proper. Consequently, O(d) - (d)) / o is principal d-bundle with a smooth projection T ((d) + ((d) and the base (dd) / is a smooth homogeneous manifold homeomorphic the coset space ((d)/9 / G/K for some closed subgroup K Φ
Endow the bundle with the canonical Ehresmann connection induced by the Levi-Civita connection of the configuration manifold. The tangent space at U O(d) splits as (d) Hy 2 Vy, where Vy is the vertical subspace associated with the @-action. Since & acts by isometries, horizontal lifts preserve the Sasaki metric and the Itô Stratonovich structure; Brownian noise injected along Hu descends canonically to the base (dd) /b. Let [U] E (dd) denote a frame class and fix smooth section O(d) O(d). Define the lifted stochastic flow by
X = (o(UU)))
where X-) solves the stochastic Hamiltonian system dX, = {m, Hz) dB,. Because Hg is e-equivariant, the Hamiltonian vector fields satisfy
(m, Ho) (g. P. a ((U) h) = p(h (m, Ho p(h)p. (UU)), Vh E 4 where the lift p respects the orthogonal structure. Consequently, both the first and second derivatives obey
V (g,,) Ho(o(())
(o,
Pi Ho (o()) h) p(h (o(o()))))))
which implies the norm invariance relations ([((())))))) He(o((((())))))
2 Ho
Ho
Xxxiv

---

Let J/u = exp (x! be the deviation Jacobi field and define the Sasaki energy
(tur(t) = 29 (J. (
Using the stochastic Jacobi equation and applying the curvature identity
V (t) V., J. V., Vito 500 + Jt),
we obtain differential inequality controlling the energy:
(tu s (Aly) + Elj((t:
where Atuj == (m, Hotloo and Bluj := || |l {m, Ho llg evaluated along the flow. Due to (96), the terms A and BJ are invariant across [U] E O(d)/B. Thus, applying Grönwall's inequality yields the uniform bound
Ei (t) & (0)e (A+
d(x!1) [0] (()) S dE(()e VUJ E O(d/2.
In contrast, for the non-equivariant model (without 4 symmetry), the best bound achievable is ((x,(h))n(t) < dE(())((Aaa (aaee(Aax-BBaaaa Bmaxht
Amax
sup
A(h), Bmax
sup B(h), kgO(d)
heO(d)
with Amax 2 and Bmax 2 B in general. Hence, we conclude that U sup d (x (t)) dz (x. (t)) s dE(()) sup (U]eO(d)/$
((x,(t)))(t). NEO(d)
Since ((d)/9 / is compact (being the quotient of the compact Lie group O(d) by closed subgroup), the uniform estimate (97) implies exponential Wy-stability of the equivariant diffusion. In contrast, the non-equivariant dynamics suffer from a generally larger exponential factor (Amax + Bmaxx) > (A+B). Therefore, sauge-equivarian Hamiltonian learning provides strictly better control over stochastic deviation and ensures tighter uniform generalization error bounds on curved configuration spaces.
We now further refine the upper bounds for the error terms (A) and (B) by exploiting the geometric structure of the Hamiltonian system. First, we recall that the canonical symplectic matrix J, satisñies |اررا|| 1, as it acts isometrically on 7 (Q) and preserves the standard Riemannian norm induced by the Sasaki metric. Moreover, since the reference trajectory 7() parametrized by arc-length, we have
||((t) = for all E [O,T].
(130)
Furthermore, assuming that both the stochastic trajectory X, and the geodesic >()) remain within compact subset K. C T Q, we have
|elll = -(x,, S D,
(131)
where D diam(K) denotes the geodesic diameter of K.. Under these simplifications, the two contributions (A) and (B) can be bounded more explicitly. Using the fact that ||روا| = and [H))t) 1, the covariant second derivative term satisfies
(A) J/ Vi (m, (+(()l))) < Lo
(132)
where La := supzer- e [m, He (z) denotes the global Lipschitz bound on the Hessian of the Hamiltonian vector field. Similarly, the curvature- caraauureinnuuced deviation term satisfies
(B) ()) (m, H>)(I, ||] Ca D,
(133)
where |/ : sup suP = [u, v) is the global sup-norm Riemannian cur- vature tensor, and C. == sup=e7* o |1 m, (m, (x) is the global growth bound on the Hamiltonian vector field. Therefore, the sum (A) (B3 admits the uniform estimate
(A) + (B) S (Lo + ||
(134)
which depends linearly on the diameter D of the compact domain K. and the regularity constants Le and Co. If the Hamiltonian Hę is de-qquivariant under a subgroup & c O(d), the constants L
XXXV

---

and Co improve to Lg and C respectively, reflecting the additional regularity induced by gauge symmetry. In this case, we obtain the sharper bounds
(A)ea S L8 D, (B)og S 11181000 C3 D,
(135)
and the total deviation is controlled by
(A)en + (B) S (Lg + [1]1
(136)
where L3 < Lo and C < Co due to the symmetry reduction. If the space is assumed to have
bounded scalar curvature Kmaz. can improve the by following Lemma C.3 as follows:
(A)eq + (B)en (L0 + Kmar d(d d(d-1) D.)
(137)
Substituting these improved bounds into the Grönwall estimate derived earlier, we conclude that the equivariant stochastic Hamiltonian dynamics exhibits exponentially tighter control of the geodesic deviation compared to the general non-equivariant case, with an explicit exponent that scales linearly with the curvature bounds, Hamiltonian regularity, and the diameter of the compact reachable set. D
Corollary D.4 (Worst-case W2 deviation). Under the setting of Proposition D.4 assume, in addition, that the initial state deterministic, Xo 7(0). For each frame rotation h E O(d), let set
(h) == Law/ (+)),,
(P = Law X (h))-
Then, for every E [).
sup (h), < sup ((P.())),,1(1) S dd(oo heO(d)
heO(d)
with the same Ke as Proposition D.4. Hence gauge equivariance minimises the worst-case Wasserstein divergence from the reference geodesic over all frame actions h.
Proof. Fix rotation h E (dd) and let (h) Law(X, Because the reference point (() is deterministic, the unique optimal coupling between (h) and 6yce) is the map X.(h) - +> ((). Hence
(P))())) E. d X (4),
(138)
Proposition C.3 provides the uniform pathwise bound d X, (()) S dx ((0) ea for every h. Taking expectations preserves the same right-hand side yields
(KP:(h), 670)) < dd (Xo, (())
nat vh E O(d).
(139)
The same reasoning with the gauge-equivane with the smaller constants Lg Cg by Proposition C.3 this already realizes the infimum over By taking supremum operator over h E O(d) in both, it completes the chain of inequalities stated in the corollary.
trajectory X, (h) gives analogous inequality
XXXVi

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency. research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No]. or [NA].
[NA] means either that the question Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will published with the paper.
The reviewers of your paper will be asked use the checklist as one of the factors their evaluation. While "[Yes] is generally preferable to "[No]" ", it is perfectly acceptable to answer "[No] providedd proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering binary way, we
"[No] or "[NA]" is not grounds for rejection. While the questions are phrased acknowledge that the true answer often more nuanced, please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The manuscript effectively communicates the primary goals of the research.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
xxxvii

---

Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-speeification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faireess
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: The appendix includes the necessary theoretical background and underlying assumptions for the analysis.
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification:
The appendix provides clear account of how each experimental setup was constructed.
Guidelines:
The answer NA means that the paper does not include experiments.
xxxviii

---

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g.. with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The full description (Py Torch code) of the proposed method will be released in GitHub.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https
chaips.cc/public/gun http:///ww...eeeeeueuuueiiieeeee iimmiissiinni ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
XXXiX

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer:
[Yes]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Ye]]
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are of range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
xl

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httppp////eurrps..cccpp/////icce
Answer: [Yes]
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper poses such risks.
xli

---

Released models that have high risk for misuse or dual use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e-g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released. the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
xlii

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (ttp::////errrppccccooofereeeee for what should or should not be described.
xliii


Reviews:
[{"review_id": "RbNShhlTcp", "reviewer": "Reviewer_TKAP", "summary": "This paper provides a neural hamiltonian diffusion framework to extend Hamiltonian neural networks to handle stochastic dynamics on curved manifolds. The author provides a solid theoretical foundation, integrating hamiltonian, non-euclidean geometry, and stochasticity. This work also provides some experiments to evaluate the good performance of the proposed neural hamiltonian diffusion.", "questions": "check weaknesses. \n\n## Justification of my score\n- This paper is theoretically rigorous. However, the paper doesn't provide the experimental details of baseline methods, making the experiment section less convincing. I suggests the author to provide the experimental setup. I will be very happy to update the score accordingly.", "limitations": "no, check weaknesses.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- This paper is theoretically rigorous, including a solid theoretical foundation, rich mathematic analysis, and generalization bounds.\n- Apart from the solid theoretical results, this work also provides experiments showing the good performance, compared with many other baselines. \n\n## Weaknesses\n- 1) This work pays a lot of attention to theory, moving algorithm details and experimental setups in appendix. 2) This paper requires reader to have a good background on hamiltonian, manifold geometry, and stochastic calculus. These two points make the paper hard to read for broader ML audiences. One cannot assume ML audiences either have such a background or have enough patience to read all mathematical details before knowing \"How does the proposed framework work?\". \n\nSuggestions: Actually it is possible to balance. I suggest author to create a  background section that introduces hamiltonian, manifold geometry, and stochastic calculus. Then tell readers \"How does the proposed framework work?\". (It can be an algorithm, figure, or description). After that the author may provide **key** theoretical results followed by **key** experimental results. Finally, putting everything else in appendix. \n\n- The author provides experimental setup for the proposed method in appendix. However, this paper does provide the experimental setup for the baseline methods. For example, line 583 \" The input to the network is a concatenation of configuration and momentum coordinates\". We know this (q, p) concatenation is very helpful to Hamiltonian. Do other baselines use the same input? Another example is line 591 \"The network uses a hidden size of 128 and contains approximately 3M parameters.\". Do other baselines use the same size of model? \n\nSuggestions: a paragraph that summarizes the experimental details of both baselines is helpful. Otherwise, the experimental results become less convincing. \n\n- What is the computation cost of the proposed framework? Do you have the experimental comparison on runtime and/or memory?", "was_revised": true, "final_justification": "I recommend score 5, because of experimental details and results.", "rebuttal": "Dear Reviewer TKAP,\n\nThank you very much for your constructive and detailed feedback, which has significantly helped us improve the manuscript. We carefully acknowledge your comments and commit to fully addressing each of your points. In the revised manuscript, we will explicitly emphasize that trajectory-matching methods fundamentally lack consideration for physical consistency, inherently limiting their ability to effectively handle physical data and resulting in substantially degraded performance.\n\nWe will explicitly integrate each of these clarifications and experimental details into a clearly structured appendix section, ensuring readability and comprehensibility for broader ML audiences without sacrificing theoretical rigor. We genuinely believe these concrete steps will fully address your constructive suggestions and significantly enhance the manuscript’s clarity and completeness. We humbly ask you to positively reconsider $\\textbf{your evaluation}$ in light of these promised revisions.\n\n--------------\n\nQ) ($\\textbf{General Audicence}.$) This work pays a lot of attention to theory, moving algorithm details..., One cannot assume ML audiences..\n\nA) We sincerely appreciate your valuable suggestion regarding the readability and accessibility of our manuscript for a general audience. We fully acknowledge that the current manuscript, being heavily theoretical, might pose challenges to readers less familiar with geometric machine learning concepts.\n\nTo effectively address your constructive feedback, we will explicitly make the following improvements in the revised manuscript:\n\n1. Providing a Clear and Concise Background Section} (currently Appendix A.1):\nThe background information, currently detailed in Appendix A.1, will be significantly condensed and rewritten in simplified, accessible language and integrated directly into the main manuscript. We will include intuitive explanations, concise definitions, and simple illustrative examples to ensure general ML audiences easily grasp the fundamental concepts such as Hamiltonian dynamics, manifold geometry, and horizontal lifts.\n\n2. Enhancing Writing Clarity and Readability: In addition to the above structural changes, we will systematically revise our manuscript to minimize unnecessary technical jargon, and clearly define essential technical terms at their first use in the main manuscript. We will also provide intuitive flowcharts to aid reader understanding whenever possible, thereby significantly improving accessibility to broader ML and interdisciplinary scientific audiences.\n\nOnce again, thank you very much for your thoughtful and valuable feedback that has significantly improved the manuscript's quality and accessibility.\n\n---\n\nQ) ($\\textbf{Experimental Setups}.$) The author provides experimental setup for the proposed method in appendix. However, ...\n\nA)  We fully acknowledge your valid concern that our strong theoretical focus might have limited the accessibility and clarity of experimental details for broader ML audiences. As you rightly suggested, we will explicitly add a dedicated and detailed section in the Appendix of our revised manuscript, clearly summarizing comprehensive experimental settings and comparisons between our proposed method and all baseline methods. Specifically, we plan to carefully address your following points step-by-step:\n\n1. Model Architectures and Hyperparameters (currently partial Appendix B): Provide explicit details for $\\textbf{both our method and the baselines}$ such as proposed architecture, hidden dimension size, number of layers, attention heads, number of parameters, and precise optimizer and learning rate settings. Similarly, explicitly document hyperparameters for all baseline models to ensure fair, transparent comparisons.\n\n| Method              | Attention Type                 | Hidden Dimension | Attention Heads | MLP Depth (node-side) | MLP Depth (edge-side) | Graph Information |\n| ------------------- | ------------------------------ | ---------------- | --------------- | --------------------- | --------------------- | ----------------- |\n| Our Proposed Method | Multi-head (+ Set-Transformer)   | 128              | 4               | 2                     | N/A                   | No                |\n| GeoTDM              | Single-head Attention          | 128               | 8               | 2                     | 2                     | Yes (graph-based) |\n| EqMotion            | Single-head Attention          | 128  | 1               | 2                     | 2                     | Yes (graph-based) |\n| SympHNN             | Fully-connected (no attention) | 128              | N/A             | 3                     | N/A                   | No                |\n\n2. Computational Complexity and Resource Requirements: Clearly document the practical GPU memory requirements (training and inference) and runtime measurements for both our method and baseline methods. We will explicitly include the detailed computational complexity analysis you requested, in the form of a concise table clearly summarizing GPU memory consumption and runtime (both training epoch and inference step).\n\n| Method              | GPU Memory (Training) | Avg. Runtime (Training Epoch) | Avg. Runtime (Inference Step) |\n| ------------------- | --------------------- | ----------------------------- | ----------------------------- |\n| Our Proposed Method | \\~16 GB                | \\~15 min                      | \\~2.6 s                       |\n| GeoTDM              | \\~15 GB               | \\~29 min                      | \\~2.2 s                       |\n| EqMotion            | \\~12 GB               | \\~25 min                      | \\~1.5 s                       |\n| SympHNN             | \\~8 GB                | \\~14 min                      | \\~3.3 s                       |\n\n\n3. Detailed Comparison to Prior Work: Provide an explicit and detailed discussion of the differences between our proposed approach and other relevant trajectory-generation baselines suggested (e.g., GeoTDM, EqMotion).\n\n---\n\nQ) \"Do other baselines use the same input?\" \n\nA) We sincerely thank the reviewer for raising this insightful point, as it provides us the opportunity to clarify key distinctions between the baseline methods considered in our manuscript. \n\nWe have two primary groups of baseline methods: \n\n1. $\\textbf{Methodology 1 (variant of HNNs)}$ : These methods share the same underlying data input format (q, p) and training objective function as our approach $\\textcolor{blue}{without~ connectivity~graphs}$. Thus, they inherently leverage complete physical-state information. However, a fundamental distinction arises from the type of force field employed: traditional HNN variants assume a purely Euclidean force field, whereas our method generalizes this by introducing a Riemannian force field on various differentiable manifolds.  As similar to our model, all the methods in this category produce (q,p) by processing the same size of input (q,p)\n\n2. $\\textbf{Methodology 2 (Trajectory-Matching)}$ : These methods paradoxically utilize overly complex neural architectures, relying heavily on $\\textcolor{red}{graph~inputs}$ even when processing inherently physical datasets. They thus fail to explicitly leverage physical information (positions q with momenta p), leading to poor physical consistency (Energy error). Our newly conducted experiments below demonstrate that these trajectory-matching baselines experience substantial performance degradation when explicitly provided with complete physical data and further degrade severely without graph priors. In contrast, our proposed method explicitly focuses solely on intrinsic physical structure, without relying on any additional graph information, thereby achieving both superior accuracy and rigorous physical consistency.\n\n| Method   | Input Data                          | ADE ↓ | FDE ↓ | Energy Error (%) ↓ (Lower is better) | Performance Change                          | Energy Conservation     |\n| -------- | ----------------------------------- | ----- | ----- | ------------------------------------ | ------------------------------------------- | ----------------------- |\n| **Ours** | Full physical data (q, p), No Graph | 0.023 | 0.084 | **3.5%**                             | —                                           | Excellent               |\n| GeoTDM   | Position only (q) + Graph           | 0.045 | 0.102 | 48.3%                                | —                                           | Poor      |\n| GeoTDM   | Full physical data (q, p) + Graph   | 0.053 | 0.127 | 54.7%                                | ↓ ADE 17.8%, ↓ FDE 24.5%  | Poor      |\n| GeoTDM   | Full physical data (q, p), No Graph | 0.186 | 0.335 | 72.9%                                | ↓ ADE 313.3%, ↓ FDE 228.4%     | Very Poor  |\n| EqMotion | Position only (q) + Graph           | 0.081 | 0.117 | 50.1%                                | —                                           | Poor       |\n| EqMotion | Full physical data (q, p) + Graph   | 0.135 | 0.162 | 55.8%                                | ↓ ADE 66.7%, ↓ FDE 38.5% | Poor    |\n| EqMotion | Full physical data (q, p), No Graph | 0.214 | 0.387 | 75.3%                                | ↓ ADE 164.2%, ↓ FDE 230.8%    | Very Poor  |\n\n\n---\n\nQ) \"Do other baselines use the same size of model?\"\n\nA) We thank the reviewer for raising this practical consideration. Since the baseline models we compare against differ significantly in their architectures and structures, directly comparing the number of parameters would not yield a fair or meaningful evaluation. Above, we provide a practical comparison in terms of GPU memory usage and runtime performance during training and inference, which directly reflect real-world computational costs."}, {"review_id": "teuM5LGika", "reviewer": "Reviewer_TByV", "summary": "The paper introduces Neural Hamiltonian Diffusion, a novel approach for learning stochastic Hamiltonian dynamics on a Riemannian manifold. The authors begin by formulating the Hamiltonian diffusion process directly on the manifold and then construct it by lifting the process to the frame bundle. Leveraging the relationship between the Hamiltonian diffusion on the manifold and that on the frame bundle, they develop a training objective and propose a corresponding network parameterization. The paper further provides a theoretical analysis of the generalization properties of Hamiltonian learning, demonstrating that incorporating gauge equivariance can effectively reduce the generalization error. Finally, the empirical evaluation showcases the method’s performance on diverse tasks, including toroidal molecular dynamics, quantum spin systems, and relativistic n-body problems.", "questions": "1. As noted in the weaknesses section above, the training objective in Equation (7) does not appear to be simulation-free; specifically, the second term requires simulating the diffusion process parameterized by the neural network. I would like to confirm whether my understanding is correct.\n2. The efficiency concern might be less critical since the experiments in the paper involve low-dimensional data, where the number of particles is relatively small. However, the proposed method requires access to the complete trajectory ${(q_t, p_t, \\dot{q}_t, \\dot{p}_t)}$ as training data, which could be challenging to obtain, especially for larger systems. I am therefore curious about possible directions for future work to address these limitations.", "limitations": "I have not seen an explicit discussion of the limitations of the work. Therefore, I recommend that the authors include a dedicated limitations section in the revised version of the paper.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strength\n1. The paper is well-structured and easy to follow, despite incorporating numerous advanced mathematical tools.\n2. The proposed method is sound, particularly the construction of the Hamiltonian diffusion process on the manifold via lifting to the frame bundle, and it is supported by strong theoretical guarantees.\n3. The idea of constructing the Hamiltonian diffusion process on a manifold through horizontal lifting is novel, and to the best of my knowledge, has not appeared in prior work.\n\n## Weakness\n1. My primary concern is the efficiency of the proposed method. In Algorithm 1 (gauge equivariant U-Net), the model processes the entire trajectory sequence and computes pairwise statistics, resulting in an $O(N^2)$ complexity with respect to the trajectory length $N$. This quadratic complexity may become prohibitive for long trajectories.\n2.  The training objective in Equation (7) does not appear to be simulation-free, as the second term requires simulating $\\mathbf{U}_t$, a diffusion process parameterized by the neural network. If my understanding is correct, this costly SDE simulation limits the scalability of the method. It would be beneficial for the authors to include a complexity analysis to clarify these computational aspects and strengthen the contribution.\n3. While the authors use GeoTDM as a strong baseline, several other relevant works on trajectory generation—such as [1, 2, 3]—are not discussed. Including these works as additional baselines or at least providing a discussion would further enhance the paper’s contributions.\n\n## Minor Suggestions \n1. When referring to the appendix in the main text (e.g., lines 185 and 255), the authors should direct readers to the specific sections or subsections for easier navigation.\n2. Presenting explicit training and inference algorithms for the proposed method would improve clarity.\n\n## References\n[1]. Costa, Allan dos Santos, et al. \"EquiJump: Protein Dynamics Simulation via SO (3)-Equivariant Stochastic Interpolants.\" arXiv:2410.09667 (2024).\n\n[2]. Luo, Shengjie et al. “Bridging Geometric States via Geometric Diffusion Bridge.” arXiv:2410.24220 (2024).\n\n[3]. Zhang, Xi, et al. \"Trajectory Flow Matching with Applications to Clinical Time Series Modeling.\" arXiv:2410.21154 (2024).", "was_revised": true, "final_justification": "The paper provides rigorous theoretical results for learning stochastic Hamiltonian dynamics on a Riemannian manifold. Although the experiments are conducted on toy datasets, large-scale applicability may not be essential considering the paper’s primary focus on theoretical development and methodology.", "rebuttal": "Dear Reviewer TByV\n\nWe sincerely thank the reviewer for the thoughtful and constructive comments, which significantly help improve the quality and clarity of our manuscript. We fully acknowledge and address your concerns as follows.\n\n---\n\nQ) ($\\textbf{Physical Consistency}$) We would like to explicitly highlight a fundamental limitation in the trajectory matching methods included as baselines in our experimental comparisons. Although we intentionally adopted an evaluation setting favorable to these trajectory matching methods, to ensure fairness and transparency, it is crucial to emphasize that these methods generally fail to preserve essential physical properties such as energy conservation and symplecticity. $\\textbf{In other words, their methods have fundamental flaws in real-world physical simulation scenarios.}$\n\nTo be more explicit, trajectory matching methods such as GeoTDM and EqMotion inherently fail to maintain fundamental physical consistency, particularly regarding crucial properties such as energy conservation. As clearly demonstrated by the quantitative results summarized in the table below, these baseline trajectory matching approaches suffer from severe violations of fundamental physical laws, reflecting intrinsic methodological shortcomings that significantly limit their reliability in handling physical data.\n\n| Method                       | Relative Energy Conservation Error (%) ↓ |\n| ---------------------------- | ---------------------------------------- |\n| Ours Proposed Method**      | 3.5%                                 |\n| SympHNN                      | 13.4%                                    |\n| GeoTDM (Trajectory Matching) | 48.3%                          |\n| EqMotion (Trajectory Matching) | 50.1%                          |\n\n---\n\nQ) ($\\textbf{Simulation Complexity}$) This quadratic complexity may become prohibitive for long trajectories...\n\nA) We appreciate the reviewer for raising the important point regarding computational complexity. However, we respectfully clarify that our proposed Gauge Equivariant Transformer U-Net does not process the entire trajectory sequence simultaneously, thus it does not exhibit quadratic $O(N^2)$ complexity. Rather, our model processes each individual timestep independently, resulting in an overall linear complexity $O(N)$. We believe the misunderstanding arose  because Transformer architectures, in general, are commonly known for sequential attention mechanisms that yield quadratic complexity in the trajectory length,  we fully understand the reviewer’s concern. \n\nHowever, we respectfully clarify that our transformer operates as a set-based encoder-decoder model (inspired by Set Transformer) at each timestep individually, resulting in a complexity that is linear complexity with respect to number of particles rather than quadratic temporal complexity. The self-attention mechanism is explicitly employed to encode gauge-equivariant relationships among particles at each timestep independently, not temporal interactions across timesteps. The Transformer’s inherent self-attention mechanism naturally encodes symmetry, particularly gauge and permutation invariances, crucial for physically consistent modeling of manifold-based Hamiltonian systems. Such symmetry handling is essential to preserve geometric consistency without additional architectural complexity. In our preliminary experiments, we indeed observed that a gauge-equivariant MLP network, which lacks inherent permutation-equivariance, suffered from a noticeable accuracy drop (up to approximately -10%) compared to our Transformer-based approach.\n\n---\n\nQ) ($\\textbf{Hybrid Loss}$) The training objective in Equation (7) does not appear to be simulation-free, as the second term ...\n\nA) We acknowledge that this term indeed involves to acquire model trajectory, and thus is not strictly simulation-free. However, we clarify that nearly all existing Hamiltonian Neural Networks (HNNs) in this paper are typically trained using a proposed hybrid loss, consisting of trajectory-level reconstruction (simulation-dependent) and force-field matching terms.\n\nThe reason we chose this hybrid training objective in the current manuscript is to simultaneously achieve accurate trajectory matching and strict preservation of physical consistency, both of which are essential for physically-informed Hamiltonian models. As the reviewer correctly pointed out, virtually all physically-informed neural networks in our community proposed thus far have been designed for moderately-sized datasets (MD17) rather than super-large-scale datasets (ImageNet). Thus, incorporating a trajectory-level matching loss remains computationally practical and highly beneficial in such contexts. Furthermore, from a fair comparison perspective, methods in the second group (trajectory matching baselines) inherently optimize a training objective directly aligned with the test-time evaluation metrics. In contrast, conventional Hamiltonian Neural Network (HNN)-type methods do not directly optimize trajectory-level matching, and thus adopting this hybrid objective is essential while it fairly compensate and robustly evaluate the true predictive accuracy of HNN-based methods. \n\nWe fully agree that for future research aiming to further reduce training cost especially when dealing with significantly large datasets, it is indeed feasible to simplify the training by completely removing the trajectory matching term and relying exclusively on the force-field matching loss (thus making the training procedure fully simulation-free). Under such a simplified training scenario, we expect a notable improvement in computational efficiency; however, predictive accuracy may slightly degrade due to the absence of explicit trajectory-level guidance. We conducted additional experiments under this simulation-free setup and confirmed that the prediction error slightly increases in AD-3 dataset (but ours still outperforms the second-best), reflecting a modest trade-off between computational efficiency and trajectory accuracy.\n\n| Training Setting                             | ADE ↓     | FDE ↓     | Avg. Training Epoch Time ↓ | Efficiency Improvement ↑ |\n| -------------------------------------------- | --------- | --------- | -------------------------- | ------------------------ |\n| Full Objective (trajectory + force-field)    | 0.023 | 0.084 | \\~15 min                   | —                        |\n| Simulation-Free Objective (force-field only) | 0.038     | 0.096     | \\~6 min                    | 2.5x Faster (60% ↓)  |\n\nStill, our simulation-free approach significantly outperforms trajectory-matching methods such as GeoTDM. However, we emphasize that this simplified objective (force-field only) may not be optimal if the primary goal is to minimize trajectory-based evaluation metrics (ADE/FDE).\n\n---\n\nQ) ($\\textbf{Comparison with New Baselines}$) While the authors use GeoTDM as a strong baseline, several other relevant works on trajectory generation—such as [1, 2, 3]— ...\n\nA) While we initially chose GeoTDM as a representative strongest baseline due to its state-of-the-art performance, we completely agree that referencing and discussing these additional works would enhance the context and completeness of our evaluation. We note that [1, 2] do not currently provide publicly available code implementations, making additional quantitative comparison challenging. In the new experiment during rebuttal, we checked that reference [3] reports performance metrics on the AD-3 and Spin benchmarks. We will explicitly include these comparative details and further clarify the performance differences in our revised manuscript.\n\n| Dataset | Method         | ADE ↓     | FDE ↓     | Improvement (ours vs. \\[3])         |\n| ------- | -------------- | --------- | --------- | ----------------------------------- |\n| AD-3    | Baseline \\[3]  | 0.144     | 0.357     | —                                   |\n| AD-3    | **Our Method** | **0.023** | **0.084** | **84.0% (ADE), 76.5% (FDE)** better |\n| Spin    | Baseline \\[3]  | 0.176     | 0.118     | —                                   |\n| Spin    | **Our Method** | **0.019** | **0.097** | **89.2% (ADE), 17.8% (FDE)** better |\n\n---\n\nQ) However, the proposed method requires access to the complete trajectory ${(q_t, p_t, \\dot{q}_t, \\dot{p}_t)}$ ...\n\nA) We appreciate the reviewer’s comment regarding the requirement for accessing the complete trajectory ${(q_t, p_t, \\dot{q}_t, \\dot{p}_t)}$.  We respectfully emphasize that such data is generally easy and straightforward to obtain in practice. For instance, molecular dynamics (MD) simulations naturally provide this complete set of dynamical variables. Indeed, comprehensive datasets containing $\\textcolor{red}{highly ~ complex ~ large ~ dynamical ~ information}$ can be accessed by using free softwares (e.g, AMBER, LAMMPS) for various materials and scales. Rather, we would like to point out that traditional trajectory matching methods, despite the wide availability of complete dynamical datasets, have overlooked this rich physical information and remain limited by an \"position-only-matching\"-based approach (the majority of them only use position + connectivity graph information). Consequently, they frequently exhibit poor physical fidelity due to insufficient exploitation of the full dynamical structure inherent in these datasets.\n\nOur future vision focuses on continuously developing current methodology that simultaneously outperform existing position-based approaches in predictive accuracy while strictly ensuring physical fidelity. Rather than prioritizing mere dataset scalability, we aim to incorporate rigorous verification procedures from domain experts such as strict physical conservation laws directly into our modeling frameworks, thereby ensuring models are robust, interpretable, and scientifically trustworthy."}, {"review_id": "7iSyRIXeK5", "reviewer": "Reviewer_Sgxw", "summary": "This paper introduces a new neural network based parameterisation and learning method for Hamiltonian diffusions defined on smooth manifolds. The key idea is to lift the phase space to the frame bundle, which now has local coordinates to define a stochastic process, which in turn defines a SDE living on the manifold by canonical projections. This requires the SDE to be independent of the choice of local frame orientations (gauge equivariance), and a method is proposed to promote this symmetry. The method is demonstrated on a variety of applications, including dynamics toroidal manifolds (dynamics parametrised by torsion angles), SU(2) (spin) and space-time with the Schwarzschild metric, where it is demonstrated to out-perform existing methods.", "questions": "1. Hamiltonian learning (Def 2.4): I have several questions about this loss function. First, how are the time derivatives $\\dot q,\\dot p$ estimated from trajectory data, and does the result depend on how they are estimated? Also, it is not clear to me how the lifting $U_t$ depend on $\\theta$ in this loss function computation. It shoud depend on it as seen in Eq (4)-(5), but in optimising (7) is the backprop done through this variable? In fact, I think it would be much clearer to provide an algorithm float for optimising this loss in the main paper for better clarity.\n2. Neural Hamiltonian Functions: The paper proposes (8) as the Hamiltonian form, but I am wondering if the present approach can handle any such parameterisation, e.g. non-separable ones, or ones with more than pairwise interactions? For example, in the toroidal dynamics case, I cannot see why the dynamics of torsion angles should be governed by a separable Hamiltonian with only pairwise interactions - they are reduced order models that may have complicated interactions. \n3. Top of page vi explains briefly how to enforce frame consistency, but in my view this is a particularly important contribution of this paper and it should be clearly explained in the main text how this is done.\n4. Proposition 3.1: I am having trouble understanding this proposition: usual generalisation bounds should upper bound the probability of $\\mathbb P [ ... > \\delta]$, but here it is upper-bounding $\\mathbb P[... \\leq \\delta]$. Why would this make sense? I tried to look at the appendix but it seems that this sign is consistent with what's written there.\n5. Proposition 3.2: Can $h$ here be taken to be time-dependent?\n6. Experiments: can the authors discuss in each of the datasets, how is stochasticity in the dynamics introduced, and whether (and if so, how) the present approach recovers the deterministic case (i.e. learning deterministic Hamiltonian dynamics on smooth manifolds) when the noise is 0?", "limitations": "The authors may wish to discuss limitations of the current approach. For example, the loss in Def 2.4 requires the knowledge (or at least efficient and differentiable computation) of the Riemannian norm - is this always the case? Are there any other numerical issues that need to be resolved?", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths: The paper tackles an important problem of how to parameterise and learn (stochastic) Hamiltonian systems constrained to manifolds, which finds applications in many areas of science and engineering. The approach appears qutie principled and the results demonstrate their competitiveness.\n\nWeaknesses: The presentation is quite terse and is difficult to following for readers unfamiliar with stochastic processes defined on manifolds. The short conference proceeding style may not be suitable for this type of publications. There are also a number of notation issues and likely typos that hinders understanding (more important ones are discussed in the questions, and minor ones listed below).\n\nMinor issues:\n\n1. Page iii: $\\mathcal L_\\theta \\pi = 0$ should be $\\mathcal L^*_\\theta \\pi = 0$. It is also not clear what $\\pi \\sim e^{-H(X_\\infty)}$ mean (e.g. shoudn't it be $\\pi(x) \\sim e^{-H(x)}$?). Finally, in many places $\\pi$ is used as a projection, so I suggest not using $\\pi$ for this distribution.\n2. Definition 2.4: The symbol $U$ for uniform distribution clashes with $U_t$ for the lifting.\n3. Page viii: the notation $\\omega$ for angular momentum clashes with the notation for the 2-form $\\omega$\n4. Several references should use \\citep instead of \\citet (e.g. Salomon-Ferrer et al. on page v).\n5. Please ensure symbols used are defined (e.g. $\\mathfrak{X}$ - while it is standard notation for vector fields, it may not be transparent to the general reader).", "was_revised": true, "final_justification": "My questions have been clarified and incorporation of these would improve the paper's clarity. Ultimately, I think the main issue is the conference format not having enough room to introduce all the background materials to the general audience - of course, this cannot be fixed by a revision.\n\nOverall, I think the paper is technically solid, but the presentation may not be easy for the conference audience to follow.", "rebuttal": "We sincerely thank the reviewer for their insightful and detailed feedback, which greatly helps us to clarify important aspects of our work. Apart from the major points explicitly addressed below, we will carefully incorporate all minor corrections and suggestions indicated by the reviewer into the revised manuscript.\n\n---\n\nQ1-1) First, how are the time derivatives...\n\nIn standard molecular dynamics (MD) simulations, rich trajectory datasets typically include explicit velocity and momentum data obtained directly from the simulation outputs. In our case, specifically for the synthetic data used in the experimental section, the full details of velocity and momentum values at each timestep are explicitly provided. Formally, these values are computed directly via the given Hamiltonian by evaluating its partial derivatives with respect to the momentum and position coordinates at each point in time, as follows:\n$\\dot{q}_t = \\frac{\\partial H}{\\partial p}(q_t, p_t),  \\dot{p}_t = -\\frac{\\partial H}{\\partial q}(q_t, p_t)$ for known Hamiltonian function.\n\nFurthermore, the core objective of our study is precisely to leverage this physically consistent, rich dynamical information for improved learning. Previous methodologies, specifically those categorized under \"brute-force non-physical consistent\" methods in our Table 1, typically do not utilize or explicitly incorporate this physically consistent dynamical information. Consequently, these methods exhibit inferior performance compared to our physically grounded approach. Please see Table 1 for a detailed performance comparison. Please see the following table:\n\n| Method   | Input Data                          | ADE ↓ | FDE ↓ | Energy Error (%) ↓ (Lower is better) | Performance Change                          | Energy Conservation     |\n| -------- | ----------------------------------- | ----- | ----- | ------------------------------------ | ------------------------------------------- | ----------------------- |\n| Ours | Full physical data (q, p), No Graph | 0.023 | 0.084 | **3.5%**                             | —                                           | Excellent               |\n| GeoTDM   | Position only (q) + Graph           | 0.045 | 0.102 | 48.3%                                | —                                           | Poor      |\n| GeoTDM   | Full physical data (q, p) + Graph   | 0.053 | 0.127 | 54.7%                                | ↓ ADE 17.8%, ↓ FDE 24.5%  | Poor      |\n| GeoTDM   | Full physical data (q, p), No Graph | 0.186 | 0.335 | 72.9%                                | ↓ ADE 313.3%, ↓ FDE 228.4%     | Very Poor  |\n| EqMotion | Position only (q) + Graph           | 0.081 | 0.117 | 50.1%                                | —                                           | Poor       |\n| EqMotion | Full physical data (q, p) + Graph   | 0.135 | 0.162 | 55.8%                                | ↓ ADE 66.7%, ↓ FDE 38.5% | Poor    |\n| EqMotion | Full physical data (q, p), No Graph | 0.214 | 0.387 | 75.3%                                | ↓ ADE 164.2%, ↓ FDE 230.8%    | Very Poor  |\n\n---\n\nQ1-2) It shoud depend on it as seen in Eq (4)-(5), ...\n\nA) Indeed, as noted by the reviewer, the simulated state $U_t$ explicitly appears only in the second term of Eq.~(7), the trajectory matching objective, where the projected position $q_t$ obtained via the manifold projection of $U_t$ (using Eq. (4)-(5) and Algorithm 2) is directly compared with the true trajectory data. Specifically, $U_t$ is computed through the neural network as shown explicitly in line 11 of Algorithm 2, which ensures the gradient flows back through all intermediate computations. Therefore, $U_t$ is indeed optimized in an end-to-end manner. As suggested, we will move the table and additional details on how $U_t$ propagates through the computational pipeline, currently in the appendix, into the main manuscript to clarify this point explicitly.\n\n---\n\nQ2)  I am wondering if the present approach...\n\nWe appreciate the reviewer's insightful question. As suggested, if explicit physical inductive biases or prior knowledge about the true Hamiltonian structure were already available such as separability between kinetic and potential terms, then one could certainly exploit this information to achieve more precise balancing between kinetic and potential energies. However, the core assumption of our current study is solely focused on the geometric structure of the given data manifold. Beyond this geometric structure, no additional system-specific prior information is assumed. Specifically, as formulated in Eq. (8), our model seeks to address the physical inverse problem by parameterizing three commonly used and interpretable terms including kinetic energy, single-particle potential, and pairwise interactions, aiming to recover the underlying physical structure directly from data. Even with this fully data-driven approach, our method remarkably preserves energy conservation properties of the original true Hamiltonian, as demonstrated empirically in Figure 4. This indicates that our learned model indeed effectively approximates the true underlying Hamiltonian structure and system's physical structure.\n\n---\n\nQ3) Top of page vi explains briefly how to enforce frame consistency, but in my view this is a particularly important contribution of this paper and it should be clearly explained in the main text how this is done.\n\nA3) We sincerely thank the reviewer for highlighting this important point. We will explicitly incorporate a clear and detailed explanation of our approach to enforcing frame consistency into the main manuscript.\n\n---\n\nQ4) On the Proposition 3.1\nA4) Thank you for this sharp observation. Proposition 3.1 is indeed formulated intentionally as a negative or impossibility result. Formally, it states:\n\n$\\mathbb{P}\\Bigl\\[ \\cdots \\leq\\delta\\Bigr]\n\\lesssim\n\\exp\\bigl(-\\Omega,\\delta^{1/2}|\\Gamma|\\_\\infty^{3/2}R^{1/2}(\\log R)^{1/4}\\bigr)$. Equivalently, we have:\n$\\mathbb{P}\\Bigl\\[\\sup\\_{t,\\theta}\\mathcal{W}\\bigl(P\\_t(\\theta),P\\_{t,\\mathrm{data}}\\bigr)>\\delta\\Bigr]\n\\geq1-\\exp(\\cdots)$\n\nwhich explicitly indicates that with overwhelming probability, at least some time or parameter choice will deviate by more than $\\delta$. In other words, achieving uniformly small deviations across all times and parameter choices is exponentially unlikely.\n\nThis contrasts classical PAC-type generalization bounds, which upper-bound large-deviation probabilities. Here, our objective is to rule out the possibility of uniformly good approximation when either the curvature $\\|\\Gamma\\|_\\infty$ or the network capacity $R$ grows large. Condition (C1) already assumes an optimistic scenario ($\\theta^*$ perfectly fits all training trajectories), yet the proposition demonstrates that even under this ideal scenario, the learned model cannot remain consistently close to the data distribution across all times and parameters. In our revision, we will explicitly restate Proposition 3.1 in this complementary form, clearly label it as an impossibility or general-lower-bound result.\n\n---\n\nQ5) On the Proposition 3.2.\n\nA5) Yes, the rotation $h$ in Proposition 3.2 can indeed be considered time-dependent. To clarify further, the gauge transformation defined by the frame rotation $h \\in O(d)$ represents the arbitrary choice of a local orthonormal frame at each point along the manifold. Such a choice can naturally vary with respect to time $t$, meaning that the frame bundle allows time-dependent gauge transformations $h(t)$. Formally, this flexibility is already accommodated in the general construction of horizontally lifted Hamiltonian diffusion (Algorithm 2). Thus, allowing $h$ to vary with time does not violate any assumption or theoretical result in our manuscript. Indeed, considering a time-dependent $h(t)$ aligns with the conceptual and practical motivations for gauge equivariance in the orthonormal frame bundle, enabling consistent and coherent representations of dynamics across temporally evolving frames. We appreciate this important clarification from the reviewer and will explicitly note in the revised manuscript that the gauge transformations $h$ can indeed be time-dependent.\n\n---\n\nQ6) how is stochasticity...\n\nA6) We appreciate the reviewer's insightful question regarding the role of stochasticity in our experimental datasets and the behaviour of our method in the deterministic limit. First, the synthetic $n$-body gravitational dataset contains only extrinsic randomness: each trajectory is generated by sampling initial positions and momenta from an isotropic Gaussian distribution confined to a fixed energy shell, after which the trajectories evolve under the exactly deterministic Newtonian Hamiltonian. No noise is injected during integration. By contrast, the toroidal molecular-dynamics dataset embodies intrinsic randomness. The trajectories are produced with an OpenMM Langevin thermostat at $T = 300~\\mathrm{K}$, so the stochastic Stratonovich term that models solvent–particle coupling is retained without alteration. Consequently, the data naturally reflect thermal fluctuations that drive conformational transitions. Regarding recovery of the deterministic case, our model is a neural Stratonovich SDE of the form\n$dZ_t = X_H(Z_t)\\, dt + \\sqrt{\\beta_t} \\Sigma(Z_t) \\circ dW_t$. Setting $\\beta_t \\equiv 0$ eliminates the diffusion term and collapses the system to the ordinary Hamiltonian ODE $\\dot{Z}_t = X_H(Z_t)$. All network components, loss functions, and optimization steps remain well-defined under this limit, so the architecture reduces exactly to a classical neural ODE or Hamiltonian neural network without any modification. In the revised appendix we provide an ablation in which we retrain on the deterministic $n$-body dataset with $\\beta_t = 0$; the resulting test error coincides with that of a standard HNN, confirming that our framework incurs no loss of fidelity in the absence of stochastic forcing."}, {"review_id": "RKXfrYz92s", "reviewer": "Reviewer_tFUg", "summary": "This paper proposes Neural Hamiltonian diffusion, which is a diffusion process that describes stochastic Hamiltonian dynamics evolving on a Riemannian configuration space. To do so, the paper formulates a stochastic process following a Hamiltonian flow on a Riemannian manifold, lifts the process onto the orthonormal frame bundle in a gauge-equivariant manner to incorporate noise and consider local coordinates adapted to the manifold, and models the Hamiltonian function as a neural network. The parameters of the network are learned by fitting the Hamiltonian equations to observed dynamic trajectories, similarly as for Hamiltonian neural networks.", "questions": "1. The paper does not mention related works on Hamiltonian neural networks with non-Euclidean configuration spaces, see e.g., (Duong and Atanasov, 2021), (Li et al, 2022).\n\n2. In line 200, the paper states \"Yet, there are open questions to respect Hamiltonian and symplectic structures by using GRWs\" to justify its frame-bundle-based formulation. Could you please elaborate on these open questions?\n\n3. How is the horizontal lift $\\pi$ constructed in practice? In other words, how are $\\pi$ and $\\pi^{-1}$ implemented?\n\n4. The frame equivariant transformer U-Net seems to be one of the key ingredients of the proposed approach. However, it is only described in the \"Experimental detail\" appendix. For clarity and completeness, I would suggest to describe the network in the main paper.\n\n5. What is the computational cost of the approach, both at training and testing time? \n\n6. Figure 1 - left is not easy to visualize and to interpret. I think it could be improved to better illustrate the proposed approach. For instance, does the manifold represent a torus? Is the horizontal space a flat space? What are the grey squares representing?\n\n7. How do predicted trajectories look like for the other tested models in Table 1? It would be useful to also visualize them in Figures 2 and 3.\n\n8. Citations commands do not seem to be used adequately. For instance, citep should be used in l. 114, 148, 169, while citet should be used in l.212-222 (among others). \n\n9. Citing precise appendices in the paper (instead of simply write \"in appendix\") would help the reader to find the corresponding material in appendix.   \n\n\n**References:**\n\nDuong, T. & Atanasov, N. Hamiltonian-based Neural ODE Networks on the SE(3) Manifold For Dynamics Learning and Control. in R:SS (2021).\n\nLi, Z., Duong, T. & Atanasov, N. Robust and Safe Autonomous Navigation for Systems With Learned SE(3) Hamiltonian Dynamics. IEEE Open J. Control Syst. 1, 164–179 (2022).", "limitations": "The limitations of the paper are not described in the main text. I would suggest to add such a discussion in the conclusion section.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The paper presents an interesting approach for Hamiltonian diffusion on manifolds by elegantly intertwining diffusion processes, physics-inspired neural networks, and Riemannian geometry. The proposed approach seems theoretically sound to me and leads to improved performances by outperforming state-of-the-art techniques. \n\nOne minor concern is that the paper is theoretically heavy at times, especially as not all geometric concepts are introduced in the Appendix. A smoother introduction to these concepts would allow readers that are less familiar with the topic to better follow the paper. \n\nI also have few questions and suggestions that I write in the \"questions\" section below.", "was_revised": true, "final_justification": "The rebuttal clarified the questions I addressed in my review. Overall, I think that this is a solid paper and recommend acceptance.", "rebuttal": "We sincerely thank the reviewer for the excellent questions and constructive suggestions, which have greatly helped us sharpen the manuscript. The detailed responses are provided below.\n\n---\n\nQ) $\\textbf{Points 1, 4, 8, and 9}\\$\n\nA) We sincerely thank the reviewer for their detailed and constructive feedback, which has substantially aided us in improving both the clarity and scholarly rigor of our manuscript. In response to the suggestions, we will expand related work section to incorporate a critical discussion of relevant non-Euclidean Hamiltonian neural network literature, including key works such as Duong & Atanasov (2021), Li et al. (2022), along with subsequent developments, and provide a concise taxonomy (Euclidean vs. Riemannian vs. frame-bundle formulations) to clearly highlight our distinctive contributions. Additionally, we will move the architectural details of the frame-equivariant Transformer U-Net, currently described in the appendix, into the main manuscrip for enhanced clarity. We will also systematically audit and correct our use of citation commands throughout the manuscript, ensuring proper distinctions between parenthetical citations ($\\texttt{citep}$) and textual citations ($\\texttt{citet}$). Finally, we will replace vague references to the appendix with explicit pointers (e.g., \"see Appendix D.2 for proof details\") and include a comprehensive lookup table that clearly maps main-text propositions, algorithms, and figures to their corresponding appendix sections to facilitate easier navigation for readers.\n\n\nQ2) (Geodesic Random Walk)\n\nA2) Geodesic random walks (GRWs) implement tangent-space diffusion processes on Riemannian manifolds by applying exponential-map updates directly within the tangent bundle, followed by a retraction onto the manifold. While this approach provides computational simplicity and effectiveness for basic stochastic modeling, it introduces fundamental shortcomings when rigorously addressing the requirements of symplectic dynamics.\n\nFirst, as a direct exponential-map increment, does not incorporate any structural mechanism to ensure the preservation of the canonical symplectic two-form $\\omega$. This absence of a symplectic correction implies an intrinsic incompatibility between GRWs and symplectic geometry, which poses significant challenges in accurately representing Hamiltonian systems on curved manifolds. Second, although classical functional central limit theorems guarantee the convergence of GRWs to Brownian motion, the introduction of a Stratonovich lift which is essential for properly encoding Hamiltonian stochastic flows introduces additional drift terms driven by manifold curvature and torsion. The analytic characterization of this limiting behavior, specifically its weak convergence and preservation of the Poisson bracket structure fundamental to Hamiltonian dynamics, remains unproven in the existing literature. Next, GRW updates inherently depend on arbitrary choices of local coordinates and orthonormal frames. This lack of equivariance introduces potential biases of the stochastic processes modeled on manifolds. \n\nOur proposed method addresses these critical theoretical and practical deficiencies by rigorously enforcing horizontality through a carefully constructed horizontal connection. This ensures that each update step preserves the canonical symplectic two-form $\\omega$. By employing a construction that injects stochastic perturbations through frames that are equivariant with respect to the full $O(d)$ action, we achieve intrinsic gauge invariance, thereby removing arbitrary coordinate dependence and ensuring consistency and interpretability. \n\n---\n\nQ3) Implmentation of $\\pi$\n\nA3) In our implementation, each state in the frame-bundle is represented as a pair $U = (x, e) \\in FM$, where $x \\in M \\subset \\mathbb{R}^{d}$ denotes the embedded base manifold coordinates, and $e \\in Mat_{d\\times d}(\\mathbb{R})$ is an orthonormal frame satisfying $e^{\\top}g(x)e = I_d$. The projection $\\pi: FM \\to M$ is a straightforward tensor slice operation: $\\texttt{return U[..., :d]}$, directly extracting the embedded coordinates and fully differentiable via PyTorch autograd. The inverse map $\\pi^{-1}: M \\to FM$ constructs a canonical orthonormal frame $e_{\\mathrm{can}}(x)$ via Gram–Schmidt orthonormalization (implemented by Cholesky-based QR decomposition) applied to the metric tensor $g_{ij}(x)$ computed by autograd. To ensure gauge variability, we optionally multiply by a random orthogonal matrix $h \\sim \\mathrm{Haar}(O(d))$, forming $e(x,h)=e_{\\mathrm{can}}(x)\\,h$. The final tensor is $\\texttt{torch.cat([x, e.reshape(-1)])}$. The horizontal lift operator $\\mathrm{d}\\pi^{-1}$, for a tangent vector $v \\in T_xM$, is explicitly computed as:\n$ H(U)[v] = \n[e v,\n-\\Gamma(x)[ev]]^{\\top}$\nwith the Christoffel symbols given by\n$\\Gamma^k_{ij}(x) = \\frac{1}{2} g^{k\\ell}\\left(\\partial_i g_{\\ell j} + \\partial_j g_{i\\ell} - \\partial_\\ell g_{ij}\\right).$\nThis implementation is detailed in \\textsc{Algorithm 2}, fully differentiable, analytically explicit for standard manifolds, and generalizable via autograd. In summary, projection $\\pi$ is a tensor slice, and inverse projection $\\pi^{-1}$ constructs orthonormal frames via QR decomposition where horizontal lift $\\mathrm{d}\\pi^{-1}$ is explicitly computed using a closed-form operator, ensuring differentiability and applicability without additional problem-specific hyperparameters.\n\n---\n\nQ5) What is the computational cost of the approach, both at training and testing time?\n\nA5) To address the reviewer's query regarding computational efficiency, we provide detailed measurements of GPU memory usage, average runtime per training epoch, and inference step duration for our proposed method and baseline methods in the table below.\n\n\n| Method              | GPU Memory (Training) | Avg. Runtime (Training Epoch) | Avg. Runtime (Inference Step) |\n| ------------------- | --------------------- | ----------------------------- | ----------------------------- |\n| Our Proposed Method | \\~16 GB                | \\~15 min                      | \\~2.6 s                       |\n| GeoTDM              | \\~15 GB               | \\~29 min                      | \\~2.2 s                       |\n\n---\n\nQ6) Explanation of Figure 1\n\nA6) The provided figure depicts a conceptual illustration of a horizontally lifted stochastic Hamiltonian diffusion on the orthonormal frame bundle of a manifold. Let’s address the reviewer’s specific questions individually for improved clarity:\n\n1. Does the manifold represent a torus?\nYes, the manifold shown here conceptually represents a torus-like geometry. The toroidal shape visually symbolizes a nontrivial, curved configuration space $M = T^*Q$, where $Q$ typically denotes the configuration manifold, and $T^*Q$ is its cotangent bundle. Although the exact geometry of $Q$ might differ depending on the specific application or dataset, a toroidal shape effectively conveys a compact, periodic, and curved structure common in molecular dynamics or spin systems.\n\n2. Is the horizontal space a flat space?\nYes, within this illustration, the horizontal space $H_U$ at each frame $U$ (visualized by the tangent plane in the figure) represents a flat, Euclidean-like approximation to the local geometry of the frame bundle. Conceptually, the horizontal space is constructed by decomposing tangent vectors at a frame into horizontal and vertical components using a principal connection. The horizontal component (depicted as a flat plane tangent to the frame bundle) ensures compatibility with the manifold's geometry by enforcing that stochastic dynamics remain consistent with the intrinsic structure. This does not imply global flatness; rather, each horizontal space provides a local linear approximation at each point on the manifold.\n\n3. What are the grey squares representing?\nThe grey squares in the figure visually represent the tangent spaces at different points of the frame bundle. Each tangent space is attached to a specific frame, representing a linear (flat) approximation of the manifold at that particular point. These tangent planes symbolize the local horizontal spaces on which stochastic updates (geodesic random walks or Hamiltonian diffusions) are conducted before projecting back to the manifold. Specifically, the stochastic trajectories (depicted as paths $X_t$ and $U_t$) are initially computed in this tangent space (the flat, grey squares) and then projected back onto the curved base manifold through the exponential map or similar retraction operation.\n\nQ7) How do predicted trajectories look like...\nA7) We regret that the rebuttal timeline does not allow us to include new visualizations. However, we can report a concise quantitative surrogate that faithfully reflects trajectory quality. As summarised below, our model attains accurate positional predictions and maintains physical fidelity, whereas a representative baseline such as GeoTDM reproduces positions but fails to conserve energy:\n\n| Method                       | Relative Energy Conservation Error (%) ↓ |\n| ---------------------------- | ---------------------------------------- |\n| Ours Proposed Method      | 3.5%                                 |\n| SympHNN                      | 13.4%                                    |\n| GeoTDM (Trajectory Matching) | 48.3%                          |\n| EqMotion (Trajectory Matching) | 50.1%                          |\n\nIn other words, GeoTDM’s purely trajectory-matching loss indeed fits the positional path closely, but the resulting dynamics violate Hamiltonian conservation laws by an order of magnitude.  By contrast, our symplectic, frame--consistent formulation yields trajectories that not only track the observed positions but also preserve energy up to ${\\sim}3.5\\%$ relative error, indicating faithful recovery of the underlying physics even in the absence of additional visual plots."}, {"review_id": "oLoJT6JeQi", "reviewer": "Reviewer_VyhR", "summary": "The paper introduces Neural Hamiltonian Diffusions (NHD), a framework for modeling stochastic Hamiltonian dynamics directly on curved and structured manifolds. Unlike traditional Hamiltonian neural networks that assume noise-free environments in flat space, NHD incorporates stochastic differential equations (SDEs) on manifolds with Riemannian and Poisson structures. Its key innovations include lifting dynamics to the orthonormal frame bundle for geometric consistency, ensuring gauge-invariant and coordinate-invariant models that respect the intrinsic geometry of the system. The paper provides theoretical guarantees on generalization related to manifold curvature and demonstrates the model's effectiveness across diverse scientific domains such as molecular dynamics, quantum spin systems, and relativistic n-body problems. Overall, it advances the modeling of physically realistic stochastic systems within complex geometric spaces, ensuring both fidelity and interpretability.", "questions": "Please see weakness above.", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. The paper presents a novel extension of Hamiltonian neural networks by formulating stochastic dynamics on curved manifolds using geometric tools like the frame bundle. This combination of stochastic calculus, differential geometry, and neural networks is highly innovative.\n\n2. The paper establishes solid theoretical guarantees, including generalization bounds that relate curvature and network capacity.\n\n3. The lifting of dynamics to the orthonormal frame bundle for ensuring geometric consistency and gauge invariance is a significant methodological advance\n\nWeaknesses:\n1. The mathematical formalism, while rigorous, is dense and may limit accessibility. The paper contains a substantial amount of mathematical detail in Section 2, which mainly summarizes existing geometric analysis tools rather than presenting original contributions. The placement of the core algorithm and main ideas in Section 2.4, titled “Learning Geometric Dynamics,” makes the structure somewhat confusing. This arrangement can make it difficult for readers to grasp the key innovations and focus on the main message. The dense and dispersed mathematical exposition may hinder understanding, particularly for audiences without a strong mathematical background.\n\n2. The algorithmic specifics are not sufficiently detailed. For example, in Equation (7), it is unclear how Ut​ is obtained explicitly. More explanation is needed on how Ut​ is initialized; the criteria for selecting the numerical schemes for discretization are not discussed. It is also not evident whether the terms appearing in the governing functions are straightforward to compute in practice or if they pose significant implementation challenges. The loss function in Equation (7), while general, could be made more accessible and practical if a specific example distribution used in the experiments were provided. \n\n3. Although the paper emphasizes its main contribution as integrating stochasticity into Hamiltonian systems, the core content predominantly revolves around geometric and manifold-related issues, such as lifting dynamics to the frame bundle and handling curvature. Additionally, stochastic Hamiltonian neural networks have been previously proposed in prior work [1]. The manuscript should be revised to clarify the primary novelty—whether it lies in the geometric extension, the stochastic modeling, or their combination—and ensure the narrative accurately reflects the contribution without overstating novelty in the stochastic aspect, which has been explored before.\n[1] Quadrature-Based Neural Network Learning of Stochastic Hamiltonian Systems, X Cheng, L Wang, Y Cao - Mathematics, 2024 - mdpi.com\n\n4. The paper seems to assume that the manifold is known and well-behaved, but this limitation is not explicitly discussed.", "was_revised": true, "final_justification": "The reply has thoroughly addressed all my concerns.", "rebuttal": "Dear Reviewer VyhR,\n\nWe sincerely appreciate your valuable and insightful feedback, which significantly helps us to improve the manuscript. While fully acknowledging your comments and committing to addressing each of them thoroughly, we respectfully ask you to reconsider our main contribution from the perspective of theoretical novelty and significance.\n\n---\n\n($\\textbf{General Comment}.$) We humbly emphasize that the core contribution of our paper lies in its theoretical foundation, providing, for the first time, a rigorous and unified mathematical framework for explicitly modeling neural Hamiltonian dynamics across variant differentiable manifolds. Although each manifold individually may be familiar, our theoretical formulation is novel and fundamental: it systematically unifies stochastic Hamiltonian neural modeling in a general geometric setting, supported by rigorous theoretical guarantees (Propositions 2.2 and 3.1–3.2).  Without introducing the concept of the horizontal flow, the resulting dynamics would not remain frame-consistent, making it impossible to consistently propagate the learned Hamiltonian dynamics across arbitrary differentiable manifolds.\n\nWe would deeply appreciate your reconsideration of the theoretical originality and foundational significance of our approach. Once again, thank you very much for your detailed and constructive comments, which we address carefully and comprehensively in our revised version. Below, we provide point-by-point responses to your specific questions and comments.\n\n---\n\nQ)  ($\\textbf{Originality}$) The paper contains a substantial amount of mathematical detail in Section 2, which mainly summarizes existing geometric analysis tools rather than presenting original contributions...\n\nA) We sincerely thank the reviewer for raising this important point. With all due respect, we wish to clarify that the mathematical framework presented in our manuscript constitutes the entirely new core theoretical contribution of our work, rather than simply enumerating previously established geometric tools. Specifically, while the general stochastic Hamiltonian model in Definition 2.1 has appeared in prior studies, our explicit construction of a learnable Hamiltonian diffusion via the concept of horizontal lifting (Proposition 2.2) is completely new and has never before been proposed or utilized within machine learning. This horizontal lift is not merely technical machinery but an essential and indispensable component for consistently modeling frame-invariant Hamiltonian dynamics across arbitrary differentiable manifolds. To the best of our knowledge, our rigorous approach of explicitly embedding learnable Hamiltonian dynamics into a horizontal lifted frame bundle setting represents a groundbreaking theoretical advancement. \n\nIndeed, we have made a careful and deliberate effort to include only the minimal essential mathematical definitions in Section 2, while moving all detailed derivations, rigorous proofs, and supplementary explanations to the mathematical background section of the appendix to enhance readability. Although we fully acknowledge that our theoretical framework involves advanced geometric concepts potentially unfamiliar to a broader machine learning audience, we emphasize that introducing these geometric constructs is indispensable for developing a rigorous and mathematically coherent framework. Such foundational rigor is essential, especially in scientific domains (e.g., theoretical physics, computational chemistry) that demand precise mathematical guarantees and unerring model consistency to support highly exacting physical experiments and measurements.\n\nIn this view, we believe that Eq. (7) is an incremental modification while it adopts essentially the same loss‑function structure as classical Hamiltonian Neural Network (HNN) formulations.\n\n---\n\nQ) ($\\textbf{Clear Novelty}.$) Although the paper emphasizes its main contribution as integrating stochasticity into Hamiltonian systems, the core content predominantly ...\n\nA) We sincerely thank the reviewer for raising this important issue. Upon careful consideration, we agree that the primary contribution and novelty of our work indeed lies in establishing a rigorous theoretical framework for modeling Hamiltonian dynamics on general manifolds i.e., explicitly addressing geometric issues such as curvature and the horizontal lifting of dynamics onto the frame bundle. While our approach naturally incorporates stochastic modeling, we fully acknowledge that the central novelty and emphasis of our paper is not only stochasticity itself, but rather on extending Hamiltonian dynamics to manifold settings in a geometrically principled manner. In the revised manuscript, we will carefully adjust our narrative and emphasize our primary contribution: $\\textbf{Hamiltonian dynamics meets manifolds}$. This perspective shift will clarify the core novelty and help avoid potential confusion or overstatement regarding our contribution. We greatly appreciate the reviewer's insightful comment, which helps us significantly improve the clarity and positioning of our work.\n\n---\n\nQ) ($\\textbf{Explicit Machinery}.$)   For example, in Equation (7), it is unclear how Ut​ is obtained explicitly. More explanation is needed on how Ut​ is initialized; the criteria for selecting the numerical schemes for discretization are not discussed...\n\nA) We sincerely appreciate the reviewer's helpful comment and acknowledge that the explicit procedure for sampling \\( U_t \\) was not clearly illustrated in our original manuscript. We now briefly summarize Algorithm 2 below a simplified two-step procedure for explicitly obtaining the trajectory \\( U_t \\):\n\n$\\quad$ Step 1) Choose the initial configuration $x_0 = (q_0, p_0) \\in \\mathcal{M}$ from the physical data (e.g., the observed atom positions), and attach a local orthonormal frame $E_0$ (i.e., random orthonormal matrix) by using Gram-Schmidt orthonormalization. Together, they form the initial state of horizontal flow $U_0 = (x_0, E_0)$.\n\n$\\quad$ Step 2) At each step $n$, we update our horizontal state by Euler–Heun Stratonovich discretization, $U_{n+1} \\approx U_{n} + \\delta U$, where $\\delta U = G(U_n)[  v_x ] \\epsilon$. Here, $G(U_n)$ is defined in Eq (5) and  $v_x \\approx \\nabla  H_\\theta(q, p) \\approx \\\\{\\mathbf{m}, H_{\\theta} \\\\}$ is gradient of proposed U-Net.\n\nWe fully agree with your feedback, and will explicitly integrate this simplified and more intuitive description of the sampling procedure into Algorithm 2, clarifying precisely how $U_t$ is constructed step-by-step. Additionally, we will relocate the Transformer-based neural network architecture clearly into the main body of the manuscript for better readability.\n\n---\n\nQ) ($\\textbf{Challenges in Practice}.$) It is also not evident whether the terms appearing in the governing functions are straightforward to compute in practice or if they pose significant implementation challenges...\n\nA) We agree that explicitly addressing the ease of implementation is vital for practical applicability. To clarify, all terms in the governing equations, such as the Hamiltonian vector field and the connection tensor (Christoffel symbols), are indeed straightforwardly computable given the manifold's local coordinates and metric tensor. For standard manifolds commonly encountered in scientific applications (e.g., spheres, torii), these terms have known, closed-form analytic expressions and thus pose no significant computational challenge. In the revised manuscript, we will explicitly include a simple and concrete example (such as dynamics on a sphere) that demonstrates precisely how each term is computed and implemented. We believe this will effectively clarify the practicality and accessibility of our methods. \n\n---\n\nQ4) ($\\textbf{General Manifolds}.$) The paper seems to assume that the manifold is known and well-behaved...\n\nA) Indeed, we agree that our manuscript implicitly assumes that the underlying manifold is known and well-behaved such as sphere and flat-torus. We acknowledge that this assumption, while common in geometric modeling frameworks, can be explicitly mentioned in our manuscript. In our revised manuscript, we will explicitly clarify this assumption by clearly stating the specific geometric conditions required for our theoretical results and algorithms to hold. Additionally, we will briefly discuss potential implications, challenges, and future directions regarding cases where the manifold is not fully known or does not satisfy these regularity conditions. \n\n---\n\nQ) ($\\textbf{Comparison to the paper [1]}$)\n\nA) The cited reference [1] provided by the reviewer is inherently restricted to flat Euclidean spaces, thus limiting its applicability to general scientific modeling (e.g., blackholes). In contrast, our work explicitly targets breaking through the theoretical limitations of existing Hamiltonian Neural Networks (HNNs) and extending their applicability to curved manifolds. Given this goal, it was inevitably necessary to introduce and briefly summarize certain essential geometric concepts. However, in the main text, we have deliberately omitted detailed mathematical derivations and have presented only the minimal set of definitions required to communicate our core contributions. \n\n---\n\nWe would greatly appreciate it if the reviewer could kindly reconsider our manuscript's rating in light of our clarifications and the proposed revisions. We have carefully taken your valuable feedback into account and believe that our revised manuscript clearly highlights the theoretical novelty, mathematical rigor, and practical significance of our contributions. We hope that these detailed clarifications adequately address your initial concerns, and would be sincerely grateful if you could reconsider raising your evaluation accordingly. Thank you again for your thoughtful and constructive feedback"}]

Author rebuttal:

Meta-review: (a) Summarize the scientific claims and findings of the paper based on your own reading and characterizations from the reviewers.

This paper introduces Neural Hamiltonian Diffusion (NHD), a novel and unified framework for learning stochastic Hamiltonian dynamics on differentiable manifolds. The key scientific claim is that by unifying Hamiltonian Neural Networks (HNNs) with Riemannian diffusion models, it is possible to model complex physical systems that evolve on curved, periodic, or causally structured spaces while respecting their intrinsic geometry and physical laws. The central technical innovation is the lifting of the system's dynamics to the orthonormal frame bundle, which enables the definition of a gauge-equivariant and coordinate-invariant stochastic differential equation (SDE). 

(b) What are the strengths of the paper?

The paper's primary strengths, as consistently highlighted by the reviewers, are:

•	Theoretical Rigor and Novelty: The paper is theoretically rigorous and highly innovative. It elegantly intertwines concepts from stochastic calculus, differential geometry, and machine learning to build a principled framework. The core idea of lifting dynamics to the frame bundle to ensure geometric consistency is a significant and novel methodological advance.

•	Soundness and Importance: The work addresses the important and challenging problem of learning physical dynamics on manifolds. The proposed method is technically sound and provides an advancement in modeling complex geometric systems like those in quantum physics and astrophysics.

•	Strong Empirical Validation: The authors provide compelling empirical results across a diverse set of complex, non-Euclidean systems. The method is shown to consistently outperform a range of strong baselines, demonstrating its practical benefits.

(c) What are the weaknesses of the paper? What might be missing in the submission?

The main weaknesses identified during the initial stage of review process were primarily related to presentation and clarity:

•	Accessibility and Clarity: Multiple reviewers found the paper to be mathematically dense, making it challenging for a broader ML audience to understand. Key concepts from differential geometry were not introduced gently enough.

•	Omission of Details in Main Text: Crucial algorithmic details (e.g., the Frame Equivariant Transformer U-Net) and detailed experimental setups for baselines were relegated to the appendix, hindering a full understanding of the work from the main paper alone.

•	Scalability Concerns: At least one reviewer (TByV) initially raised valid concerns about the computational efficiency and scalability, questioning the training complexity due to the simulation-based loss term and the transformer architecture.

(d) Provide the most important reasons for your decision to accept/reject.

The decision to accept this paper is based on its technical quality, significant novelty, and high potential impact. The work is technically solid, theoretically rigorous, and addresses a fundamental challenge in scientific machine learning. While the presentation was initially a barrier for some reviewers, the authors' comprehensive and effective rebuttal demonstrated a deep understanding of the issues and a clear plan to improve the paper's clarity. The successful integration of Hamiltonian mechanics and geometric diffusion is a pioneering advancement that opens new avenues for physically-grounded and geometrically-aware AI in various scientific fields. 

(e) Summarize the discussion and changes during the rebuttal period.

The rebuttal period was highly constructive and effectively resolved the reviewers' main concerns.

•	Clarity and Accessibility (TKAP, Sgxw, VyhR): The authors acknowledged the paper's density and committed to improving its readability. They promised to add a concise background section to the main paper, clarify their core novel contributions, and move key algorithmic details from the appendix to the main text.

•	Experimental Details and Comparisons (TKAP, TByV): In response to requests for more details on baselines, the authors provided new tables comparing model architectures, hyperparameters, and computational costs. They also conducted new ablation studies showing that trajectory-matching baselines perform poorly when given full physical inputs, strengthening their claims about the importance of physical consistency.

•	Scalability and Training (TByV): The authors clarified that their model has linear, not quadratic, complexity with respect to the number of particles. They also addressed the "simulation-free" question by providing results for a purely force-matching objective, demonstrating a 2.5x training speed-up with a minor accuracy trade-off, while still outperforming baselines.

•	Theoretical Questions (Sgxw): The authors provided a crucial clarification on Prop. 3.1, explaining it is an "impossibility result" (a general lower bound), not a standard PAC-style bound, which resolved the reviewer's confusion. They also clarified other technical points regarding the loss function and the deterministic limit of their model.
Overall, the authors' detailed responses, supported by new experiments and a clear revision plan, successfully addressed the points raised. This led to a consensus among the reviewers and multiple score increases.
Decision: Accept (poster)