PaperlayerView exact revision on arXiv

Source-linked AI summary

Predictive Geometry of Hidden Trajectories in Transformers

Timur Mudarisov, Mikhail Burtsev, Tatiana Petrova, Radu State

arXiv:2609.37717v1cs.LGcs.CL

TL;DR

The paper addresses how terminal next-token loss constrains intermediate hidden states, where prior empirical accounts do not derive relevant directions from that objective. It studies loss-to-go geometry through a pullback Fisher operator and tokenwise sensitivity, finding useful signals across several model-reduction settings.

  • Problem

    Prior empirical accounts characterize hidden-state structure but do not derive relevant directions from terminal prediction loss or provide loss-aware measures of hidden-state and token importance.

    Source: §1 · pp. 1–2

  • Method

    The paper studies local loss-to-go geometry for intermediate states of a fixed trained decoder-only transformer, using pullback Fisher curvature and a tokenwise sensitivity score.

    Source: §2 · pp. 2–3, §1 · p. 2

  • Results

    The geometry predicts perturbation sensitivity, supports matched-budget layerwise rank allocation and competitive structured pruning, and improves low-rank student recovery with stronger autoregressive distillation objectives.

    Source: §5 · p. 9

  • Takeaways & Limitations

    Near successful trajectories, terminal loss induces a thin, anisotropic set of prediction-relevant hidden-state directions that can be measured for compression and distillation.

    Source: §5 · p. 9

    The theory is scoped to a fixed trained decoder-only pre-norm transformer and local states within the neighborhood where its geometric assumptions apply.

    Source: §3 · p. 3, §3.3 · p. 6

Abstract

from arXiv · show

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.

1 Introduction

The paper asks which intermediate hidden-state directions are constrained by terminal next-token loss, extending empirical accounts of representation structure with a loss-based geometry. It derives a pullback-Fisher view of hidden-state and tokenwise sensitivity and evaluates its use in model reduction.

1 Introduction

  • Prior work finds hidden representations structured and compressible, but does not derive relevant directions from terminal prediction loss or provide loss-aware measures of hidden-state and token importance.The paper contrasts its objective-based account with empirical characterizations of representation structure.Source: §1 · pp. 1–2
  • Near successful validation trajectories, layerwise loss-to-go curvature is dominated in the low-loss regime by a pullback Fisher operator.The operator identifies output-sensitive directions alongside approximately prediction-null directions in hidden-state space.Source: §1 · p. 2
  • The pullback-Fisher spectrum measures output-sensitive directions by layer, suggesting data-dependent rank allocation instead of uniform allocation.The operator’s token-fiber traces define a local score of each token’s contribution to target-prediction curvature.Source: §1 · p. 2
  • For causal transformers, the tokenwise score is supported only on causal ancestors of the target and is controlled by downstream Jacobian couplings.The paper uses it as a loss-aware alternative to attention-based pruning heuristics [26] [17].Source: §1 · p. 2
  • The geometry is evaluated for perturbation sensitivity, layerwise rank allocation, structured token pruning, and low-rank student distillation.The paper frames these as loss-aware signals for preserving prediction-relevant hidden directions under compression.Source: §1 · p. 2

2 Problem Statement

The problem is to characterize intermediate hidden states that can still yield low terminal prediction loss when continued through a fixed trained transformer. The paper asks how their local geometry and tokenwise distribution can inform loss-aware model reduction.

2 Problem Statement

  • The setup fixes a trained decoder-only transformer and target position, then considers its hidden-state trajectory and final prediction readout.The analysis concerns learned forward dynamics, not the model’s training dynamics.Source: §2 · p. 2
  • A layerwise loss-to-go evaluates the terminal loss after continuing an intermediate state through the remaining transformer, and its low-loss set contains states whose continued loss is at most ε.The terminal loss is applied after the final readout.Source: §2 · p. 3
  • The paper asks whether successful trajectories occupy lower-dimensional pullback low-loss regions and how prediction-relevant geometry is distributed across token positions.These questions motivate loss-aware layerwise low-rank compression and token-sparse computation.Source: §2 · p. 3

3 Theory

The theory studies local hidden-state geometry for a fixed trained decoder-only transformer and derives consequences for tokenwise sensitivity and geometry-guided reduction. Near successful low-loss trajectories, output-sensitive curvature is distinguished from prediction-null directions.

