Predictive Geometry of Hidden Trajectories in Transformers

Timur Mudarisov; Mikhail Burtsev; Tatiana Petrova; Radu State

cs.LGSep 29, 2026

Predictive Geometry of Hidden Trajectories in Transformers

Authors: Timur Mudarisov, Mikhail Burtsev, Tatiana Petrova, Radu State

Organizations: University of Luxembourg, Luxembourg

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.

Figures & tables

Panel of Figure 1 — Figure 1 : Geometry of low-loss hidden trajectories. Top: hidden-state trajectory and layerwise low-loss sets $\mathcal{A}_{\varepsilon}^{\ell}$ . Bottom left: local pullback geometry near $X_{\ell}^{\ast}$ . Bottom right: prediction-null foliation and observable quotient.

Panel of Figure 2 — Figure 2 : Experiment 1: anisotropy and effective observable dimension across layers. Left: KL sensitivity of equal-norm observable and null perturbations across layers. Right: Effective observable rank across layers for multiple models.

Panel of Figure 3 — Figure 3 : Experiment 2: geometry-guided $\Delta\mathrm{PPL}$ advantage on WikiText for different spectral thresholds $\tau$ . Each panel reports $\Delta\mathrm{PPL}_{\mathrm{adv}}=\Delta\mathrm{PPL}_{\mathrm{uniform}}-\Delta\mathrm{PPL}_{\mathrm{geom}}$ ; larger values indicate that geometry-guided rank allocation degrades perplexity less than the best matched uniform baseline.

KD baseline	Pairs	Wins	Mean $\Delta_{\mathrm{Geo}}\uparrow$	Median $\Delta_{\mathrm{Geo}}\uparrow$	SVD mean	ASVD mean
FKL	14	8	-38.4	-12.6	-35.0	-41.7
RKL	14	9	21.3	12.3	44.5	-1.9
SkewKL	14	10	59.3	8.6	66.7	51.9

Table 1 : Geometry advantage over matched KD baselines.

\Delta_{\mathrm{Geo}}=\mathrm{PPL}(\mathrm{KD})-\mathrm{PPL}(\mathrm{KD+Geo})

; higher is better.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Model name in paper	Hugging Face checkpoint	Model reference	Params
SmolLM2-1.7B	HuggingFaceTB/SmolLM2-1.7B	[ 4 ]	1.7B
LLaMA-3.2-1B	meta-llama/Llama-3.2-1B	[ 9 ]	1B
Qwen2.5-3B	Qwen/Qwen2.5-3B	[ 27 ]	3B
Phi-3 Mini	microsoft/Phi-3-mini-4k-instruct	[ 1 ]	3.8B
Mistral-7B	mistralai/Mistral-7B-v0.1	[ 14 ]	7B
LLaMA-2-7B	meta-llama/Llama-2-7b-hf	[ 23 ]	7B

Appendix

Table 2 : Model checkpoints used in the experiments.

Figure 4 : Local hidden-state geometry predicts terminal loss changes. We perturb intermediate residual-stream states at different radii and compare the true terminal loss change to local predictors induced by the downstream map. For supervised CE, the linear-only and linear-plus-quadratic predictors are nearly indistinguishable, indicating that CE is first-order dominated in this perturbation regime. The quadratic-only CE predictor is substantially weaker, as expected because supervised one-hot CE is generally not stationary at the reference prediction. In contrast, the teacher-KL quadratic predictor isolates the pullback-Fisher curvature and becomes strongly predictive at moderate perturbation radii. Curves show model-averaged $R^{2}$ ; for visualization, $R^{2}$ values are clipped to $[-1,1]$ before aggregation to avoid degenerate-regression outliers.

Figure 5 : Fisher token saliency is competitive with standard pruning baselines. We prune a fraction of token updates according to each baseline score and report the Fisher advantage $\Delta\mathrm{CE}_{\mathrm{baseline}}-\Delta\mathrm{CE}_{\mathrm{Fisher}}$ . Positive values indicate that Fisher pruning produces a smaller CE increase than the baseline. Fisher substantially outperforms random pruning and gives modest, consistent improvements over attention, $J^{\top}J$ , diagonal-Fisher, and gradient-times-activation criteria in the aggregate. Activation norm is less stable across pruning fractions, performing better at small drop fractions but worse at the largest drop fraction. The $y$ -axis uses a symlog scale with linear threshold $0.1$ to show both small differences near zero and large improvements over weak baselines.

