This paper was formerly known as

• Inheritune: Training Smaller Yet More Attentive Language Models (v2)
• Pre-training Small Base LMs with Fewer Tokens (v1)

Inheritune is an efficient training method for developing smaller, high-performing language models by inheriting knowledge from larger pre-trained models. Our approach addresses a critical inefficiency in standard transformer architectures: attention matrices in deeper layers often degenerate to single-column patterns, creating "lazy layers" that contribute minimally to model performance.

This paper has been published as a journal version in TMLR and was reviewed on OpenReview.

• 📌 Used in the wild: This paper is used by 21 Hugging Face models (HuggingFace page).
• 🌍 Notable adoption: Gumini (multilingual) English and Korean models are trained following the Inheritune recipe (Gumini report).
• 📰 Media feature: Covered by Marktechpost (article).
• 🎥 Video tutorial: A YouTube tutorial by Prince Canuma is available (link)

.
├── GPT2-experiments/   # GPT-2 training/analysis experiments for the paper
├── lit-gpt/            # Code adapted from lit-gpt / small-LM training utilities
├── analysis/           # Attention rank computation, softmax analysis, plotting scripts
├── images/             # Figures used in the README / paper artifacts
├── README.md           # Main project documentation
├── poster.pdf          # Project poster (PDF)
└── attention-collapse-demo/ # toy examples of attention collapse in simple setting

The rank of many attention matrices collapses to near rank-1 during training. Refer demo notebook.

Figure. 1: Toy example of rank collapse in pure self-attention network.

Figure. 2: Rank analysis of GPT-2 Medium (355M).

Figure. 3: Rank analysis of LLaMA-3 8B reveals that nearly half of the attention heads (in red) exhibit rank collapse.

Inheritune follows a simple three-step process:

1. Inherit: Copy potent early transformer layers from a larger pre-trained model.
2. Train: Continually train the inherited layers on the pre-train dataset.
3. Expand: Progressively grow the model until reaching desired performance.

This approach leverages the knowledge already captured in early layers while eliminating lazy deeper layers.

Figure. 4: Overview of the Inheritune training recipe using a 24-Layer GPT-2 Medium model example.

Our 24-layer variant converges faster and matches the validation loss of the full 48-layer model despite being 50% smaller.

Figure. 5: GPT-2 XLarge variants derived using Inheritune converge faster and match the final validation loss of the full-sized model, despite having much fewer layers.

Table 1. Downstream evaluation of GPT-2 XLarge (1.5B) trained from scratch vs. a 24-layer model trained with Inheritune (Ours).

Figure 6: GPT-2 Large^† variants derived using Inheritune converge faster and match the final validation loss of the full-sized model, despite having much fewer layers.

Table 2. Downstream evaluation of GPT-2 Large^† (680M) trained from scratch vs. a 16-layer model trained with Inheritune (Ours).

If you find this work helpful, please consider citing us:

@article{
sanyal2026when,
title={When Attention Collapses: How Degenerate Layers in {LLM}s Enable Smaller, Stronger Models},
author={Sunny Sanyal and Ravid Shwartz-Ziv and Alex Dimakis and sujay sanghavi},
journal={Transactions on Machine Learning Research},
issn={2835-8856},
year={2026},
url={https://openreview.net/forum?id=2zQn0bUoPf},
note={}
}

The training code for small language model 1B-2B is mainly adapted from litgpt. The code for GPT2 experiments are mainly adapted from Sophia and nanoGPT. The llama image was created using DALLE.