Source-linked AI summary
The Remarkable Robustness of LLMs: Stages of Inference?
Vedang Lad, Jin Hwa Lee, Wes Gurnee, Max Tegmark
TL;DR
The paper asks how LLM computation is organized across depth despite evidence for both iterative refinement and specialized circuits. It studies layer deletion and adjacent-layer swaps across model families, finding localized sensitivity that motivates a four-stage inference framework. The framework describes detokenization, feature engineering, prediction ensembling, and residual sharpening, while remaining an aggregate account with approximate stage boundaries.
Problem
Existing accounts emphasize either iterative layer-wise refinement or specialized circuits, leaving their relationship unclear.
Method
The study deletes layers, swaps adjacent layers, and synthesizes behavioral and mechanistic evidence into a four-stage framework for decoder-only LLMs.
Results
Interventions reveal a repeatable depth-wise structure across architectures and scales, with detokenization, feature engineering, prediction ensembling, and residual sharpening as the proposed stages.
Takeaways & Limitations
LLM inference appears to be an ordered composition of specialized computational regimes rather than a flat pipeline.
Takeaways & Limitations
Stage boundaries are approximate, stages may co-occur within layers, and the framework captures aggregate trends rather than individual token paths.
Abstract
from arXiv · showhide
We investigate the robustness of Large Language Models (LLMs) to structural interventions by deleting and swapping adjacent layers during inference. Surprisingly, models retain 72-95% of their original top-1 prediction accuracy without any fine-tuning. We find that performance degradation is not uniform across layers: interventions to the early and final layers cause the most degradation, while the model is remarkably robust to dropping middle layers. This pattern of localized sensitivity motivates our hypothesis of four stages of inference, observed across diverse model families and sizes: (1) detokenization, where local context is integrated to lift raw token embeddings into higher-level representations; (2) feature engineering, where task- and entity-specific features are iteratively refined; (3) prediction ensembling, where hidden states are aggregated into plausible next-token predictions; and (4) residual sharpening, where irrelevant features are suppressed to finalize the output distribution. Synthesizing behavioral and mechanistic evidence, we provide a framework for interpreting depth-dependent computations in LLMs.
1 Introduction
The paper frames LLM inference as a tension between iterative refinement and specialized circuits, then investigates depth-dependent computation through layer interventions and proposes four stages.
- Motivation: Mechanistic interpretability contrasts iterative hidden-state refinement with specialized circuits, leaving their integration unresolved.Iterative inference emphasizes overlapping updates, while circuit-oriented findings identify distinct functional components.
- Related evidence: Prior layer-wise findings—including middle-layer knowledge storage, fine-tuning effects, and sparsity transitions—suggest a structured computation that remains incompletely characterized.These observations motivate integrating behavioral and mechanistic evidence.
- Approach: The study deletes individual layers and swaps adjacent layers to characterize localized effects and bridge top-down and bottom-up views of decoder-only computation.These interventions are synthesized with prior interpretability findings into a depth-wise framework.
- Four-stage framework: The proposed framework hypothesizes four stages: detokenization, feature engineering, prediction ensembling, and residual sharpening.Early layers integrate local context, middle layers construct features, later layers produce predictions, and final layers suppress noisy components.
2 Experimental Protocol
The experimental protocol compares five decoder-only model families across scales using standardized preprocessing and large random-token samples, then measures responses to layer swaps and zero ablations.
- Models: Experiments cover Pythia, GPT-2, Qwen 2.5, LLaMA 3.2, and Microsoft Phi models ranging from 124M to 6.9B parameters.All are decoder-only transformers, but their attention and MLP execution orders differ.
- Data: Evaluation uses one million tokens from random Pile sequences unless an experiment specifies otherwise.
- Interventions: Layer swaps execute adjacent transformer blocks in reversed order, while recorded outcomes include KL divergence, loss, top-1 accuracy, entropy, and benchmark performance.The swap at layer ℓ runs block ℓ+1 before block ℓ.
- Interventions: Zero ablations provide baselines for each swap while preserving the corresponding swap ordering.Zero ablation is used instead of mean ablation to maintain consistency with the intervention order.
3 Robustness
Layer interventions reveal strong depth-dependent robustness: first and final layers are especially sensitive, whereas middle layers tolerate deletion and adjacent reordering.
- Intervention results: Middle layers are remarkably robust to both deletion and minor order changes.This pattern appears in intervention metrics and similar benchmark-task trends.
- Intervention results: Figure 5 compares swap and drop effects using KL divergence and top-1 consistency, alongside CKA-based representational similarity across layers.The CKA plots show block-like structure in GPT-2 XL and Pythia 2.8B, with similar trends reported across other families and sizes.
- Intervention results: Swapping adjacent intermediate layers is less harmful than ablating them, suggesting that some forward-pass operations may be commutative.The commutativity interpretation requires further experimentation.
- Intervention results: First-layer interventions are catastrophic across every model family and size, producing very high-entropy predictions.Some models show similarly catastrophic effects when swapping the final and penultimate layers, while GPT-2 largely preserves predictions.
4 Stages of Inference Hypothesis
Layer-wise interventions and mechanistic analyses support four depth-dependent inference stages: early detokenization, intermediate feature engineering, prediction ensembling, and final residual sharpening.
- Overview: CKA reveals block-like, depth-dependent representation structure consistent with shared computation motifs across adjacent layers and staged robustness to interventions.These blocks align with the observed localized effects of layer deletion and swapping.
- Stage 1: Detokenization: Early layers integrate nearby tokens into coherent entities through local attention and subjoiner heads, supporting a detokenization stage.Subjoiner heads transfer information from earlier tokens, especially when the final token completes a multi-token word.
- Stage 2: Feature Engineering: Intermediate layers construct semantic features: WiC probe accuracy peaks near the midpoint, while intermediate predictions remain high-entropy before later consolidation.This gap indicates that linearly accessible features are present before they are fully used for confident next-token prediction.
- Stage 3: Prediction Ensembling: Mid-to-late layers increasingly convert semantic features into concrete predictions as intermediate distributions converge toward the final output.Prediction neurons emerge around the midpoint, and the sharpest KL decline coincides with their density peak near 85% of model depth.
- Stage 3: Prediction Ensembling: Prediction-neuron ensembles outperform individual neurons and the model average, supporting a prediction-ensembling mechanism involving overlapping computational pathways.The model region is also highly robust to layer interventions, consistent with redundancy or self-repair.
- Stage 4: Residual Sharpening: In the final layers, suppression neurons rise as prediction-neuron density declines, potentially removing obsolete features and down-weighting improbable tokens.This late mechanism coincides with continued output convergence and decreasing entropy.
5 Related Work
Prior work frames transformer computation through both mechanistic circuits and iterative, depth-dependent refinement, while ablation studies document substantial redundancy.
- Mechanistic interpretability uses circuit analysis to identify reusable transformer mechanisms, including circuit reuse, self-repair, function vectors, and long-context retrieval.
- The iterative inference hypothesis describes layers as incrementally updating representations and predictions, with logit-lens analyses tracking these changes across depth.
- BERT layer-ablation studies find substantial redundancy, with many neurons, attention heads, and feedforward components removable at limited accuracy cost.
6 Concluding Remarks
The paper proposes a four-stage framework for LLM inference, supported by behavioral and mechanistic analyses, while acknowledging that its depth-dependent boundaries are approximate.
- Residual architecture may make models resilient to local layer interventions by distributing computation across multiple pathways.This interpretation is linked to prediction and suppression neurons, self-repair, and overlapping representations.
- The framework identifies detokenization, feature engineering, prediction ensembling, and residual sharpening as an ordered composition of specialized computational regimes.It is grounded in analyses of structural interventions, attention patterns, neuron function, and residual-stream dynamics.
- The framework is intended as a unifying perspective for interpreting, auditing, and intervening on language models.The authors position it as a foundation for deeper investigation of emerging capabilities such as reasoning.
- Stage boundaries are approximate, stages may co-occur within a layer, and individual tokens may follow distinct processing paths.The framework captures aggregate trends and does not isolate the factors behind model-specific differences.
A.1 Experiment Diagram
The experiment diagram defines two layer-wise interventions: fully removing a layer or swapping the execution order of adjacent layers.
- The study applies zero ablation by fully removing a layer so residual connections skip it entirely.
- The study applies adjacent layer swaps by permuting the positions of two neighboring layers.The ablation is performed across all layers, while the swap intervention concerns adjacent layers.
A.2 Centered Kernel Alignment (CKA)
The CKA analysis measures similarity across layers using last-token representations sampled from the Pile dataset.
- CKA is computed across layers from last-token representations sampled from Pile sequences up to 512 tokens.The analysis uses batch size 128 and unbiased CKA.
A.3 Benchmark Tasks Performance After Layer-Wise Intervention
Benchmark results show catastrophic degradation after intervening on the first layer and, in some models, the final layer, while intermediate layers remain comparatively robust.
- First- and last-layer interventions generally cause catastrophic benchmark deterioration, whereas intermediate-layer interventions preserve robust performance.This trend is observed on HellaSwag, ARC-Easy, and LAMBADA; the figures mark each model’s baseline with a dotted horizontal line.
- The benchmark evaluation covers HellaSwag, ARC-Easy, and LAMBADA using models modified by layer-wise interventions.
A.4 Cosine Similarity Analysis of Swapped Layers
Cosine-similarity analysis of adjacent-layer swaps shows that layer components often preserve their projections despite relocation, while stage-specific patterns reveal iterative computation and specialized later-layer roles.
- The analysis compares self-similarity, index similarity, and adjacent similarity to characterize layer stubbornness, positional correspondence, and neighboring overlap.These metrics are used to assess commutativity and the nature of computations across inference stages.
- High self-similarity across components indicates that swapping a layer’s position often preserves how it projects onto the residual stream.This pattern appears in both Pythia 1.4B and GPT-2 XL, despite smoother trends in GPT-2.
- Attention Heads: Attention heads diverge across model families: Pythia heads become more position-sensitive, whereas GPT-2 heads become increasingly redundant.Pythia attention-head metrics converge toward orthogonality, while GPT-2 metrics converge toward similarity in the latter half of the network.
- MLPs: In stage 2, MLP index similarity exceeds adjacent and self-similarity, indicating overlap with the layer’s previous neighbor and supporting iterative inference.The overlap is more pronounced in Pythia, possibly because its fewer layers must complete stage 2.
- MLPs: In stage 3, MLP metrics except self-similarity converge toward orthogonality, while high self-similarity suggests specialized computations.Together, these patterns indicate that stage-3 MLPs retain their own projections while differing from neighboring or replaced layers.