Research map / Architectures and efficiency / Attention and transformers

Predictive Geometry of Hidden Trajectories in Transformers

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Derives a pullback Fisher operator as the local curvature of layerwise loss-to-go in decoder-only transformers, and uses its spectrum and tokenwise traces for sensitivity analysis, rank allocation, token pruning, and KD-regularized low-rank student recovery.

Problem

Which hidden-state directions and token positions in a trained decoder-only transformer are constrained by the terminal next-token loss, and how can this loss-aware geometry be measured and used for compression and distillation.

Approach

Near successful trajectories, the Hessian of the layerwise loss-to-go is approximately a pullback Fisher operator K = J^T F J from the output distribution to the intermediate residual stream; its spectrum gives observable vs prediction-null directions, and its per-token traces give a causally supported tokenwise saliency score.

Key results

• Geometry-guided rank allocation helps most at tau=0.05 and 0.1 for several small/medium models (Gemma 2B, Qwen2.5 3B, SmolLM2 1.7B, OLMo-3 7B).
• Adding geometry regularizer to reverse KL improves PPL in 9/14 matched settings (mean gain 21.3); to skew KL in 10/14 (mean gain 59.3).
• Adding geometry to forward KL does not improve average PPL (mean gain -38.4, 8/14 wins).

Surprising findings

• Geometry regularization hurts forward KL on average (mean Delta -38.4), so it helps only with stronger autoregressive KD objectives.
• Attention-score pruning remains a strong baseline; kappa-based pruning is only competitive with, not clearly better than, attention.

Limitations

• Supervised cross-entropy has a nonzero first-order term, so the pure Fisher quadratic is not a full approximation; left to future work.
• Local results hold only near successful trajectories and low loss; the error terms scale with sqrt(loss).
• Structured token pruning does not clearly beat attention-based pruning.

Experimental setup

Datasets: WikiText OpenWebText FineWeb

Benchmarks: WikiText OpenWebText FineWeb

Metrics: perplexity KL divergence perplexity increase effective observable rank (participation ratio) win rate

Baselines: uniform SVD uniform LoRA-style low-rank allocation attention-score pruning gradient-activation saliency random pruning FKL RKL SkewKL

Builds on: forward KL distillation reverse KL distillation skew KL distillation Pearlmutter Hessian-vector product activation-aware SVD

Closest papers by meaning

• Constrained Belief Updates Explain Geometric Structures in Transformer Representations (ICML 2025 · similarity 0.87)
• Learning Theory of Transformers: Local-to-Global Approximation via Softmax Partition of Unity (NeurIPS 2026 · similarity 0.87)
• Linear Transformers Implicitly Discover Unified Numerical Algorithms (NeurIPS 2025 · similarity 0.87)
• Dynamical Properties of Tokens in Self-Attention and Effects of Positional Encoding (NeurIPS 2025 · similarity 0.87)
• Transformers Represent Belief State Geometry in their Residual Stream (NeurIPS 2024 · similarity 0.86)
• Dynamics of the Transformer Residual Stream: Coupling Spectral Geometry to Network Topology (NeurIPS 2026 · similarity 0.86)