Towards Cross Domain Generalization of Hamiltonian Representation via Meta Learning

TL;DR

This work models a graph neural network with a graph neural network and employs a meta learning algorithm to enable the model to gain experience over a distribution of systems and make it adapt to new physics, thereby overcoming the limitations of system-specific models.

Abstract

Recent advances in deep learning for physics have focused on discovering shared representations of target systems by incorporating physics priors or inductive biases into neural networks. While effective, these methods are limited to the system domain, where the type of system remains consistent and thus cannot ensure the adaptation to new, or unseen physical systems governed by different laws. For instance, a neural network trained on a mass-spring system cannot guarantee accurate predictions for the behavior of a two-body system or any other system with different physical laws. In this work, we take a significant leap forward by targeting cross domain generalization within the field of Hamiltonian dynamics. We model our system with a graph neural network (GNN) and employ a meta learning algorithm to enable the model to gain experience over a distribution of systems and make it adapt to new physics. Our approach aims to learn a unified Hamiltonian representation that is generalizable across multiple system domains, thereby overcoming the limitations of system-specific models. We demonstrate that the meta-trained model captures the generalized Hamiltonian representation that is consistent across different physical domains. Overall, through the use of meta learning, we offer a framework that achieves cross domain generalization, providing a step towards a unified model for understanding a wide array of dynamical systems via deep learning.

Figures (15)

• Figure 1: The relative error throughout the time rollout for the meta-trained, pre-trained, and random-initialized model at adaptation step 50 for the predicted coordinates (top) and energy (bottom). The solid line and the shaded area each represent the average and the standard error of 10 runs.
• Figure 2: The relative error across the adaptation step for the meta-trained, pre-trained, and random-initialized model at adaptation step 50 for the predicted coordinates (top) and energy (bottom). The solid line and the shaded area each represent the average and the standard error of 10 runs.
• Figure 3: System dynamics through time rollout 0 to 10, generated from the meta-trained, pre-trained, and random-initialized model each after 50 adaptation steps. Each row (a) to (d) corresponds to the predictions with mass-spring, pendulum, Hénon-Heiles, and the magnetic mirror system respectively.
• Figure 4: The measured CKAs in the last layer between the model before adapting (i.e. right after finishing meta-training or pre-training), and during the adaptation up to 1000 steps. The CKA values are transformed to 1-CKA to make a clearer visualization. The solid line and the shaded area each represent the average and the standard error of 10 runs.
• Figure S.1: The relative error for the meta-trained, pre-trained, and random-initialized model with respect to the predicted coordinates (top) and energy (bottom) using different integrators. The solid line and the shaded area each represent the average and the standard error of 10 runs.