Predictive Geometry of Hidden Trajectories in Transformers

2026-09-29

AI Summary

This study addresses the unclear geometric properties of hidden states in Transformers trained solely with terminal losses, particularly regarding their constraints from downstream computation. We analyze the local second-order geometry of layer-wise losses and employ the pullback Fisher operator to identify output-sensitive directions and predictive null spaces, thereby constructing an observable residual stream subspace. Furthermore, we propose a token-level curvature score based on Fisher-weighted sensitivity, serving as a loss-aware alternative to attention magnitude for enabling non-uniform hierarchical rank allocation, which is efficiently estimated via matrix-free Jacobian-vector products. Evaluated on datasets such as WikiText, this approach effectively predicts perturbation sensitivity, provides competitive structured token pruning signals, and significantly enhances the recovery performance of low-rank student models during autoregressive distillation.

Abstract

Decoder-only transformers are trained only through a terminal next-token prediction loss, yet this loss constrains every intermediate hidden state through the fixed downstream computation. We formalize this constraint by studying layerwise loss-to-go functions: the terminal loss obtained by continuing a candidate hidden state through the remaining transformer blocks. Around successful validation trajectories, we show that the local second-order geometry of these functions is governed, up to low-loss residual terms, by a pullback Fisher operator on hidden-state space. Its spectrum identifies output-sensitive directions and approximately prediction-null directions, yielding a local observable subspace of the residual stream. For causal transformers, the same geometry induces a tokenwise curvature score: a Fisher-weighted sensitivity of the target logits to perturbations of each token's hidden state. This score vanishes outside the causal ancestor set of the target and is controlled by downstream Jacobian couplings, making it a loss-aware alternative to attention magnitude. We estimate these quantities using matrix-free Jacobian-vector and vector-Jacobian products and evaluate them across decoder-only language models on WikiText, OpenWebText, and FineWeb. Empirically, the induced geometry predicts perturbation sensitivity, supports nonuniform layerwise rank allocation, yields competitive structured token-pruning signals, and improves low-rank student recovery when added to stronger autoregressive distillation objectives such as reverse KL and skew KL. These results support a predictive-geometric view of transformer computation: near successful trajectories, the terminal loss induces a thin, anisotropic set of output-relevant hidden-state directions that can be measured and exploited for compression and distillation.

Problem

• Transformer hidden states
• predictive geometry
• Fisher information
• next-token prediction loss
• model compression

Innovation

• pullback Fisher operator
• loss-to-go geometry
• tokenwise curvature score
• matrix-free Jacobian products
• knowledge distillation