💬 NLP

When Attention Collapses: How Degenerate Layers in LLMs Enable Smaller, Stronger Models

The paper introduces "Inheritune," a training method that leverages the discovery of redundant "lazy layers" caused by attention collapse in large language models to construct smaller, more efficient models that match or exceed the performance of their larger counterparts.

Original authors: Sunny Sanyal, Ravid Shwartz-Ziv, Alexandros G. Dimakis, Sujay Sanghavi

Published 2026-02-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sunny Sanyal, Ravid Shwartz-Ziv, Alexandros G. Dimakis, Sujay Sanghavi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Zombie" Layers

Imagine you hire a team of 48 experts to solve a complex mystery. You expect each expert to add a unique piece of the puzzle.

Experts 1–20: They are brilliant. They analyze clues, connect dots, and find hidden patterns. They are doing exactly what you paid them for.
Experts 21–48: Here is the surprise. As the paper reveals, these later experts have stopped thinking. They have all decided to say the exact same thing: "I think the answer is just the average of everything we've seen so far."

In the world of Large Language Models (LLMs), this is called Attention Collapse.

The Metaphor: Imagine a classroom where the teacher asks, "Who knows the answer?"
- In the early rows (early layers), students raise their hands with specific, different answers.
- In the back rows (deep layers), every single student raises their hand and says, "I think the answer is... everything." They stop distinguishing between specific words and just blur them all together.
The Result: These deep layers are called "Lazy Layers." They take up space and cost money to run, but they aren't actually learning anything new. They are just "zombies" going through the motions.

The Discovery: We Are Wasting Space

The researchers looked at famous models like GPT-2 and found that in many of these deep layers, the "attention" (the ability to focus on specific words) has collapsed into a single, boring pattern. It's like having a 48-story building where the top 24 floors are just empty storage rooms. You are paying for the height, but you aren't using the space.

The Solution: Inheritune (The "Smart Apprentice" Method)

Instead of trying to fix the lazy experts, the authors decided to build a smaller, smarter team from scratch using a recipe they call Inheritune.

Think of it like training a master chef's apprentice:

The Mentor (The Big Model): You have a massive, 48-layer model that has already learned a lot.
The Selection: You don't take the whole team. You only take the top 16 experts (the early layers) who are still sharp and doing great work. You ignore the lazy ones at the bottom.
The Training: You give this small team of 16 experts a fresh set of data and say, "Go learn!"
The Growth: If they get stuck or need more help, you don't just add random new people. You carefully add more layers, but you train them specifically to avoid becoming lazy.

Why This is Magic

Usually, if you shrink a model (make it smaller), it gets dumber. But because the original big model was "bloated" with lazy layers, the small model is actually cleaner and more efficient.

The Analogy: Imagine a marathon runner carrying a heavy backpack full of bricks (the lazy layers). If you take the bricks off, the runner becomes lighter and faster.
The Result: The paper shows that this smaller model (with half the layers) can run just as fast, and often finish the race better, than the giant model that was carrying the extra weight.

The Key Takeaways

Big isn't always better: Deep layers in AI models often stop working properly and just "average out" information.
Lazy layers are useless: They don't hold any special knowledge; they are just dead weight.
Inheritune is the fix: By starting with the "smart" early layers of a big model and training a smaller model to grow without becoming lazy, we can create compact, powerful AI that is cheaper to run and faster to train.

In short: The paper teaches us that sometimes, to build a stronger AI, you don't need to make it bigger. You just need to cut out the lazy parts and train a leaner, sharper version.

1. Problem Statement: Attention Collapse and Lazy Layers

The authors identify a significant structural inefficiency in standard decoder-style Large Language Models (LLMs): Attention Collapse.

Phenomenon: In many pre-trained LLMs (e.g., GPT-2, LLaMA), the attention matrices in deeper layers degenerate into near rank-one structures. Instead of learning complex, discriminative relationships between tokens, these layers perform near-uniform averaging across the sequence.
Lazy Layers: The authors term these structurally degenerate layers "lazy layers." They are redundant, contribute little to the model's representational power, and fail to learn meaningful token interactions.
Root Cause: The paper provides a theoretical analysis (Theorem H.1) linking rank-1 collapse to vanishing gradients. When an attention head concentrates on a single "sink" token (e.g., the first token), the softmax Jacobian saturates, causing gradients for the Query ( $W_Q$ ) and Key ( $W_K$ ) weights to vanish. This makes the collapse "sticky," preventing the layer from recovering its expressiveness during training.
Empirical Evidence: Analysis of GPT-2 (Medium, Large, XLarge) and other modern models (LLaMA-3, Falcon) reveals that a substantial portion of attention heads in deeper layers collapse to rank-1. For instance, in GPT-2 XLarge, many layers beyond the halfway point exhibit near-complete rank collapse.

