Source-linked AI summary
Neural Relational Inference for Interacting Systems
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling, Richard Zemel
TL;DR
Interacting systems are difficult to model when observations omit their underlying relations. NRI uses a variational auto-encoder with a discrete latent interaction graph and GNN-based dynamics, recovering physical interactions unsupervised while predicting real motion and sports dynamics. The approach is limited because it is trained for static graphs, requiring dynamic re-estimation when interactions evolve.
Problem
Interacting-system observations often provide individual trajectories without the underlying interactions or dynamical model.
Method
NRI jointly infers discrete latent edge types and learns system dynamics with a GNN-based variational auto-encoder from observational trajectories.
Results
NRI accurately recovers ground-truth interaction graphs in physical simulations and predicts long-term dynamics in motion-capture and sports-tracking data with interpretable edge types.
Takeaways & Limitations
The model provides an unsupervised way to learn interaction structure alongside dynamics across physical, motion-capture, and sports data.
Takeaways & Limitations
NRI is trained for static interaction graphs, so evolving interactions require dynamically re-estimating the latent graph during application.
Abstract
from arXiv · showhide
Interacting systems are prevalent in nature, from dynamical systems in physics to complex societal dynamics. The interplay of components can give rise to complex behavior, which can often be explained using a simple model of the system's constituent parts. In this work, we introduce the neural relational inference (NRI) model: an unsupervised model that learns to infer interactions while simultaneously learning the dynamics purely from observational data. Our model takes the form of a variational auto-encoder, in which the latent code represents the underlying interaction graph and the reconstruction is based on graph neural networks. In experiments on simulated physical systems, we show that our NRI model can accurately recover ground-truth interactions in an unsupervised manner. We further demonstrate that we can find an interpretable structure and predict complex dynamics in real motion capture and sports tracking data.
1. Introduction
Interacting systems produce complex dynamics, yet observations often provide only individual trajectories without the underlying interactions or dynamical model. NRI addresses this gap by inferring an explicit latent interaction structure while learning system dynamics unsupervised.
- Observations of interacting systems often contain individual trajectories but omit the underlying interactions and dynamical model.
- Manual interaction annotation is possible for selected tasks but is tedious, motivating unsupervised learning of interactions shared across tasks.
- NRI learns dynamics with a GNN over a discrete latent graph while inferring explicit edge types that cluster interactions.
- The probabilistic formulation incorporates graph-structure priors such as sparsity in a principled manner.
- NRI recovers ground-truth physical interactions accurately without supervision and predicts many future time steps in motion-capture and NBA tracking data using few edge types.
2. Background: Graph Neural Networks
Graph neural networks operate on graph-structured data by passing local messages between nodes and edges. Their message-passing operations transform node and edge representations through concatenation, aggregation, and neural-network updates.
- GNNs operate directly on graph-structured data by passing local messages between connected components.
- The formulation supports directed graphs, while an undirected edge can be represented by assigning opposite directed edges.
- A node-to-edge operation combines embeddings of nodes connected by an edge to construct an edge representation.
- An edge-to-node operation aggregates embeddings from incoming edges to update a node representation.
- Node- and edge-specific neural networks follow these representation changes, enabling multiple rounds of message passing.
3. Neural Relational Inference Model
NRI jointly infers a discrete interaction graph from object trajectories and learns the system dynamics conditioned on that graph. Its encoder–decoder design uses graph neural networks, probabilistic edge inference, and multi-step prediction to make interactions identifiable in forecasting.
- Model overview: NRI jointly trains an encoder that infers interactions from trajectories and a decoder that learns dynamics conditioned on the inferred interaction graph.The model is formalized as a variational autoencoder with a discrete latent graph.
- Encoder: The encoder represents each pairwise interaction as a discrete edge type and predicts its categorical distribution with a GNN over the fully connected graph.Multiple message-passing rounds allow edge representations to incorporate information from other nodes.
- Decoder: The decoder predicts future trajectories with a GNN conditioned on the latent graph, using separate neural networks for different edge types.For physical simulations, the decoder exploits the Markovian state defined by location and velocity.
- Sampling: The discrete latent variables are trained with a continuous Concrete relaxation, which enables reparameterized gradient estimates during optimization.The relaxation approaches one-hot categorical samples as the temperature τ approaches zero.
- Avoiding degenerate decoders: The decoder predicts multiple future steps and uses one MLP per edge type so that interaction information cannot be ignored as easily as in single-step prediction.The experiments use M = 10 predicted steps before reinserting the correct previous input.
- Recurrent decoder: For non-Markovian applications, NRI augments GNN message passing with a GRU-based recurrent decoder that models dependence on the full trajectory history.The recurrent update receives aggregated messages, the current input, and the previous hidden state.
4. Related Work
Prior graph-based approaches either assume a graph, infer interactions implicitly, or use soft attention, whereas NRI explicitly infers a latent graph with discrete relation types. This supports interpretable structures and principled priors such as sparsity.
- Graph-based dynamics: Earlier graph neural network models learned dynamics with a known graph or represented interactions implicitly rather than inferring an explicit interaction structure.This distinction separates NRI from methods that assume connectivity or encode interactions only within message functions.
- Attention and latent graphs: Attention-based GNNs focus message aggregation on selected neighbors, while NRI performs explicit inference over a latent graph structure.The latent graph supports discrete relation types rather than only soft attention weights.
- Interpretability and priors: NRI’s latent graph formulation allows prior beliefs such as sparsity and yields an interpretable discrete structure with multiple relation types.The paper contrasts this formulation with interaction-inference work in causal reasoning and computational neuroscience.
5. Experiments
Experiments evaluate NRI on simulated physical systems and real motion, sports, and trajectory-prediction settings, comparing interaction recovery and future-state prediction against supervised, correlation-based, LSTM, and graph baselines. NRI recovers interactions accurately, predicts long-term dynamics, and benefits from dynamically re-evaluated graphs for motion capture.
- Physics simulations: NRI is evaluated on springs, charged particles, and Kuramoto oscillators, with known interactions enabling assessment of both dynamics learning and edge-type recovery.The simulations include 50k training examples and provide ground-truth interaction structure for evaluation.
- Physics simulations: The unsupervised NRI model greatly surpasses correlation-based and other baselines for interaction recovery, approaching the supervised benchmark on springs.Results are reported as averages over five runs with standard error.
- Physics simulations: NRI predicts the next 20 unseen time steps after observing 49 steps and outperforms LSTM baselines for long-term prediction.For Kuramoto dynamics, LSTMs continue short-term waveform shapes smoothly but fail to model long-term interacting-system dynamics.
- Physics simulations: 82.6% edge accuracy can still match the true-graph NRI MSE for charged particles because distant particles have weak effects on future predictions.The result illustrates that interaction-recovery accuracy and predictive accuracy can diverge when some interactions have small dynamical effects.
- Motion capture data: On motion capture, dynamic graph re-evaluation improves prediction over a static baseline, while NRI outperforms the fully connected graph and both outperform LSTM baselines.The learned graph depends on the phase of motion, motivating encoder re-evaluation at every test time step.
- Pick and Roll NBA data: On basketball tracking, NRI outperforms an LSTM and matches the full graph while learning edge types that separate ball and ball-handler relations from other player connections.This latent structure identifies a semantic organization centered on the ball and ball handler in the pick-and-roll play.
6. Conclusion
NRI jointly infers relational structure and learns dynamics, recovering interaction graphs in simulations and modeling physical, motion-tracking, and sports data. Its static-graph training limits direct handling of interactions that change over time.
- NRI simultaneously infers relational structure and learns the dynamical model of an interacting system.
- The model accurately recovers ground-truth interaction graphs in physical simulations without supervision.
- NRI models interacting physical, motion-tracking, and sports-analytics dynamics at high precision while learning reasonably interpretable edge types.
- For dynamic interactions, a trained NRI model can re-estimate the latent graph, but training remains limited to static graphs.The paper identifies explicit modeling of dynamic latent interactions during training as future work.
- A simple LSTM captures oscillator waveform shape but fails to model interaction-driven phase dynamics, whereas NRI closely captures the coupled behavior.The LSTM often goes out of sync by up to half a wavelength in later steps.
A.2. Spring simulation variants
Additional spring-simulation tests show that unsupervised NRI can identify an empty graph and distinguish three edge types, while visualizations extend evaluation to motion capture and basketball trajectories.
- Spring simulation variants: 98.4% accuracy identifies “no interaction” edges when NRI is tested on 1,000 simulations of five non-interacting particles.
- Spring simulation variants: 99.2% accuracy identifies the correct edge types when a third spring type with coupling constant k_ij = 0.5 is added.All three edge types are sampled with equal probability in this variant.
- Motion capture visualizations: Motion-capture predictions capture the overall movement form with high precision across 30 future steps conditioned on 49 observed steps.Small errors can accumulate over long prediction sequences, such as a misplaced toe node in frame 30.
- NBA visualizations: NBA trajectory predictions remain semantically reasonable even when they do not match the true future path.One example predicts a defender going over a screen rather than switching, which is plausible despite differing from reality.
B. Simulation data
The simulations cover spring-connected particles, charged particles, and phase-coupled oscillators, using differentiable or externally integrated dynamics to generate trajectories and support model evaluation.
- Spring model: Spring simulations use 5 or 10 particles in a 2D box, with each particle pair randomly connected by a spring with probability 0.5.Connected particles interact through Hooke’s-law forces.
- Spring model: A differentiable leapfrog implementation serves as a ground-truth decoder and enables comparison with a learned decoder.
- Charged particles model: Charged-particle simulations assign positive or negative charges, with every pair interacting through attraction or repulsion.The interaction strength can be weak when particles remain far apart.
- Charged particles model: Force clipping stabilizes charged-particle generation but makes the trajectories not exactly physically accurate and can create zero gradients for the simulation decoder.
- Kuramoto model: The Kuramoto simulation models nonlinear phase-coupled oscillators whose dynamics depend on intrinsic frequencies and coupling strengths.One-dimensional trajectories are generated by solving the differential equation with fourth-order Runge–Kutta integration at step size 0.01.
C.1. Vectorized implementation
The implementation vectorizes graph message passing and provides MLP, CNN, and recurrent encoder-decoder variants for predicting interacting-system dynamics.
- Vectorized implementation: Node-to-edge and edge-to-node message passing can be evaluated in parallel for efficient vectorized computation.For large sparse graphs, sparse-dense multiplication can yield an O(E) algorithm.
- Vectorized implementation: Message-passing matrices encode incoming and outgoing edge connectivity and can be precomputed when the graph is static.
- MLP encoder: The MLP encoder uses a two-layer 256-dimensional network with batch normalization, dropout, and ELU activations.Its node2edge module concatenates receiver and sender features, while edge2node sums incoming edge features.
- CNN encoder: The CNN encoder uses attention with 1D convolutions to handle changing trajectory sizes and interactions that are strong for only part of a trajectory.
- Decoder variants: The decoder implementations include MLP single-step prediction and a recurrent GRU-style update for sequential dynamics.
D. Experiment details
The experiments use jointly trained interaction-inference and dynamics models, with physical simulations evaluating edge recovery and path prediction against correlation, recurrent, supervised, and graph-based baselines.
- Training procedure: Testing uses categorical discrete edge types, while training uses the concrete distribution with τ = 0.5 and teacher forcing every 10th time step.
- Datasets: Physical simulation datasets contain 50k training instances and 10k validation and test instances, with 49-step training trajectories and longer test trajectories.
- Baselines: The edge recovery experiments compare NRI with static, LSTM, correlation-based, simulator-based, and supervised baselines.
- Baselines: Correlation baselines infer interaction edges by thresholding trajectory correlations and selecting the threshold variant with the best ground-truth graph match.
- Path prediction: The path prediction experiments evaluate NRI alongside recurrent and other baseline models.
D.2. Motion capture data experiments
The motion-capture experiments use normalized tracked-point trajectories, an MLP encoder, and an RNN decoder, with shortened training and constrained recurrent-baseline capacity.
- Data: The motion-capture dataset contains 8,063 frames for 31 tracked points, with 49-frame training and validation samples and 99-frame test samples.
- Data: Features are normalized to maximum absolute value 1, and the reported test results use the last 50 frames of each 99-frame test sample.
- Model: The model uses an MLP encoder and an RNN decoder because the motion-capture dynamics are not Markovian.
- Training and baselines: Training runs for 200 epochs with batch size 8, while the joint LSTM baseline uses 128 hidden units and batch size 1 to fit GPU memory.
D.3. NBA experiments
The NBA experiments model pick-and-roll trajectories whose dynamics change across the event, so evaluation predicts only the final eight steps after encoding the first 17.
- Data: Each NBA example is a 25-step pick-and-roll trajectory subsampled from 25-frames-per-second SportVU data.
- Evaluation: Because behavior differs before, during, and after the pick-and-roll, testing encodes the first 17 time steps and predicts the final 8.
- Training design: Training an autoencoder with matching 17- or 25-step input and prediction lengths creates a train-test mismatch that produces poor predictive performance.
- Model design: The NBA model uses a CNN encoder, an RNN decoder, and two edge types; ten or twenty edge types performed worse on validation data, probably because of over-fitting.