Goal: learn the continuous-time dynamics of three gravitating bodies from trajectory data, in a way that generalizes to unseen initial conditions by building the physical symmetries of the problem (rotation, translation, particle permutation) directly into the model.
The three-body problem is one of the oldest unsolved problems in physics — three gravitating bodies create chaotic, unpredictable trajectories that defy closed-form solutions.
We use Neural ODEs to learn the dynamics directly from trajectory data. Crucially, we build in the symmetries of the problem: rotational equivariance, translational invariance, and energy conservation.
By constraining the neural network architecture to respect these physical symmetries, we achieve dramatically better generalization on unseen initial conditions compared to unconstrained models.
The network $f_\theta$ takes invariant inputs — the pairwise distances $d_1, d_2, d_3$, the velocity inner products, and the velocity norms — and outputs a vector of invariant scalar coefficients. The $k$-th output is the coefficient $a_k$:
Distances and inner products kill translations and rotations. Permutation equivariance over planets $2$ and $3$ is enforced by symmetrizing the network's evaluation over the two orderings.
The new position and velocity are linear combinations of the raw equivariant vectors using these invariant coefficients:
Combining equivariant quantities with invariant weights yields an equivariant update. Translation equivariance of $p_{\text{new}}$ requires $\sum_{k=1}^{3} a_k = 1$ (enforced via softmax); translation invariance of $v_{\text{new}}$ requires $\sum_{k=7}^{9} a_k = 0$. Conservation of momentum is restored by recentering: $v_i \leftarrow v_i - \bar v + c$.
Slight additional modifications ensure permutation invariance, not included here.