Model	Method	Rank	Compression	Init PPL $\downarrow$	Final PPL $\downarrow$	$\Delta_{\mathrm{init}}\uparrow$
Qwen2.5-0.5B	asvd+fkl	16	40.546x	2.44e+05	581.4	2.43e+05
Qwen2.5-0.5B	asvd+rkl	16	40.546x	2.44e+05	653.0	2.43e+05
Qwen2.5-0.5B	asvd+skewkl	16	40.546x	2.44e+05	551.2	2.43e+05
Qwen2.5-0.5B	svd+fkl	16	40.546x	2.15e+06	690.9	2.15e+06
Qwen2.5-0.5B	svd+fkl_geom	16	40.546x	2.15e+06	675.9	2.15e+06
Qwen2.5-0.5B	svd+rkl	16	40.546x	2.15e+06	761.5	2.15e+06

Appendix

Table 3 : Completed KD-geometry sweep results for Qwen2.5-0.5B. Lower is better for PPL and KD divergence; higher is better for improvement over compressed initialization.

Model	Method	Rank	Compression	Init PPL $\downarrow$	Final PPL $\downarrow$	$\Delta_{\mathrm{init}}\uparrow$
SmolLM2-135M	asvd+fkl	32	10.868x	1.00e+08	713.0	1.00e+08
SmolLM2-135M	asvd+fkl_geom	32	10.868x	1.00e+08	706.3	1.00e+08
SmolLM2-135M	asvd+rkl	32	10.868x	1.00e+08	831.4	1.00e+08
SmolLM2-135M	asvd+rkl_geom	32	10.868x	1.00e+08	808.3	1.00e+08
SmolLM2-135M	asvd+skewkl	32	10.868x	1.00e+08	690.9	1.00e+08
SmolLM2-135M	asvd+skewkl_geom	32	10.868x	1.00e+08	693.1	1.00e+08

Appendix

Table 4 : Completed KD-geometry sweep results for SmolLM2 models. Lower is better for PPL and KD divergence; higher is better for improvement over compressed initialization.

Model	Init.	Rank	RKL	RKL+Geo	$\Delta_{\mathrm{Geo}}\uparrow$	SkewKL	SkewKL+Geo	$\Delta_{\mathrm{Geo}}\uparrow$
Qwen2.5-0.5B	ASVD	32	598.1	601.8	-3.8	503.3	495.5	7.8
Qwen2.5-0.5B	ASVD	64	418.2	409.2	9.1	343.2	357.5	-14.3
Qwen2.5-0.5B	SVD	16	761.5	756.7	4.8	690.9	682.3	8.6
Qwen2.5-0.5B	SVD	32	648.9	617.1	31.8	556.4	556.4	0.0
Qwen2.5-0.5B	SVD	64	521.0	534.2	-13.2	462.4	453.8	8.6
SmolLM2-135M	ASVD	32	831.4	808.3	23.1	690.9	693.1	-2.2

Appendix

Table 5 : Paired geometry comparisons for stronger KD baselines. Lower PPL is better;

\Delta_{\mathrm{Geo}}>0

means that adding geometry improves final PPL.

Panel of Figure 6 — Figure 6 : Experiment 3: $\Delta\mathrm{PPL}$ on WikiText under structured pruning for four drop fractions $\rho\in\{0.05,0.10,0.20,0.30\}$ . We compare our geometry-guided kappa criterion (red) to pruning based on attention scores (blue), a gradient-activation saliency baseline (purple), and random pruning (green); lower $\Delta\mathrm{PPL}$ is better.

Panel of Figure 7 — Figure 7 : Direct validation of the local pullback-Fisher approximation for top-Fisher perturbation directions on WikiText and FineWeb. For each dataset, the left panel shows the Spearman correlation between the observed teacher-KL increase and the quadratic prediction $\frac{1}{2}\alpha^{2}v^{\top}K_{\ell}v$ for each model separately, as a function of the relative perturbation radius $\rho$ . The right panel shows the mean trend across models, with error bars denoting uncertainty across models. Correlations are highest at the smallest perturbation radii and generally decrease as $\rho$ grows, consistent with the local nature of Theorem 1 .

Model	Seq. len.	Examples	Model layers	GPU hours
LLaMA-3.2-1B	512	32	16	$\approx 0.8$
SmolLM2-1.7B	512	32	24	$\approx 1.0$
Gemma-2B	512	32	26	$\approx 1.2$
Qwen2.5-3B	512	32	36	$\approx 2.0$
Phi-3 Mini	512	32	32	$\approx 2.2$
LLaMA-2-7B	512	32	32	$\approx 4.0$

Appendix

Table 6 : Representative compute cost for geometry estimation on one dataset. GPU hours are approximate wall-clock estimates for computing all reported geometry statistics on the selected layer subset. Actual runtime depends on GPU type, attention implementation, and model checkpoint.

