Keep the Directions That Change the Prediction

TL;DR for operators

Transformer compression should account for how changes to intermediate representations affect final predictions, not merely how accurately those representations can be reconstructed. Research across ten pretrained language models finds that equally sized hidden-state perturbations can produce orders-of-magnitude differences in prediction error. A local Fisher-based measure helps identify directions and token updates worth preserving. It improves layerwise compression decisions for several models and offers competitive structured pruning. Geometry-regularized distillation also improves most completed reverse-KL and skew-KL comparisons, but not consistently under forward KL. These are promising engineering signals rather than guaranteed deployment savings: the measurements are expensive, and their theoretical protection is local.

Equal-sized changes can produce unequal predictive damage

A model engineer reducing a transformer’s memory or computational footprint must decide which representation dimensions, layers, and token updates to preserve. Conventional compression methods often prioritize weight reconstruction, activation statistics, or representation magnitude. Those measurements are accessible, but they do not directly quantify damage to the model’s final predictions.

Consider two changes to an intermediate hidden state with identical numerical magnitude. One may substantially alter the next-token probability distribution, while the other has almost no measurable effect. A compression method that evaluates both changes primarily through reconstruction error may treat them similarly despite their different predictive consequences.

This separation is visible in Predictive Geometry of Hidden Trajectories in Transformers, by Mudarisov and colleagues.¹ In controlled experiments across pretrained decoder-only models, perturbations along leading prediction-sensitive directions generally caused orders-of-magnitude greater terminal teacher-KL divergence than equal-norm perturbations along approximately insensitive directions (Section 4.2, Figure 2).

The researchers also found that effective observable rank was substantially below the full hidden-state dimension, generally increasing with network depth.

The findings suggest that intermediate capacity is not equally valuable to the final prediction objective. They also create a more demanding engineering question: how can prediction-sensitive directions be measured without exhaustively testing every possible perturbation?

The final prediction defines a local geometry

The paper starts by evaluating each intermediate hidden state through the remaining, unchanged transformer layers. This creates a layerwise loss-to-go: the terminal prediction loss obtained by continuing computation from that particular state.

The relevant object is not the hidden state’s magnitude but the way downstream computation transforms its perturbations into changes in output probabilities.

The authors formalize this relationship through the pullback Fisher operator:

Here, $\bar{J}_\ell$ is the Jacobian mapping intermediate hidden-state changes into centered final logits, while $F^\ast$ measures the sensitivity of the output probability distribution.

The resulting operator assigns different local curvature to different hidden-state directions. Directions with large curvature strongly influence the predictive distribution; directions with little curvature have comparatively weak first-order influence.

Theorem 1 establishes the connection to terminal-loss curvature through a Hessian decomposition:

The residual becomes controlled when the reference prediction closely matches its target distribution and the relevant downstream second derivatives are bounded.

This qualification matters. The Fisher operator is not universally the exact Hessian of supervised cross-entropy. Indeed, the experiments distinguish teacher-KL perturbations, for which quadratic curvature becomes informative at measurable local radii, from supervised cross-entropy changes that are predominantly first-order in the tested regime.

A related theoretical result characterizes first-order prediction-null directions. These lie in the kernel of the centered downstream logit Jacobian, so sufficiently small perturbations along them have no first-order effect on the target prediction.

That does not mean these dimensions can simply be deleted. Larger perturbations may introduce higher-order effects, and changing model parameters is not equivalent to injecting a small perturbation into a fixed hidden trajectory.

Likewise, the paper’s effective observable rank measures spectral concentration, not the exact rank to which a transformer layer can safely be compressed.

Token importance depends on downstream computational paths

The geometric framework extends from hidden-state directions to individual token positions.

The paper defines tokenwise Fisher-Jacobian saliency using the sensitivity of final target logits to a token’s intermediate representation, weighted by the output Fisher matrix:

Unlike attention magnitude alone, this measure explicitly incorporates the token’s downstream influence on the terminal predictive distribution.

The formal analysis also establishes a causal-support property: when no downstream computational path connects a token position to the target prediction under the standard causal mask, its saliency vanishes. Here, causal describes dependencies in the transformer’s computation graph, not causal identification in the econometric sense.

The authors make these quantities measurable through matrix-free Jacobian operations, stochastic trace estimation, and iterative spectral methods. This avoids constructing prohibitively large curvature matrices, although it does not eliminate the computational expense.

