activation geometry tracks the model's predicted latent dynamics

part of reasoning under uncertainty in humans and language models
neurips 2026 · case study: gpt oss 20b, low vs high reasoning effort

< all sections > model playground

reasoning traces are long and open-ended, which makes it hard to know where to look inside the model. the behavioral model from the main results narrows the search. once fitted, it makes concrete predictions about how the agent's beliefs should evolve round by round: the bayesian posterior $\mathbf{p}_t$, the agent's distorted memory $\mathbf{h}_t$, and the policy $\boldsymbol{\pi}_t$ it chooses from.

we do not try to find these fitted variables inside the network. we ask a simpler question: do the model's internal choice representations change across rounds the way the fitted latents say they should? as a case study we take gpt oss 20b under low and high reasoning effort. use the round slider below to step through one game. every figure on this page follows it.

round 11 / 15

01one game, two predictions

both conditions see exactly the same evidence, so the bayesian posterior is identical. around round 11 it flips from favoring C to favoring A. the fits disagree about what the model does with that evidence. under high reasoning memory is near-perfect ($\beta \approx 0$), so $\mathbf{h}_t$ tracks the switch and the policy sharpens onto A ($\kappa_f \gg 0$). under low reasoning memory is stubborn ($\beta < 0$), so $\mathbf{h}_t$ stays on C and the policy stays diffuse ($\kappa_f \approx 0$).

that gives a specific prediction: with high reasoning the internal representations should reorganize around A as evidence accumulates; with low reasoning they should stay anchored near C.

fig. 1 the example game: evidence and optimal bayesian posterior

top: the button sampled on each round and its color (outlined cells were available, blank ones occluded). bottom: bayesian posterior $\mathbf{p}_t$ that each button is the biased one.

02predicted latents vs activations

we recorded the residual stream of gpt oss 20b (24 blocks, $d = 2880$) while it played 50 games per reasoning level. at each round $t$ and layer $\ell$, the activation $\mathbf{a}^{(\ell)}_t$ is the average over the last 5 tokens of the reasoning trace. for each choice $k$ we computed a centroid $\mathbf{c}_k$: the mean activation over rounds on which the model chose that button. for the example game we then turned each round's distances to the four centroids into a similarity vector:

$$d_k(t) = \big\lVert \mathbf{a}_t - \mathbf{c}_k \big\rVert_2, \qquad s_k(t) = \frac{e^{-d_k(t)}}{\sum_{j=1}^{4} e^{-d_j(t)}}$$

centroid similarity, computed in the full 2880-dimensional space

$\mathbf{s}(t)$ lives on the same simplex as the model's policy $\boldsymbol{\pi}_t$, but it comes only from raw activations. the two can be compared directly.

fig. 2 fitted memory $\mathbf{h}_t$ and policy $\boldsymbol{\pi}_t$ next to activation similarity $\mathbf{s}(t)$

rows are buttons A–D; columns are rounds, with the final choice last. the bars compare, for the selected round, what the behavioral model predicts ($\boldsymbol{\pi}_t$) with what the activations show ($\mathbf{s}(t)$). activations start at round 2.

03activation trajectories

we ran pca on layer-1 activations (close to the peak of choice decodability), averaged within each round, choice and reasoning level. the top four components capture more than 60% of the variance in both conditions. projecting the example game into this basis shows its trajectory.

the geometry matches the prediction. under high reasoning the trajectory starts near C and follows the posterior flip with a slight lag: the closest centroid $\arg\min_k d_k(t)$ first leaves C at round 12, passes briefly through B, and settles on A from round 14. under low reasoning it stays anchored near C for the whole game.

fig. 3 example game in the layer-1 pca basis

A B C D choice centroids $\mathbf{c}_k$ · trajectory colored by closest centroid · dashed: distance to each centroid at the selected round

2d views of a higher-dimensional space: distances are computed in the full 2880-d space, so a point can be closest to a centroid that looks farther away here. pc units are arbitrary.

04choices are linearly decodable

this kind of geometry only means something if the activations actually encode the choice. a ridge probe per layer predicts the one-hot choice $\mathbf{y}_t$ from the activation:

$$\hat{\mathbf{y}}_t = W_\ell^{\top} \mathbf{a}^{(\ell)}_t, \qquad W_\ell = \arg\min_{W} \; \big\lVert Y - A^{(\ell)} W \big\rVert_F^2 + \alpha \lVert W \rVert_F^2$$

$\alpha$ by leave-one-out over $10^{-3} \dots 10^{6}$; held-out $R^2$ over 5 folds

the probes decode the model's choice at every round with held-out $R^2 > 0.8$ in both conditions, peaking in early-to-mid layers. the representation is specific to each condition: a probe trained under one reasoning level and tested on the other (swapped control) drops to around zero.

fig. 4 choice decodability across layers

5-fold cross-validated $R^2$ of ridge probes predicting the 4-way choice from residual-stream activations after each block. swapped controls (sc) train on one condition and test on the other.

methods in brief

activations: residual stream after the mlp of each of the 24 blocks, at every round (15 sampling rounds plus the final choice), averaged over the last 5 tokens of the reasoning trace. 50 games per reasoning level.
probes: closed-form ridge regression per (layer, reasoning level), regularization $\alpha$ chosen by leave-one-out over $10^{-3}, \dots, 10^{6}$, held-out $R^2$ averaged over 5 folds and the 4 target dimensions.
pca: on conditional averages (one row per round × choice × reasoning level) at layer 1, fitted separately per reasoning level. pc1 explains 27.8% (low) and 19.6% (high) of the variance; the top four, ~64% and ~61%.
centroids: mean conditional-average activation per choice in the full 2880-d space; similarity $s_k \propto e^{-d_k}$ from euclidean distances.

05what this does and doesn't show

the analysis is correlational. it shows that the activation geometry moves the way the fitted latents predict; it does not show that these subspaces are used to make the choice. targeted causal interventions are the natural next test. the broader point is methodological: an algorithmic-level model turns the open-ended question "what is this reasoning trace computing?" into a few concrete predictions that can be checked against activations.

mockup · probe $R^2$ curves, centroid similarities and pca trajectories were read back from the paper's vector figure 5 for this preview, and should be replaced by the source arrays before publishing. evidence, posterior and fitted latents are exact.