[Paper] From Kepler to Newton: Inductive Biases Guide Learned World Models in Transformers

Source: arXiv - 2602.06923v1

Overview

This paper asks whether a generic, off‑the‑shelf Transformer can move beyond raw prediction and actually learn the physical laws that generate the data it sees. By injecting three modest inductive biases—spatial smoothness, training‑time stability, and temporal locality—the authors show that Transformers can evolve from mere curve‑fitters into models that recover Kepler’s ellipses and even Newton’s force law, opening a path toward AI‑driven scientific discovery.

Key Contributions

• Minimal bias recipe: Identifies three lightweight architectural/practical tweaks (continuous regression, noisy‑context training, and a short attention window) that together enable Transformers to learn true physical abstractions.
• Keplerian world model: Demonstrates that with spatial smoothness and stability alone, a Transformer can infer the elliptical orbits of planets from trajectory data.
• Newtonian force discovery: Shows that adding temporal locality forces the model to abandon pure curve‑fitting and recover the underlying inverse‑square force law.
• Systematic ablation: Provides a clear experimental breakdown of how each bias contributes, contrasting against prior “AI Physicist” methods that rely on heavy domain‑specific priors.
• Open‑source baseline: Releases code and pretrained models, making it easy for developers to reproduce and extend the experiments.

Methodology

1. Data generation: Simulated planetary motion under Newtonian gravity, producing sequences of (x, y) positions sampled at regular time steps.
2. Model architecture: Starts from a vanilla Transformer encoder‑decoder.
3. Inductive bias 1 – Spatial smoothness: Instead of predicting discrete tokens, the model outputs continuous coordinates via a regression head, encouraging smooth trajectories.
4. Inductive bias 2 – Stability via noisy contexts: During training, the input context window is corrupted with Gaussian noise, teaching the network to be robust to small prediction errors that would otherwise compound over long roll‑outs.
5. Inductive bias 3 – Temporal locality: The self‑attention mask is limited to a fixed‑size recent window (e.g., the last 5 timesteps), embodying the physical assumption that the next state depends primarily on the current state, not the entire history.
6. Training & evaluation: Models are trained to minimize mean‑squared error on next‑step prediction. After training, the learned internal representations are probed (e.g., via linear read‑outs) to see if they encode orbital parameters or force vectors.

Results & Findings

• Baseline Transformers (no biases) achieve low prediction error but their internal states do not correspond to any recognizable physical quantities.
• Adding spatial smoothness + stability reduces error further and, crucially, the latent vectors linearly map to orbital elements (semi‑major axis, eccentricity), effectively learning Kepler’s law of ellipses.
• Including temporal locality pushes the model to represent the underlying acceleration field; a simple linear probe recovers an inverse‑square radial force, matching Newton’s law of gravitation.
• Ablation studies confirm that removing any one bias degrades either the physical interpretability or the predictive performance, underscoring their complementary roles.

Practical Implications

• Automated model discovery: Engineers can embed these biases into generic sequence models to let them infer governing dynamics from sensor logs (e.g., robotics, aerospace telemetry) without hand‑crafting physics equations.
• Improved simulation tools: Learned world models can serve as fast surrogates for high‑fidelity simulators, enabling rapid prototyping and control‑loop testing.
• Explainable AI: By forcing models to expose interpretable latent variables (orbital parameters, forces), developers gain insight into why a model predicts what it does—valuable for safety‑critical systems.
• Cross‑domain transfer: The bias recipe is domain‑agnostic; it could be applied to economics (discovering market dynamics), biology (inferring kinetic laws), or any time‑series where underlying causal laws exist.

Limitations & Future Work

• Synthetic data focus: Experiments are limited to idealized, noise‑free planetary simulations; real‑world measurements (e.g., atmospheric drag, sensor drift) may pose additional challenges.
• Scalability of attention window: Temporal locality works well for low‑order dynamics but may need adaptation for systems with longer memory (e.g., fluid dynamics).
• Interpretability pipeline: The current linear probing step is simple; more robust methods for extracting symbolic equations from latent spaces remain an open problem.
• Future directions: Extending the framework to multi‑agent or relativistic regimes, integrating with reinforcement learning agents that act on the learned world model, and exploring automated bias selection via meta‑learning.

Authors

• Ziming Liu
• Sophia Sanborn
• Surya Ganguli
• Andreas Tolias

Paper Information

• arXiv ID: 2602.06923v1
• Categories: cs.LG, cs.AI, physics.class-ph
• Published: February 6, 2026
• PDF: Download PDF