Disentangling Representation Evolution in Transformers through Directional Decomposition

AI summary

Overview

Research area: Interpretability and geometric analysis of Transformer language models, with applications to model compression and pretraining.

Technical level: Intermediate. Readers need basic familiarity with Transformer architecture (residual streams, attention value vectors, MLP sub-layers) and elementary vector geometry (projections, orthogonality), but the core ideas are intuitive once those pieces are in place.

Scope: The paper proposes an orthogonal decomposition of learned Transformer updates into "parallel" (magnitude-scaling) and "perpendicular" (direction-changing) components, then tests the behavioral importance of each across inference-time editing, post-training compression, and from-scratch pretraining.

What This Paper Is About

Every Transformer layer adds a learned update to the running hidden state. That update can be split into two geometric pieces: a part pointing in the same direction as the incoming state (which just rescales it) and a part pointing in new directions (which changes what the representation means). The paper asks whether the rescaling part is actually necessary, given that residual connections already scale representations for free, or whether it is functionally redundant. The authors answer this by surgically scaling each component in two different representation spaces and measuring the damage.

Key Contributions

1. A directional decomposition framework that splits learned Transformer updates into parallel and perpendicular components, applied both in residual space (sub-layer updates relative to the incoming hidden state) and in attention value space (the pre-output-projection aggregate relative to the token's own value vector).

2. The exclude-self value-space intervention, which separates a token's direct self-message from cross-token aggregation so that parallel scaling can be tested on contextual information alone without destroying the token's identity carrier.

3. Empirical characterization of a directional and site-dependent asymmetry: perpendicular edits are acutely damaging everywhere, while parallel edits are broadly tolerated — most of all in value space when the self-message is preserved.

4. Extensions beyond inference-time editing: perpendicular compression error predicts compressed-model quality better than parallel error, and suppressing parallel attention components during pretraining improves validation loss and downstream accuracy across model scales from 296M to 2.7B parameters.

Main Findings

• Parallel components are pervasive. Across pretrained models like Qwen3-4B and Qwen3-30B-A3B, learned updates consistently contain large components aligned with the incoming representation — ratios of parallel to perpendicular magnitude above one — persisting across depth and generation steps.

• Perpendicular edits are hyper-fragile. Scaling the perpendicular component away from 1 causes steep perplexity surges immediately; complete removal causes catastrophic breakdown, reaching tens of thousands of points in value space and residual attention and millions in residual MLP.

• Parallel edits are broadly tolerated. Model behavior stays near baseline across a wide range of parallel scaling factors, consistent with the interpretation that parallel updates act as adaptive gain control rather than carrying new semantic content.

• Robustness is site-dependent, not a generic property of "parallelness." Value-space exclude-self scaling stays nearly flat across a parallel scale range of 0 to 1 (under a point of perplexity change on Qwen3-1.7B, at most 1.3 points up to scale 3). Residual attention is intermediate; residual MLP is the most fragile.

• Preserving the self-message is essential. Naive parallel removal on the full value aggregate drops average zero-shot accuracy by 8 to 10 points. Exclude-self removal trails baseline by only 1.5 points on Qwen3-1.7B and matches it within 0.1 points on Qwen3-30B-A3B, with slight gains on several individual tasks.

• Long-context behavior follows the same pattern. On Llama-3.2-3B across 4k to 12k tokens (13 RULER tasks), full-aggregate parallel removal at scale 0 collapses by 26 to 40 points, while exclude-self removal retains over 75 percent at 4k and stays substantially more resilient at 12k.

• Stability is not about minimizing perturbation size. Auditing the post-output-projection signal shows the exclude-self value edit perturbs outputs more than residual-space removal (49.10 percent versus 25.15 percent of unedited branch norm), yet preserves capabilities better. What matters is preserving self-identity routing and residual propagation, not isotropic distance.

• Perpendicular compression error tracks quality. Across 4-bit AWQ and Wanda pruning (unstructured, 4:8, 2:4), perpendicular error curves cleanly rank compression regimes in line with downstream fidelity, while parallel error curves interleave and fail to discriminate. Perpendicular distortion dominates total compression error (r > 0.97).

• Suppressing parallel attention during pretraining helps. From-scratch GPT-style models on OpenWebText show lower validation loss trajectories from early training, with both interventions improving downstream accuracy and the value-space variant strongest (+0.7 points at 1.4B, +1.5 points at 2.7B).

Methodology in Plain English

The researchers start from a simple geometric fact: any update vector can be split into one piece that points along a reference vector and another piece that is perpendicular to it. They apply this split in two places inside a Transformer. The first is the residual stream, where the reference is the hidden state coming into a sub-layer and the update is what that sub-layer adds. The second is inside attention, where the reference is the current token's own value vector and the update is the aggregated value message that would be passed to the output projection.

They then run controlled experiments during the forward pass, leaving all weights frozen. They multiply the parallel piece by one scale factor and the perpendicular piece by another, sweeping each independently and measuring perplexity and benchmark accuracy. Because a token's self-message is naturally parallel to its own value vector, naive parallel removal in value space wipes out the token's identity signal along with contextual information. To avoid this confound, they introduce an "exclude-self" version that decomposes only the non-self (cross-token) aggregate and restores the self-message unchanged.

They also show that these edits can be viewed as changes to the attention diagonal, giving closed-form expressions for what the edit does to the attention map. For compression, they take the difference between compressed and uncompressed updates and decompose that error relative to the baseline update direction, comparing how well parallel versus perpendicular error predicts model quality. For pretraining, they train GPT-style models from scratch with parallel removal enforced during both training and evaluation, so there is no mismatch between the two.

Why This Matters

Impact on research. The work reframes how researchers think about what Transformer layers actually compute: not just "what information is added" but "in which direction relative to existing representations." It argues that isotropic L2 error — the standard objective and diagnostic in quantization and pruning — conflates benign magnitude shifts with disruptive semantic rotations, and that directional error is the more causally relevant quantity. It also offers a geometric explanation for why architectures that decouple contextual routing from magnitude modulation (gated attention, query-key normalization) work well.

Real-world applications:

• Model compression tooling. Quantization and pruning pipelines could adopt perpendicular error as a selection criterion, potentially producing smaller models with better preserved capability at the same bit-width or sparsity.

• Deployment-time model editing and pruning diagnostics. Practitioners can use parallel/perpendicular decomposition to predict which compressed checkpoints will actually behave well before running expensive evaluations.

• Efficient pretraining. Suppressing redundant parallel attention components during training offers a low-cost inductive bias that improves downstream accuracy, relevant to anyone training language models at scale.

• Architecture design. The findings support decoupling self-representation amplification from cross-token contextual mixing, informing attention variants and normalization schemes.

Industry relevance. Training and serving large language models is dominated by cost. A diagnostic that predicts post-compression quality without full evaluation, plus a pretraining intervention that improves generalization at 296M to 2.7B scales, is directly relevant to teams doing quantization, distillation, and efficient training. The observation that attention spends parameter capacity on redundant rescaling also suggests architectural simplifications with real compute implications.

Future Directions

• Test generalization beyond decoder-only LMs. The authors note their experiments cover decoder-only language models; whether the same directional patterns hold for encoder-decoder models, vision Transformers, or multimodal architectures is untested.

• Turn analysis tools into practical methods. The edits currently serve mainly as diagnostic probes. Converting the directional decomposition into a deployable training objective or a compression-aware loss requires further validation.

• Build directional objectives into compression. Given that perpendicular error predicts degradation, an open question is whether directly penalizing perpendicular distortion during quantization-aware training or pruning produces better compressed models than minimizing uniform L2 error.

• Connect layer-level geometry to behavior. The paper measures geometry at the layer level and capability at the output level; bridging these more directly — attributing specific downstream behaviors to specific perpendicular subspaces — remains open.

• Refine large-scale pretraining interventions. The gains at 1.4B and 2.7B are promising but modest; whether the inductive bias compounds at frontier scales, and how it interacts with gated attention or other modern components, is unresolved.

Target Audience

This paper is most useful to interpretability researchers studying how Transformer representations evolve across layers, and to engineers working on model compression, quantization, and efficient pretraining who want a more predictive diagnostic than global L2 error. It also speaks to architecture researchers interested in why gated attention and normalization choices help, and to graduate students with a working knowledge of Transformer internals who want an example of using simple linear-algebra geometry to generate testable hypotheses about model behavior.

Authors’ abstract

Transformer representations evolve through learned additive transformations that either preserve their current direction or redirect it. We study this evolution as a functional geometry, decomposing learned updates into parallel and perpendicular components. Across pretrained models, we find substantial parallel components beyond the residual identity path. We then apply the decomposition in two spaces: to attention and MLP updates relative to the hidden state, and to attention value aggregation relative to the current token's value. Targeted edits reveal a strongly space-dependent asymmetry: exclude-self value-space parallel manipulation is markedly more robust than residual-space and perpendicular counterparts, preserving the direct self message while scaling only the non-self aggregate. The same decomposition gives a component-resolved description of compression-induced update error: perpendicular error separates compression methods more clearly than parallel error. Extensive experiments further demonstrate that full-aggregate parallel suppression during from-scratch pretraining lowers validation-loss trajectories and improves downstream averages, with the value-space variant strongest. Together, these results connect representation geometry to editing robustness, compression diagnosis, and training-time intervention. Code is available in the \href{https://github.com/Shwai-He/Transformer-Geometry}{project repository}.