2. Methodology: Inheritune

To address this inefficiency, the authors propose Inheritune, a novel training recipe designed to build smaller, stronger models by eliminating structural redundancy.

Core Principles:

Potent vs. Lazy: Early layers in deep LLMs typically possess high-rank (potent) attention matrices, while deeper layers often degenerate into lazy layers.
Progressive Growth: Instead of training a small model from scratch or simply pruning a large one, Inheritune initializes a compact model with the potent early layers of a larger reference model and progressively expands it.

The Inheritune Algorithm:

Inherit: Initialize a target model ( $M_{tgt}$ ) with the first $l$ layers (including embeddings and LM head) from a larger pre-trained reference model ( $M_{ref}$ ). These initial layers are selected based on high average rank (potent layers).
Train: Train $M_{tgt}$ on the training dataset for a specified number of steps.
Grow: If the target model's validation loss is worse than the reference model, incrementally add the next contiguous block of layers from the reference model (or initialize new layers) and continue training.
Iterate: Repeat the train-grow cycle until the smaller model matches or surpasses the validation loss of the larger reference model.

Key Distinction from Baselines:

Unlike pruning, Inheritune involves a growth phase and re-training.
Unlike knowledge distillation, it does not rely on soft targets from a teacher but rather inherits structural weights.
Unlike stacking (copying weights to initialize deeper layers of the same model), Inheritune specifically avoids inheriting "lazy" layers, starting instead with the most effective early representations.

3. Key Contributions

Novel Analysis of Attention Degeneration: The paper provides the first systematic empirical and theoretical analysis of attention rank collapse in standard decoder-only LLMs, defining "lazy layers" and demonstrating their functional ineffectiveness.
Inheritune Training Recipe: A simple yet powerful method to initialize and grow smaller models that outperform larger counterparts by bypassing structural inefficiencies.
Comprehensive Evaluation: Extensive experiments across GPT-2 variants (355M to 1.5B parameters) and datasets (OpenWebText, FineWeb) validating the method against random initialization, warm-start baselines, and extended training.

4. Experimental Results

The authors evaluated Inheritune on GPT-2 Medium, Large, and XLarge models, as well as custom variants trained on FineWeb.

Performance vs. Larger Models: Models trained with Inheritune using significantly fewer layers (e.g., a 16-layer model derived from a 24-layer reference) achieved comparable or lower validation losses than their full-sized counterparts trained from scratch for the same number of steps.
Performance vs. Same-Size Baselines: Inheritune-trained models consistently outperformed same-sized models trained from scratch. Notably, they reached the performance of same-sized baselines trained for twice as many steps in roughly half the time.
Downstream Tasks: On zero-shot benchmarks (ARC-E, PIQA, SciQ, HellaSwag, LAMBADA), Inheritune models achieved higher average accuracy than both larger reference models and same-sized random-initialized models.
Mitigation of Collapse: Post-training analysis showed that Inheritune models did not develop lazy layers. Their attention matrices maintained higher ranks throughout the network, confirming that the method successfully mitigates attention collapse.
Initialization Sensitivity: Ablation studies showed that initializing both Attention and MLP submodules together yielded the best results, while initializing only one component was less effective. Initializing with "lazy" layers (deep layers of the reference) performed no better than random initialization.

5. Significance and Impact

Efficiency by Design: The paper challenges the assumption that deeper models are inherently better. It demonstrates that much of the depth in current LLMs is structurally redundant due to attention collapse.
New Scaling Axis: Inheritune offers a new path for model compression and efficiency. It suggests that by removing "lazy" structural inefficiencies, one can create compact models that are more performant per parameter than their larger, inefficient counterparts.
Theoretical Insight: The connection between rank-1 attention collapse, sink tokens, and vanishing gradients provides a theoretical foundation for understanding why deep transformers often fail to utilize their full depth, guiding future architectural designs.
Practical Application: The method is computationally feasible (requiring only standard training steps with a progressive growth strategy) and can be applied to various decoder-style architectures to create more efficient, "attentive" language models.

In conclusion, the paper argues that attention collapse is a pervasive structural flaw in deep LLMs, and Inheritune is an effective solution that leverages the "potent" early layers of large models to train smaller, highly efficient models that avoid these degeneracies entirely.