Hamiltonian Neural Networks for a Spring-Mass System

What it does

Models a spring-mass system with a Hamiltonian Neural Network (HNN) and compares it with a standard (baseline) neural network on accuracy and physical consistency. Both learn the motion from 50 noisy measurements of a single period; then both are asked to predict five periods. The baseline's energy drifts away from its true value; the HNN keeps it.

Part of my research internship with Prof. Saurav Sengupta, IIT (ISM) Dhanbad.

The cell below trains both networks and then integrates them, as in the notebook. Watch the arrows (each model's learned vector field) line up with the circle while the losses fall, then the two trajectories go round. With 200 hidden units it is the notebook's own run: the same noisy data and starting weights (from the notebook's seeds), trained the same way. By default it replays a recorded run of this code, which takes about 12 seconds; choose Train live in your browser to compute it on your device instead (a few minutes at 200 units; 64 units, starting from this site's own random numbers, take 20 to 40 seconds).

Hamiltonian mechanics

For a block on an ideal spring, let $q$ be its position and $p = m\dot{q}$ its momentum. The Hamiltonian $\mathcal{H}(q, p)$ is the quantity whose derivatives give the motion:

$$\dot{q} = \frac{\partial \mathcal{H}}{\partial p}, \qquad \dot{p} = -\frac{\partial \mathcal{H}}{\partial q}$$

Defined this way, $\mathcal{H}$ never changes along a trajectory, and it has units of energy: it is the system's total energy, $\mathcal{H} = \tfrac{1}{2}kq^{2} + \dfrac{p^{2}}{2m}$, which an ideal spring conserves.

Data

With $k = 2$ N/m, $m = 0.5$ kg and amplitude 1, the motion is $q = \sin\omega t$ with $\omega = \sqrt{k/m} = 2$, so one period is $\pi$ s and the orbit in the $(q, p)$ plane is the circle $q^{2} + p^{2} = 1$ (energy 1).

• 50 points over one period: inputs $(q, p)$, targets $(\dot{q}, \dot{p})$.
• Noise to test robustness: 2.5% of the standard deviation added to $(q, p)$ and 5% to $(\dot{q}, \dot{p})$.

Two networks

Both are fully connected with two hidden layers of 200 tanh units, Xavier-initialised.

The baseline NN maps $(q, p) \to (\dot{q}, \dot{p})$ directly and is trained on

$$\mathcal{L}_{\text{baseline}} = \left\| \dot{q}_{\text{pred}} - \dot{q} \right\|^{2} + \left\| \dot{p}_{\text{pred}} - \dot{p} \right\|^{2}$$

It is a black box for the mapping: nothing tells it that energy should be conserved.

The Hamiltonian NN maps $(q, p)$ to a single number $\mathcal{H}_\theta$, and its predictions are that number's derivatives, computed with automatic differentiation: $\dot{q} = \partial \mathcal{H}_\theta / \partial p$, $\dot{p} = -\partial \mathcal{H}_\theta / \partial q$. Its loss is

$$\mathcal{L}_{\text{HNN}} = \left\| \frac{\partial \mathcal{H}_\theta}{\partial p} - \dot{q} \right\|^{2} + \left\| \frac{\partial \mathcal{H}_\theta}{\partial q} + \dot{p} \right\|^{2}$$

So it learns the total energy itself, and any motion it predicts conserves its own $\mathcal{H}_\theta$ by construction.

Both train with Adam (learning rate $10^{-3}$, full batch) for 4,000 epochs. The final losses in the notebook: 0.0042 for the baseline, 0.0022 for the HNN.

Five periods with Euler

To turn predicted derivatives into a trajectory, both models are integrated with the explicit Euler scheme, $q_{t+1} = q_t + \dot{q}_t\,\Delta t$ and $p_{t+1} = p_t + \dot{p}_t\,\Delta t$, from the first point of a 50,000-point test set, for 5 periods (50,000 steps): five times longer than anything in the training data.

Results

The baseline NN fails to conserve the total energy of the system, while the HNN is very consistent in conserving it. In the notebook's run, the baseline's energy grows to about three times its true value over the five periods and its trajectory spirals outwards; the HNN's stays within a few percent of 1, on the true orbit.

The cell reproduces that run from the notebook's starting point: final losses 0.0041 and 0.0022, and the HNN ends five periods about 0.01 from where the notebook's does. Its baseline drifts further (energy about 6): that drift amplifies tiny rounding differences, and PyTorch itself, run on a CPU from the same start, ends at about 5.

How far and which way the baseline drifts depends on the random start (with 64 units the cell's baseline spirals inwards instead); that it drifts doesn't. The exact equations, integrated with the same Euler steps, gain about 2% over five periods: that rise comes from the scheme, not from a network.