Source-linked AI summary
Jump to Conclusions: Short-Cutting Transformers With Linear Transformations
Alexander Yom Din, Taelin Karidi, Leshem Choshen, Mor Geva
TL;DR
Intermediate transformer representations are difficult to interpret and use because predictions rely on final-layer representations. The paper learns linear mappings that cast hidden states across layers, outperforming identity projections for representation and prediction estimation. The method also improves early exit and supports linear sub-module approximation, while experiments are limited to English and the linear-transform choice.
Problem
Only final-layer representations are used for prediction, obscuring the interpretability and utility of intermediate hidden representations.
Method
The method fits context-free linear transformations between hidden representations at different transformer layers to approximate later representations and predictions.
Results
7.9% of GPT-2 layers and 5.4% of BERT layers are additionally saved at 95% accuracy, while mat improves prediction estimation over id across models, scales, and data distributions.
Takeaways & Limitations
Linear mappings enable more faithful inspection of intermediate representations and improve computation efficiency through early exiting.
Takeaways & Limitations
The experiments use only English data, and the paper focuses on linear transformations rather than affine or nonlinear mappings.
Abstract
from arXiv · showhide
Transformer-based language models create hidden representations of their inputs at every layer, but only use final-layer representations for prediction. This obscures the internal decision-making process of the model and the utility of its intermediate representations. One way to elucidate this is to cast the hidden representations as final representations, bypassing the transformer computation in-between. In this work, we suggest a simple method for such casting, using linear transformations. This approximation far exceeds the prevailing practice of inspecting hidden representations from all layers, in the space of the final layer. Moreover, in the context of language modeling, our method produces more accurate predictions from hidden layers, across various model scales, architectures, and data distributions. This allows "peeking" into intermediate representations, showing that GPT-2 and BERT often predict the final output already in early layers. We then demonstrate the practicality of our method to recent early exit strategies, showing that when aiming, for example, at retention of 95% accuracy, our approach saves additional 7.9% layers for GPT-2 and 5.4% layers for BERT. Last, we extend our method to linearly approximate sub-modules, finding that attention is most tolerant to this change. Our code and learned mappings are publicly available at https://github.com/sashayd/mat.
1. Introduction
Transformer language models optimize final-layer representations for prediction, leaving intermediate representations harder to interpret and use. The paper learns linear mappings across layers, improving representation and prediction estimates while supporting efficiency gains through early exit and sub-module approximation.
- Motivation: Hidden representations matter because interpreting them can illuminate transformer decision-making, while predicting from them can reduce computational cost.Only final representations ordinarily form predictions, although every layer creates new representations.
- Prior approach: The prevalent output-embedding projection is effective but oversimplifies computation by assuming all layers operate in the same space.The paper addresses this limitation by learning transformations between layer-specific representation spaces.
- Method: The method fits a linear regression from hidden representations at layer ℓ to layer ℓ′, replacing identity propagation between layers.The context-free, single-representation mapping substantially improves approximation over id and suggests greater linearity in transformer inference than identity mappings reveal.
- Empirical motivation: 15%-40% accuracy gains at most layers show that mat estimates language-model predictions more accurately than naive id projections across GPT-2 and BERT experiments.The comparison spans two data sources and varied model scales and architectures, with mappings also generalizing across data distributions.
- Efficiency: 7.9% of GPT-2 layers and 5.4% of BERT layers are additionally saved at 95% accuracy when early exit uses mat instead of id.Early exit dynamically stops inference at a selected layer and predicts from that layer’s representation.
- Sub-module analysis: Linear approximation of attention loses the least precision among attention, FFN, and layer-normalization sub-modules, suggesting possible compute reduction through parallelizable non-contextual inference.The sub-module analysis applies the same linear-mapping methodology to each component.
3. Linear Shortcut Across Blocks
The method learns linear mappings that cast hidden representations from layer ℓ to a later layer ℓ′, bypassing intervening computation. Across GPT-2 and BERT, these mappings approximate target-layer representations better than directly reading hidden states in a shared space.
- 3.1. Method: Linear regression maps a hidden representation from source layer ℓ to any subsequent target layer ℓ′ while skipping intervening transformer blocks.The mapping is learned from corresponding hidden-representation pairs collected across input sequences and positions.
- 3.2. Baseline: The id baseline leaves hidden representations unchanged and assumes that all layers operate in the same linear space.This contrasts with mat, which learns transformations between layer-specific representation spaces.
- 3.3. Quality of Fit: The evaluation fits mappings on Wikipedia sentences and measures uniformly coordinate-averaged r2-scores on held-out validation examples.For BERT, representations are extracted at masked-token positions because predictions are obtained from masked tokens.
- 3.3. Quality of Fit: mat consistently achieves higher coordinate-averaged r2-scores than id across GPT-2 and BERT layer pairs.The comparison evaluates mapped representations against those produced by full inference.
- 3.3. Quality of Fit: The advantage of mat is especially pronounced in BERT, where id fails to map representations between most layers.This indicates that hidden representations can change substantially from one transformer block to the next.
4. Linear Shortcut for Language Modeling
The study evaluates whether linear mappings can cast intermediate representations into final-layer substitutes for language-model prediction. Across GPT-2 and BERT, mat generally improves predictive agreement over identity mappings, while alternation schemes can further improve selected layers.
- Next Token Prediction: GPT-2 mat improves Precision@k and Surprisal over id across layers, with Precision@1 exceeding id by more than 20% through layer 44.The evaluation compares substitute predictions with final-layer outputs using Precision@k and Surprisal.
- Masked Token Prediction: For BERT masked-token predictions, Table 1 compares top-5 outputs from mat and id at layers 4, 12, and 24, marking human-judged contextually plausible tokens.At the final layer, mat and id produce the same predictions.
- Masked Token Prediction: In BERT, mat again outperforms id, whose precision is near zero across k in the first ten layers, while mat reaches 8%-52% precision there.Mat improves id’s Precision@1 by more than 17% at most layers, and more than 25% of predictions appear from layer 3 onward.
- Alternation Schemes: Alternating transformer inference with linear mappings can outperform direct mat mappings, although the best-r2-score selection works well early and under-achieves in the second half of layers.The experiments compare r2-informed and weighted round-robin schemes on 24-layer GPT-2 using Precision@1.
5. Method Robustness
Across model scales and data distributions, mat remains more effective than identity mapping for casting intermediate representations into predictive final-layer substitutes.
- Robustness Across Model Scales: Across GPT-2 and BERT scales, mat consistently achieves substantially higher predictive scores from intermediate layers than id.The trend holds across GPT-2 models from 12 to 48 layers and BERT models from 12 to 24 layers.
- Robustness Across Data Distributions: On a second news-sentence distribution, the study tests whether mappings learned from one distribution transfer to another using 24-layer GPT-2 and BERT.The news corpus contains 9,000 training examples and 1,000 validation examples.
6. Implication to Early Exiting
The method improves early exiting by transforming intermediate representations before prediction, enabling dynamic computation savings while maintaining a target accuracy.
- Early-exit setup: Early exit dynamically stops inference when a confidence criterion is met, using a cast intermediate representation to derive an output probability distribution.The confidence threshold varies through λ, and efficiency is measured by the average number of processed transformer layers.
- Early-exit setup: Mat is evaluated as a plug-and-play replacement for id in GPT-2 next-token and BERT masked-token prediction, alongside fixed-exit baselines.Both prediction accuracy and computational efficiency are evaluated.
- Early-exit results: At 95% average precision, mat saves 13.8% of GPT-2 layers and 20% of BERT layers, versus 5.9% and 14.6% for id.These results use 24-layer GPT-2 and BERT and compare average Precision@1 with average layers processed.
- Early-exit results: For both models and both mapping methods, dynamic early exit outperforms stopping after a fixed number of layers.The comparison is reported in the Precision@1-versus-layers-processed evaluation.
7. Linear Shortcut Across Sub-Modules
The linear-shortcut analysis extends beyond whole transformer blocks: individual sub-modules can also be approximated, with attention showing the greatest tolerance.
- Method: The method fits linear regressions from each sub-module’s input to output and uses them to define mappings across layers.The analyzed components are attention, FFN, and layer-normalization sub-modules.
- Results: From approximately layer 7 onward, every sub-module mapping achieves better Precision@1 and Surprisal than the full-block mapping mat.The comparison is made on 24-layer GPT-2 using mappings for attention, FFN, and layer normalization.
- Results: Attention is less harmed by linear approximation than FFN or layer normalization.The authors suggest that replacing FFN or normalization may erode self-attention computation after several layers.
- Interpretation: The strong attention results suggest that contextualization often exhausts itself in early layers, although late-layer attention may still matter in more delicate cases.This interpretation is presented speculatively by the authors.
8. Related Work
The paper builds on work using intermediate representations for interpretation and efficiency, while distinguishing its cross-layer linear mappings from direct propagation and related linearity methods.
- Interpretability: Prior research uses hidden representations to study prediction construction or stored features, whereas this work converts intermediate representations into final-layer form.The paper positions this conversion as a distinct interpretability approach.
- Early exiting and efficiency: Unlike prior early-exit and parallel-inference methods that propagate intermediate representations directly, this work transforms them before prediction and reports better performance than that baseline.The paper also notes that its method trains considerably fewer parameters than some alternatives.
- Linearity and layer skipping: Related work has examined layer skipping and transformer-component linearity, while Belrose et al. (2023) train affine transformations from hidden to final representations.The paper differs by training linear transformations across all layers.
- Contribution: The authors conclude that context-free, token-uniform mappings improve representation and prediction approximation and can support efficiency gains and parallel computation.Their sub-module experiments motivate breaking transformer computation into parallel tasks.
10. Limitations
The study identifies four limitations: it restricts mappings to linear transformations, does not explain the observed linear structure, evaluates only post-hoc use, and uses only English data.
- The method focuses on linear transformations, leaving affine and nonlinear mappings for future investigation.This choice is motivated by the success of the identity baseline and the question of whether transformer inference contains exploitable linearity.
- The study finds additional linear structure beyond residual connections but does not explain its cause.
- The experiments evaluate post-hoc interpretability without changing model weights, not mappings integrated during training.Training-time integration remains a direction for future work.
- The experiments use only English data, although the authors expect the language-agnostic method may generalize elsewhere.
A. Descriptions of mat_attn, mat_ffn and mat_ln1_ln2
The sub-module mappings replace selected transformer components with learned linear approximations between layers. Attention approximations operate position-wise and remove contextualization across the mapped layers, whereas feed-forward and layer-normalization replacements retain self-attention.
- The mappings mat_attn, mat_ffn, and mat_ln1_ln2 provide linear replacements for attention, feed-forward, and layer-normalization sub-modules.
- Applying a sub-module mapping replaces each corresponding component in every intermediate block between source layer ℓ and target layer ℓ′.
- The attention approximation operates independently at each position, so mat_attnℓ→ℓ′ disables contextualization between layers ℓ and ℓ′.
- The feed-forward and layer-normalization approximations retain the self-attention sub-modules and therefore continue to operate contextually.