Symplectic Neural Networks for learning non-separable Hamiltonians

Hamiltonian Neural Networks (HNNs) integrate physical priors into neural models by learning a system's Hamiltonian, improving generalization and sample efficiency. Identifying the system Hamiltonian from noisy observations of state variables is a challenging task. Moreover, for simulations to faithfully reflect the long-term behavior of Hamiltonian systems, especially energy conservation, it's essential to use symplectic integrators, which are specifically designed to preserve the system's geometric structure over time. This fidelity comes at a cost: implicit symplectic integrators are computationally intensive and make backpropagation through the ODE solver non-trivial. However, leveraging the self-adjoint property of symplectic integrators, computing gradients w.r.t. Neural Net parameters reduces to solving adjoint equations in the backward pass. In our work, we explore an alternate method of HNN training under noisy observation of trajectories with our HNN model with a symplectic implicit integrator embedded in the loop. Computationally, a predictor-corrector based ODE solver and fixed point iteration help efficient generation of gradient updates. We showcase the numerical advantage, in experiments, in system identification and energy preservation on a range of non-separable, chaotic systems and the efficient computation and memory complexity of our method.