How Do Transformers "Do" Physics? Investigating the Simple Harmonic Oscillator

TL;DR

The paper investigates how transformers model physics by using the SHO as a controlled testbed and introducing a four-criterion framework to detect internal intermediates corresponding to a specific modeling method, anchored by an in-context linear regression proxy $Y = wX$. It then applies this framework to both undamped and damped SHO trajectories, finding strong correlational and causal evidence that the matrix exponential method underpins the transformer’s predictions for the undamped case, with weaker support for linear multistep and Taylor expansion, and inconclusive results for the damped case. The significance lies in providing a scalable, mechanistic approach to reveal potential world-model computations inside transformers and in outlining how this framework can extend to higher-dimensional linear and nonlinear systems. The SHO analysis highlights how precise internal representations can align with known numerical solvers, offering a path toward understanding the physics-guided capabilities of AI systems.

Abstract

How do transformers model physics? Do transformers model systems with interpretable analytical solutions, or do they create "alien physics" that are difficult for humans to decipher? We take a step in demystifying this larger puzzle by investigating the simple harmonic oscillator (SHO), $\ddot{x}+2γ\dot{x}+ω_0^2x=0$, one of the most fundamental systems in physics. Our goal is to identify the methods transformers use to model the SHO, and to do so we hypothesize and evaluate possible methods by analyzing the encoding of these methods' intermediates. We develop four criteria for the use of a method within the simple testbed of linear regression, where our method is $y = wx$ and our intermediate is $w$: (1) Can the intermediate be predicted from hidden states? (2) Is the intermediate's encoding quality correlated with model performance? (3) Can the majority of variance in hidden states be explained by the intermediate? (4) Can we intervene on hidden states to produce predictable outcomes? Armed with these two correlational (1,2), weak causal (3) and strong causal (4) criteria, we determine that transformers use known numerical methods to model trajectories of the simple harmonic oscillator, specifically the matrix exponential method. Our analysis framework can conveniently extend to high-dimensional linear systems and nonlinear systems, which we hope will help reveal the "world model" hidden in transformers.

Figures (23)

• Figure 1: We aim to understand how transformers model physics through the study of meaningful intermediates. We train transformers to model simple harmonic oscillator (SHO) trajectories and use our developed criteria of intermediates to show that transformers use known numerical methods to model the SHO.
• Figure 2: We plot the $R^2$ of Taylor probes for the intermediate $\bm{w}$ within models trained on the task $\bm{Y} = \bm{wX}$. We see that larger models have $\bm{w}$ encoded, often linearly, with little gain as we move to higher degree Taylor probes, while smaller models do not have $\bm{w}$ encoded.
• Figure 3: We test the correlation between model performance and the encoding of $\bm{w}$ on 5 of our 25 models of evenly spaced performance quality. We plot normalized values for the error of the encoding ($1-R_w^2$) in red and the mean squared error of the model (${MSE}_M$) in blue. We find that the ability of the best performing models to in-context learn is highly correlated with their encoding of $\bm{w}$ ($R^2(MSE,w)$.
• Figure 4: Left: We plot $\max({\bar{R^2}})$ of the reverse probe from $[w, w^2] \rightarrow HS$ across all models, and find that the intermediate $\bm{w}$ can explain significant amounts of variance in model hidden states. Right: We intervene using reverse probes to make all models output $\bm{w'} = 0.5$. This intervention can either fail (16/25), be partially successful nonlinearly (2/25) or linearly (3/25), or be successful (4/25).
• Figure 5: We analyze the intermediates of our undamped harmonic oscillator models, and find all three methods encoded, with the matrix exponential method best represented. This provides initial correlational evidence for all three methods.