3 Theory

  • The theoretical setting fixes a trained decoder-only pre-norm transformer, a target position, and a softmax readout with terminal loss.The geometric claims are local to intermediate states compatible with low terminal prediction loss.Source: §3 · p. 3

3.1 Local geometry of layerwise low-loss sets

  • At low loss, the dominant local Hessian curvature is a pullback Fisher term, and the resulting low-loss region is approximately an ellipsoidal tube when the first-order term is small.Prediction-null directions leave centered target logits unchanged to first order, while the observable geometry captures output-sensitive directions.Source: §3.1 · p. 4

3.2 Tokenwise geometric saliency in transformers

  • The tokenwise score measures how much of the local output-sensitive curvature is carried by each token’s hidden state.It is the Fisher-weighted squared sensitivity of target-position logits to perturbations of that token.Source: §3.2 · p. 5, Define the tokenwise curvature score · p. 5

Define the tokenwise curvature score

  • A token contributes only if it is a causal ancestor of the target, and its contribution is smaller when its Jacobian coupling is weak.The score is therefore both sparse and graded across causal ancestors.Source: Define the tokenwise curvature score · p. 5

3.3 Geometry-guided reduction

  • The theory provides measurable operators and constructive observable subspaces and token rankings for compression, including projection onto dominant Fisher directions.The stated local reduction setting assumes both states lie within the neighborhood where the geometry theorem applies.Source: §3.3 · pp. 5–6

4 Empirical Validation and Geometry-Guided Reduction

The section tests whether pullback-Fisher geometry captures hidden-state sensitivity and whether it can guide compression and recovery. Across five experiments, it is evaluated as a local curvature model and as a signal for rank allocation, token pruning, and low-rank student distillation.

4 Empirical Validation and Geometry-Guided Reduction

  • The evaluation tests the geometry both as a description of hidden-state sensitivity and as a signal for compression and recovery.Measurements use held-out validation sequences from WIKITEXT, FINEWEB, and OPENWEBTEXT unless otherwise stated.Source: §4 · p. 6

4.1 Experiment 0: Direct validation of local predictive geometry

  • Teacher-KL’s quadratic predictor becomes strongly predictive at moderate perturbation radii, whereas supervised cross-entropy is first-order dominated.For cross-entropy, the linear-only and linear-plus-quadratic predictors nearly coincide, and the quadratic-only predictor is weaker.Source: §4.1 · p. 6

4.2 Experiment 1: anisotropy and effective observable dimension

  • Observable perturbations are orders of magnitude more harmful than equal-norm null perturbations at almost every layer, while effective observable rank grows with depth but remains below hidden dimension.The anisotropy gap typically widens toward later blocks.Source: Figure 2 · p. 7, §4.2 · p. 7

4.3 Experiment 2: geometry-guided rank allocation

  • At moderate thresholds τ = 0.05 and τ = 0.1, several small and medium models benefit from geometry-guided rank allocation over matched uniform baselines.The advantage is smaller for larger, more homogeneous model families.Source: §4.3 · pp. 7–8, Figure 3 · p. 8

4.4 Experiment 3: geometry-guided pruning

  • The κ-based pruning policy is competitive with attention and slightly better in several model/drop-fraction settings; both structured policies outperform random pruning.The comparison covers WIKITEXT, OPENWEBTEXT, and FINEWEB.Source: §4.4 · p. 8

4.5 Experiment 4: predictive-geometry distillation of low-rank students

  • Adding geometry improves reverse KL in 9/14 matched settings and skew KL in 10/14; the largest average gain is with skew KL.Geometry does not improve average perplexity when added to forward KL.Source: §4.5 · p. 9, Table 1 · p. 9

5 Conclusion

Near successful validation trajectories in the low-loss regime, terminal loss induces pullback Fisher geometry on intermediate hidden states. Its spectrum distinguishes output-sensitive from approximately prediction-null directions, while tokenwise traces characterize causal contributions.

5 Conclusion

  • Near successful trajectories in the low-loss regime, the terminal loss induces a pullback Fisher operator that separates output-sensitive from approximately prediction-null residual-stream directions.The resulting local observable structure is induced by the output distribution's geometry pulled back through the downstream computation.Source: §5 · p. 9
  • For causal transformers, tokenwise curvature vanishes outside the target's causal ancestors and is controlled by downstream Jacobian couplings.The spectrum and tokenwise traces provide loss-aware signals for rank allocation and token pruning.Source: §5 · p. 9
  • Across reduction settings, predictive geometry predicts perturbation sensitivity, supports matched-budget layerwise rank allocation and competitive structured pruning, and regularizes low-rank student recovery with stronger distillation objectives.The stated objectives include reverse KL and skew KL.Source: §5 · p. 9

