Title: Neural Hamiltonian Diffusions for Modeling Structured Geometric Dynamics
Venue: NeurIPS 2025
Accepted: yes
Keywords: geometric deep learning

Riemannian diffusion

Hamiltonian diffusion
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\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.\n\nQUESTIONS:\nPlease see weakness above.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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).\n\nLIMITATIONS:\nThe limitations of the paper are not described in the main text. I would suggest to add such a discussion in the conclusion section.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: 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).\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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?", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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?\n\nQUESTIONS:\ncheck 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.\n\nLIMITATIONS:\nno, check weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}