Explore similar work

Sep 29, 2026cs.CL

The Geometry of Inference in Transformer Residual Streams

Transformer language models build predictions through successive residual updates, but how their representations become specific to an eventual outcome remains unclear. We study this process by comparing intermediate residual states with their own final states and an empirical bank of final states from other contexts. Across six pretrained language models, the own endpoint becomes preferable to the average alternative early, while many individual endpoints remain closer. These competing sets generally shrink with depth, but their membership changes and their surviving endpoints need not become more similar to one another. Directional alignment and endpoint rank can therefore improve while Euclidean distance to the final state changes little. We develop a simple high-dimensional model that separates the roles of norm, alignment, and endpoint geometry, showing how gradual directional changes can produce sharp reductions in competition. We also prove that a straight path toward the own endpoint cannot introduce new competitors under either Euclidean or cosine distance; observed entries thus establish departures from straight-line convergence. Finally, endpoints associated with lower-ranked output tokens tend to lie farther away in cosine distance across all studied models, connecting residual geometry to output organization. Together, these findings characterize increasing geometric specificity during transformer inference and explain why distance, competitor count, and concentration of the surviving endpoints provide distinct views of that process.

Timur Mudarisov, Mikhail Burtsev, Radu State

University of Luxembourg, Luxembourg · London Institute for Mathematical Sciences, London, UK

Transformer Residual Streams Transformer Architectures

Ask PDF

Jun 8, 2026cs.LG

Trajectory Geometry of Transformer Representations Across Layers

Understanding how transformer representations evolve across layers, not merely what they encode, remains an open problem in mechanistic interpretability. We recast the transformer forward pass as a discrete population trajectory through a high-dimensional representation manifold, drawing on geometric tools from computational neuroscience. Rather than probing for pre-specified features, we characterize trajectory geometry using five metrics computed directly in the ambient space: trajectory length, curvature, a semantic convergence index, layerwise cosine similarity, and representational stability. Across three model families (GPT-2, TinyLlama, Qwen2.5) and five controlled prompt families, we report four findings. First, semantically related prompts converge significantly in middle-to-late layers (peak CI 0.41--0.58, p<0.001, Mann-Whitney U), consistent with attractor-like dynamics. Second, reasoning tasks produce trajectories of greater curvature than lexical variations (0.71--0.83 rad vs. 0.27--0.31 rad), suggesting curvature encodes computational complexity. Third, ambiguous tokens exhibit trajectory bifurcation with up to 5.6x representational separation by the final layer, absent in unambiguous controls. Fourth, layerwise cosine similarity reveals a universal three-phase structure: encoding, elaboration, and output preparation, consistent across all three architectures. All four effects vanish under shuffled-layer and random-embedding controls. We release a fully open-source, model-agnostic pipeline and argue that trajectory geometry constitutes a principled, probe-free lens for mechanistic interpretability.

Vishal Pandey, Gopal Singh, Yacine Mahdid

London, UK · Athens, GR · Montreal, CA

Transformer Architectures Layer-Wise Mechanistic Interpretability Trajectory Generation Bifurcation

Ask PDF

Aug 12, 2026cs.LG

Geometric and Behavioral Stratification in Transformer Residual Streams

Trained transformer models develop privileged bases: coordinate axes whose statistics differ from the rest of the residual stream. But what kind of direction does such a basis select? We investigate the prediction direction, the unembedding direction of the token a model currently predicts, and find that it functions as a content-defined privileged anchor. Measured with respect to this anchor, residual-stream variation is geometrically and behaviorally stratified by proximity to the prediction. The stratification holds in all eighteen models tested (dense and mixture-of-experts, 7B-120B, base and instruction-tuned). A narrow, scale-invariant prediction interface concentrates readout-relevant structure, while the vast prediction-distal complement expands with model scale. Because the prediction direction sits nearly orthogonal to the principal variance axes, variance-based analyses recover this organization only partly, and the shortfall grows with prompt heterogeneity. Anchoring reveals a steep geometric gradient: prediction-proximal regions are highly structured and cluster related prompts, while the complement is flatter and anti-discriminates among prompt groups. The interface is a narrow slice but functionally decisive. Disrupting the variance directions closest to the prediction causes immediate divergence and frequent task-frame shifts; disrupting the next level down delays divergence and preserves framing. The complement is weakly readout-aligned per direction yet causally and temporally load-bearing, and behavior is driven by direction rather than magnitude. These results establish the prediction direction as a privileged anchor distinct from previously described coordinate axes, and give a geometric account of how high-dimensional computation coexists with linear readout.

Nelson Guda

Transformer Residual Streams Residual Stream Transformer Architectures

Ask PDF