2. Pullback Fisher geometry. Define

  • When the linear term is small, the local low-loss set is approximated by an ellipsoid induced by the pullback Fisher quadratic form.High-curvature directions are thin and low-curvature directions are thick.Source: §2 · p. 13
  • Directions in the kernel of the centered-logit Jacobian are first-order prediction-null; under locally constant rank, observable variation lies on the quotient by that kernel.This quotient is equivalent to any chosen complement of the null space.Source: §2 · p. 14
  • For supervised one-hot cross-entropy, the expansion can contain a potentially important linear term, so the Fisher operator is local output curvature rather than a complete approximation to loss change.The paper uses teacher-KL for direct Fisher validation and leaves systematic evaluation of the full linear-plus-quadratic supervised-CE approximation to future work.Source: §2 · p. 14

A.3 Proof of Proposition 1

The proof details connect tokenwise curvature to causal influence, represent downstream Jacobian coupling through causal paths, and justify a spectral-tail rank rule.

A.3 Proof of Proposition 1

  • A token outside the downstream causal ancestor set of the target has no path to its logits, so its tokenwise curvature score vanishes.The result assumes a decoder-only transformer with a standard causal mask.Source: §A.3 · p. 14

3. Pathwise coupling bound. Let G(r←u)

  • The downstream Jacobian from a token at layer ℓ to target logits can be written as a sum over directed causal paths.Each path's block Jacobians are multiplied in path order.Source: §3 · p. 15

A.4 Proof of Proposition 2

  • A layerwise spectral-tail threshold bounds the discarded subspace's share of total local predictive curvature by a controlled fraction.The rank rule uses the spectral tail as a proxy for discarded output-sensitive curvature.Source: §A.4 · p. 15

B Experimental Protocol

The protocol fixes model weights, separates calibration from evaluation, and estimates curvature without forming the pullback-Fisher matrix.

B Experimental Protocol

  • Experiments use fixed weights and disjoint calibration and evaluation examples, with open decoder-only models of approximately 1B–9B parameters evaluated on WikiText, OpenWebText, and FineWeb.Target positions are sampled from the second half of non-padding context, and token-level scores are computed only over their causal ancestors.Source: B Experimental Protocol · p. 16

C Numerical Estimation of Predictive Geometry

C.1 Matrix-free multiplication by Kℓ

  • A Kℓv product applies the downstream-logit Jacobian to v, applies the output Fisher weighting, then maps back with the transposed Jacobian.The procedure uses a JVP, the Fisher multiplication, and a VJP.Source: §C.1 · p. 17

C.2 PyTorch implementation with JVP and VJP

The implementation estimates predictive geometry through downstream computations and matrix-free products, with probe-based traces and practical controls for data, precision, and memory.

C.2 PyTorch implementation with JVP and VJP

  • The PyTorch implementation represents the downstream computation as a function of the layer-ℓ hidden state and uses JVP/VJP operations for matrix-free products.Model weights remain fixed, and key-value caches that bypass the perturbation must be disabled or recomputed.Source: §C.2 · pp. 17–18

C.3 Estimating traces and effective ranks with Hutchinson probes

  • Hutchinson probes provide an unbiased trace estimate, and the participation-ratio effective rank is estimated from trace and squared-trace quantities.With m probes, the estimate averages the probe-based trace values.Source: §C.3 · p. 18

C.4 Practical implementation details

  • Teacher-KL removes the first-order term at the reference point, whereas one-hot cross-entropy may retain one; estimates are averaged across calibration examples and target positions.Curvature statistics are accumulated in float32, and bootstrap confidence intervals are used for trace-derived quantities.Source: §C.4 · p. 19

