Token-axis Predictive Decomposition: What Information Does Each Token Add to the Hidden States of a Language Model?

Interpretability research on large language models commonly relies on the decomposition along the layer axis provided by the architecture, $h_t^l=h_t^{l-1}+\Delta h_t^l$. Along the token axis, however, the architecture provides no such decomposition. We decompose the hidden state $h_t$ at position $t$ into its expectation conditioned only on the past hidden states, ${\hat{h}}_t=\mathbb{E}[h_t\mid h_{<t}]$, and the residual ${\epsilon }_t=h_t-{\hat{h}}_t$. Since the true conditional expectation cannot be computed, we approximate it with several predictors of different architectures and also examine how much the observed properties depend on the choice of predictor. At five layers of Pythia-410m-deduped, we evaluate this decomposition from three perspectives: temporal structure, linguistic information, and effect on the output. First, using a Transformer as the predictor reduces the squared error only slightly in the middle layers, but when the predictor aggregates the entire past, the autocorrelation of ${\hat{h}}_t$ exceeds that of $h_t$, while that of ${\epsilon }_t$ is nearly zero. Second, when linear probes trained on $h_t$ are applied unchanged to both components, attributes determined by the context (e.g., the document source) can be read out from ${\hat{h}}_t$ and token-local attributes (e.g., noun number) from ${\epsilon }_t$; in deeper layers, the component that carries the attributes shifts toward ${\hat{h}}_t$. Third, the KL divergence of the next-token distribution when $h_t$ is replaced with ${\hat{h}}_t$ is generally smaller for better predictors, and the output is more sensitive to reducing ${\epsilon }_t$ than to increasing it by the same amount. These results suggest that the predictive decomposition along the token axis, like the decomposition along the layer axis, can serve as a unit for analyzing internal representations.