The resulting measurements allow engineers to compare prediction-sensitive directions across layers and evaluate token-update importance using a common terminal objective.

Rank allocation benefits more consistently than universal pruning claims

The paper evaluates whether these measurements improve practical reduction decisions across ten pretrained language models, approximately 1B–9B parameters, using held-out WikiText, OpenWebText, and FineWeb sequences.

For compression, the authors allocate ranks nonuniformly using the spectral curvature retained at each layer. Layers with more distributed prediction-sensitive curvature receive different capacity allocations from layers whose curvature concentrates in fewer directions.

Several models exhibit smaller perplexity degradation than matched-budget uniform SVD and LoRA-style allocation baselines, particularly with spectral-tail thresholds of 0.05 and 0.1 (Section 4.3, Figure 3).

The improvement is architecture-dependent. Uniform allocation remains competitive for some LLaMA and Gemma variants, so the measured geometry does not establish a universally superior compression policy.

The structured pruning results are similarly bounded. Fisher-based token-update selection substantially outperforms random removal at controlled drop fractions and remains competitive with attention, gradient, and other saliency methods (Section 4.4 and Appendix D.2).

Its incremental advantage over strong alternatives is generally more modest than its advantage over random pruning.

These experiments support predictive geometry as an additional selection signal. They do not establish that it replaces simpler criteria in every architecture or under every compression budget.

Geometry-regularized student recovery depends on the distillation objective

A separate experiment examines whether predictive geometry can help compressed student models recover quality during knowledge distillation.

The researchers supplement conventional distillation and cross-entropy losses with a Fisher-weighted hidden-state alignment penalty. The intuition is that a student’s intermediate errors should receive greater penalties when those errors would substantially affect the teacher’s predictions.

The study uses three sub-billion-parameter model families, comparing matched completed configurations across initialization methods, compression ranks, and distillation objectives.

Source: Section 4.5, Table 1. Advantage equals baseline perplexity minus geometry-regularized perplexity; positive values favor geometry. Only completed matched configurations are included.

Geometry improves most reverse-KL and skew-KL configurations. Forward KL has eight positive pairs but negative mean and median advantages, demonstrating why counting wins alone gives an incomplete assessment.

Initialization also matters. Under reverse KL, the mean geometry advantage is +44.5 for SVD initialization but -1.9 for activation-aware SVD initialization.

An additional distinction limits interpretation: the implemented regularizer uses the teacher’s final normalization and language-model head as a readout-Fisher proxy, rather than propagating every intermediate mismatch through the complete downstream teacher network.

Thus, the positive distillation results validate this cheaper approximation in selected settings, not universal benefits from the full theoretical operator. Some recovered students also retain high absolute perplexity, and the experimental grid is incomplete.

Measurement cost determines whether the geometry is worth using

Predictive geometry introduces an economic trade-off beyond model quality.

The representative geometry-estimation configurations use 32 sequences of length 512. Their reported costs range from approximately 0.8 GPU hours for LLaMA-3.2-1B to 6.0 GPU hours for Gemma-2-9B (Appendix F.3, Table 6). These are approximate measurement costs, not inference-time speedups.

For engineers evaluating compression, Cognaptus draws three practical requirements from the evidence:

1. Measure predictive quality under matched budgets. Compare perplexity or cross-entropy after compression against uniform ranks and strong saliency baselines, using held-out evaluation data.
2. Include measurement and recovery costs. The benefit of selective compression must justify geometry estimation, subsequent training, and validation.
3. Evaluate the resulting deployment configuration. Parameter reduction or favorable curvature estimates do not independently establish lower inference latency, memory savings, or acceptable application quality.

The same discipline applies to distillation: choose the objective and initialization through matched experiments rather than assuming geometric regularization is always beneficial.

Preserve predictive capacity, not merely representation capacity

The paper provides a mathematically grounded explanation for why transformer hidden-state dimensions differ in predictive importance. Its local curvature results are stronger than its evidence for universally effective compression, which remains heterogeneous across architectures, interventions, and training objectives.

For model engineering, the appropriate use of predictive geometry is therefore selective. It can help identify which layers and token updates warrant closer protection, and which compressed students may benefit from additional alignment. Whether that information improves an actual deployment depends on the quality retained and the cost of obtaining it.

The central contribution is a measurable connection between intermediate representation changes and final prediction sensitivity. The remaining engineering problem is determining when that connection produces enough additional savings to justify its computational expense.

Cognaptus: Automate the Present, Incubate the Future.