D Additional Details for Local-Geometry and Token-Pruning Experiments

  • The experiments test whether pullback-Fisher geometry describes local hidden-state sensitivity and whether token-fiber traces provide a useful loss-aware pruning signal.The operator measures which hidden-state perturbations are locally visible to the terminal predictive distribution.Source: D Additional Details for… · pp. 19–20
  • The pullback-Fisher operator is the downstream Jacobian transpose, softmax Fisher matrix, and downstream Jacobian composed to measure local visibility of hidden-state perturbations.The downstream map runs the remaining blocks, final normalization, and language-model head; caches that bypass a perturbed state must be disabled or recomputed.Source: §C.2 · pp. 17–18, D Additional Details for… · p. 20

D.1 Local Taylor Validation

The section validates pullback-Fisher geometry as a local output-curvature signal, then applies its tokenwise trace to pruning and uses predictive-geometry regularization in low-rank student recovery. The pruning score improves on random and several standard baselines, while the regularizer preserves output-Fisher weighting through a readout proxy.

D.1 Local Taylor Validation

  • Teacher-KL’s quadratic predictor becomes strongly predictive at moderate perturbation radii, while supervised CE is first-order dominated in this regime.For CE, the linear-only and linear-plus-quadratic predictors nearly coincide, and the quadratic-only predictor is substantially weaker.Source: Figure 4 · p. 21, §D.1 · p. 21

D.2 Token-Fiber Saliency and Structured Token-Update Pruning

  • The tokenwise Fisher trace measures downstream sensitivity weighted by output Fisher geometry, and small scores identify token updates that are safer to prune.The experiment evaluates removal of the least-salient token updates by measuring the resulting cross-entropy increase.Source: §D.2 · pp. 21–22
  • Fisher pruning substantially outperforms random pruning and gives modest, consistent aggregate improvements over attention, J⊤J, diagonal-Fisher, and gradient-times-activation criteria.Activation norm is less stable: it performs better at small drop fractions but worse at the largest fraction.Source: §D.2 · p. 23, Figure 5 · p. 23

E.1 Low-Rank Student Construction

  • Low-rank students replace selected teacher linear layers with low-rank factors, and recovery trains those factors while freezing other student parameters.The construction compares ordinary truncated SVD with activation-aware SVD, which uses activation second moments.Source: §E.1 · p. 23

E.2 KD Objectives

  • The distillation setup compares forward KL, reverse KL, and skew KL, with reverse KL described as a stronger, more mode-seeking baseline suited to lower-capacity students.Each objective is evaluated both with and without geometry regularization.Source: §E.2 · p. 24, §E.3 · p. 24

E.3 Predictive-Geometry Regularization

  • The implemented predictive-geometry regularizer uses a readout-Fisher proxy that preserves output-Fisher weighting of hidden-state errors but is weaker than the full downstream pullback-Fisher operator.The ideal geometry term is the local pullback-Fisher quadratic.Source: §E.3 · p. 24

E.4 Experimental Grid

The distillation study evaluates geometry regularization across low-rank student settings and interprets its effect through Fisher-weighted predictive mismatch. Geometry improves stronger autoregressive objectives in most matched settings, though its benefit depends on the student errors and regularization weight.

E.4 Experimental Grid

  • The sweep evaluates SVD and activation-aware SVD initializers with forward KL, reverse KL, and skew KL, each with or without geometry regularization.It uses sub-billion-parameter SmolLM2-135M, SmolLM2-360M, and Qwen2.5-0.5B models.Source: §E.4 · p. 24

E.5 Metrics

  • The study reports final perplexity and teacher-KD divergence, plus the compression gap to the teacher and improvement over compressed initialization.Lower is better for perplexity, compression gap, and KD divergence; higher is better for improvement over initialization.Source: §E.5 · p. 25

E.6 Geometry Advantage

  • Geometry advantage is the paired difference between base-objective perplexity and perplexity with the geometry term, with positive values indicating improvement.The primary comparisons test geometry alongside reverse KL and skew KL.Source: §E.6 · p. 25

E.7.1 Exact Fisher identity in logit space.

  • In logit space, the exact teacher-student KL is a path-integrated Fisher energy.The stated identity uses the softmax Fisher matrix along the prediction path.Source: §E.7.1 · p. 25

E.7.2 Local pullback to hidden space.

  • Pulling the teacher-student KL geometry back through the teacher readout yields a local hidden-space quadratic penalty that weights errors by their predictive effects.The full version uses the downstream teacher map; the implementation uses a readout-Fisher proxy.Source: §E.7 · p. 25, §E.7.2 · p. 26, §E.7.3 · p. 26

E.7.3 When a small geometry regularizer helps.

  • Geometry regularization is expected to help when student errors have large Fisher-weighted energy and the base KD objective has not already corrected those directions.It may help little or hurt when errors are mostly prediction-null, the base objective already controls relevant directions, or the geometry weight is too large.Source: §E.7.3 · p. 26

E.7 Theoretical examination

  • The regularizer acts as a hidden-state surrogate for teacher-student predictive mismatch, weighting hidden errors by their effect on the teacher’s predictive distribution rather than Euclidean size.In logit space, the exact teacher-student KL is a path-integrated Fisher energy.Source: §E.7.1 · p. 25, §E.7.3 · p. 26

E.8 Distillation Results

  • RKL+Geo improves over RKL in 9/14 matched settings, and SkewKL+Geo improves over SkewKL in 10/14; the largest average gain is for SkewKL.Geometry is not uniformly beneficial for forward KL, and the comparisons use matched completed settings.Source: §E.8 · p. 27, Table 5 · p. 29

F.1 Pruning results

The pruning figure compares geometry-guided pruning with attention, gradient-activation, and random criteria across four drop fractions. Separate validation reports that quadratic-prediction correlations are strongest at the smallest perturbation radii and generally decline as radius increases.

F.1 Pruning results

  • Figure 6 compares geometry-guided κ pruning with attention, gradient-activation saliency, and random pruning across four drop fractions, with lower WikiText ΔPPL indicating better performance.The listed drop fractions are 0.05, 0.10, 0.20, and 0.30.Source: Figure 6 · p. 30

F.2 Direct validation of the quadratic approximation.

  • Quadratic-prediction correlations are highest at the smallest perturbation radii and generally decrease as radius grows, consistent with the local nature of the approximation.The reported correlations compare observed teacher-KL increases with quadratic predictions on WikiText and FineWeb.Source: §F.2 · p. 31, Figure 7 · p. 31

F.3 Compute costs

  • The compute-cost table gives approximate GPU-hour estimates for geometry statistics on a selected layer subset, with runtime depending on GPU, attention implementation, and checkpoint.Source: Table 6 · p. 32

G Limitations and Future Directions

The framework is deliberately local and limited to fixed trained decoder-only pre-norm transformers; the authors identify broader theory, more efficient estimation, and integration with compression methods as future directions.

G Limitations and Future Directions

  • The theory describes loss-to-go geometry near a reference hidden trajectory; understanding how local geometries fit across larger regions and emerge during training remains open.The local focus enables connections between terminal prediction loss, measurable curvature, and residual-stream observability.Source: G Limitations and Future Directions · p. 32
  • Very small perturbation radii can make Taylor-validation changes indistinguishable from noise, while very large radii can leave the local Taylor regime.The authors report results across radii and suggest adaptive layer- or model-dependent choices.Source: G Limitations and Future Directions · p. 32
  • The analysis covers fixed trained decoder-only pre-norm transformers with causal self-attention and softmax readout, leaving other architectures for future work.Proposed extensions include encoder-decoder, bidirectional, mixture-of-experts, retrieval-augmented, and explicit-memory models.Source: G Limitations and Future Directions · p. 32
  • Matrix-free JVP/VJP products estimate curvature without materializing the operator, but these estimates cost more than forward-pass statistics such as activation norms or attention weights.Better probe sharing, sketching, or amortized estimators could improve practicality at larger scale.Source: G Limitations and Future Directions · p. 32
  • Rank-allocation and pruning experiments test loss-aware signals rather than replace full compression pipelines, and direct gradient heuristics can remain competitive for fixed pruning interventions.The authors propose combining Fisher-geometric saliency with first-order criteria.Source: G Limitations and Future Directions · p. 32
  • The hidden-state geometry term complements output-level distillation and appears most useful with stronger autoregressive objectives such as reverse KL or skew KL.Further study could vary training budgets, student sizes, low-rank parameterizations, and model types.Source: G Limitations and Future Directions · p. 33

NeurIPS Paper Checklist

  • The checklist states that experiments use fixed pretrained models, report bootstrap confidence intervals or standard errors, and provide code and reproducibility instructions.The described code covers matrix-free estimation, perturbation validation, rank allocation, pruning, and baseline comparisons.Source: NeurIPS Paper Checklist · pp. 34–35
Loading 2609.37717v1…