Inheritune: Training Smaller Yet More Attentive Language Models

Authors: Sunny Sanyal, Ravid Shwartz-Ziv, Sujay Sanghavi, Alex Dimakis
Primary area: foundation or frontier models, including LLMs
Keywords: Large Language Models, Small Language Models, Attention degeneration, Efficient training, Model Initialization

Abstract:
Large Language Models (LLMs) have achieved remarkable performance across various natural language processing tasks, primarily due to the transformer architecture and its self-attention mechanism. However, we observe that in standard decoder-style LLMs attention matrices degenerate to single-column for deeper layers. Layers in this state unable to learn anything meaningful and mostly redundant; we refer to these as lazy layers. The goal of this paper is to train smaller models by eliminating this structural inefficiency without compromising performance.

Motivated by this observation, we propose Inheritune, a simple yet effective training recipe for developing smaller, high-performing language models. Smaller models trained with Inheritune inherits early transformer layers from a larger pre-trained model, then retrains and progressively expands the smaller model until it matches or exceeds the performance of the larger model. We demonstrate that Inheritune enables the training of various sizes of GPT-2 models on datasets like OpenWebText-9B and FineWeb\_Edu. Models trained with Inheritune, despite having significantly fewer layers, match or even surpass the performance of their larger counterparts. For instance, our 16-layer GPT-2 medium variant achieves comparable performance to the standard 24-layer GPT-2 medium model.

Decision: reject
Reviewer scores: 3, 6, 6, 5

Reviews and discussion:
<thread>
creation date:
1730635147193

thread_type:
official_review

<thread_object>
creator: reviewer_fevh

<content>
summary:
This paper highlights attention degeneration, where layers become rank-1 matrices, leading to ineffective learning in "lazy layers." The proposed method initializes a smaller model with a few layers from a larger pre-trained model and progressively adds layers, enabling efficient training of high-performing models. Experiments across various GPT-2 sizes show that Inheritune maintains or improves performance compared to larger models trained from scratch or with other initialization methods, establishing it as an effective strategy for compact language models.

strengths:
The paper is well-written and easy to understand, with a logical flow that supports the arguments presented. Additionally, it demonstrates the superiority of the proposed method through a variety of experiments.

weaknesses:
It seems a bit too intuitive to refer to this as an initialization method utilizing the lazy layer phenomenon, even though this paper has identified the occurrence of lazy layers. The logic of using early layers for initialization because of lazy layers feels somewhat weak. When I consider that there are 36 teacher layers and I need to initialize 18 layers, it appears quite obvious to select the first 18 layers, as they are likely to be more compatible with the embeddings.

The experiments presented in the paper are largely conducted with a limited number of steps, which raises questions about the actual performance. It seems that the approach does not significantly differ from simply trimming a larger model and fine-tuning it, leading me to wonder whether the results truly reflect superior performance.

questions:
(1) I think the phenomenon of lazy layers is similar to the attention sink phenomenon. Does approaching rank 1 mean that attention is concentrated on a specific token?

(2) In Figures 1 (c) and (f), rather than suggesting that the lazy layers are unnecessary, I suspect that the early layers are more compatible with the embeddings. Could you perhaps compare using the first nine layers with using the remaining ones randomly?

(3) This paper adopts a three-step approach of Inherit, Train, and Grow. What are the advantages of this method compared to reaching the final size all at once?

(4) I would like to inquire about the comparison between the GPT-2 Medium 16-layer variant trained from scratch and the one trained with Inheritune in Figure 2, 9. Does Inheritune effectively prevent rank collapse?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686192889

thread_type:
official_review

<thread_object>
creator: reviewer_bigz

<content>
summary:
This paper studies and proposes a better way to create smaller LLMs from existing pretrained models, in particular by examining whether the learned attention layers exhibit meaningful attention patterns or degenerate into single-column structures (“lazy layers”). They find later layers are more lazy across GPT-2 LLMs, and propose Inheritune as a method to exploit this phenomenon. Inheritune works by initializing a smaller model with the first half of the layers from a pre-trained larger LLM, training this model, and then progressively growing it by adding more blocks if necessary. This process is repeated until the smaller model's performance matches or surpasses that of the reference model. Through experiments using GPT-2 models of varying sizes, on datasets like OpenWebText-9B and FineWeb_edu, and using evaluations including language modeling perplexity and downstream LM Evaluation harness tasks, the authors demonstrate that Inheritune consistently outperforms various baselines, including larger models trained from scratch, models initialized with different techniques (stacking, hybrid stacking, half-width), and models trained using knowledge distillation. They further provide interesting analysis into Inheritune's behavior via various ablations on target modules and pretraining data mix.

strengths:
**Interesting study + method contribution**
I liked the combination of studying a natural phenomenon (lazy layers showing up in pretrained decoder LLMs), and further exploiting this to create a well-motivated method for efficiently creating smaller models. 

**Comprehensive Task Evaluation**

I appreciated how the authors evaluated not just validation-set perplexity, but also zero-shot performance on WikiText and Lambada in main experiments, and further considered popular LM Evaluation Harness tasks.

**Good Performance**

The experiments demonstrate Inheritune's strong performance compared to baselines such as randomly initializing models, stacking, hybrid stacking, and knowledge distillation.


**Interesting Ablation Studies for In-depth Analysis**

I appreciated the study into the different aspects of Inheritune, including initializing different submodules, using different pretraining data mixes.

weaknesses:
**Insufficient treatment of related work / first contribution novelty**

How does the notion of "lazy layers" relate to prior work such as Attention Sinks [1]? In particular, the claim that “Notably, this phenomenon has not been studied in the context of standard LLMs.” (L043) does not seem true? See Attention Sinks [1] at ICLR 2024, which studied and observed before that several standard decoder-only LLMs (Llama 2 7B, Pythia 12B, Falcon 7B, MPT-7B) display this “lazy layer” phenomenon. 

The method and way to exploit this lazy layer phenomenon seem novel, but I do think there needs to be better discussion on how the lazy layer contribution is novel.  

**Limited model diversity (model family and scale)**

The study and method are limited to a single model family (GPT-2), where the largest LLM evaluated is only 1.5B parameters. Overall, the paper would be stronger if the findings were shown beyond a single class of models, and on a variety of more modern and popular models (e.g., Llama 3 8B, Mistral 7B, Phi 1.5, Gemma, etc.). While I don't fault the authors for not evaluating 7B models due to budget constraints  (although doing so would increase my score); even under the 1.5B parameter budget, we have models available such as Phi 1.5 from 2023 [2].

**Motivation behind comparisons**
I found some of the comparisons a bit too "baseline", but this could be clarified with discussion on their motivation. 
* For example, if the motivation of Inheritune is to save model memory, how does the method compare to quantization techniques that can drastically reduce the parameter memory? 
* If using less layers can improve inference or generation efficiency (as a benefit over quantization techniques), it would be good to see this benchmarking analysis


[1] Efficient Streaming Language Models with Attention Sinks, https://arxiv.org/abs/2309.17453 

[2] Textbooks Are All You Need II: phi-1.5 technical report https://arxiv.org/abs/2309.05463


**--- After rebuttal revisions ---**  
I think the authors did a reasonable job clarifying the distinction between lazy layers and attention sinks, while also extending their analysis to pretrained LLMs outside of GPT-2 architectures trained on OWT. 

I still think the paper's presentation + method could be improved, e.g., explicitly using the number of observed  lazy layers to recommend how many layers should be kept (instead, they keep this main design choice for pruning at `num_layers // 2`, which makes the lazy layer connection a bit vacuous imo). 

For this I am willing to raise my score to 6.

questions:
Why are the target / smaller models initialized as k/2, if k is the number of layers in the original model? This seems a bit ad-hoc. Was there any study into if lazy layers were much more frequent in the second-half of the LLMs? Plotting or visualizing this would be support this design choice more.   
* Because you have the LLMs, some analysis into whether a layer is lazy or not, and using this to inform which layers should be preserved, might also improve the method quality. 

Is the step-count comparison fair? To make an Inheritune model, we need to have a full pretrained model in the first place. Are the steps it takes to acquire this pretrained model factored into the total step comparisons? 
* Conversely, when making the claim that Inheritune outperforms randomly initialized models with the original parameter counts, are these full models further trained to hit the same total training updates that it takes to create Inheritune models? 

In Figure 2, if the GPT2 models are decoder-only / causal, should the attention weights only be limited to lower-triangular?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700036861

thread_type:
official_review

<thread_object>
creator: reviewer_w6ma

<content>
summary:
The paper proposes a practical method "Inheritune" to create more efficient language models by reducing the layers of the original model. The authors analyze the phenomenon of "attention degeneration" in the deep layers of transformer language models and compare the proposed method against relevant baselines.

strengths:
1. The analysis of attention degeneration phenomenon is well written and easy to understand. The motivation of this work is clear.
2. The paper proposes a practical and lightweight approach to perform model pruning that can be implemented with ease.

weaknesses:
1. The scope of the degeneration analysis and experiments is limited. Only GPT-2 model variants have been discussed and evaluated. The generalization on other model architectures is unclear.
2. The choice of the initialization point n = k/2 assumes degeneration only takes place in the second half of the total number of layers, which is not supported by theoretical evidence.

questions:
Please refer to Weaknesses section.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730700064017

thread_type:
official_review

<thread_object>
creator: reviewer_zpxz

<content>
summary:
The paper introduces Inheritune, a training method designed to create smaller, high-performing language models. This method works by inheriting the initial layers from a larger pre-trained model and gradually expanding the smaller model until its performance matches or exceeds that of the larger model.

The authors investigate attention degeneration in standard large language models (LLMs) and find that rank-collapsed attention matrices often reduce to single-column structures, highlighting inefficiencies in the attention mechanism, especially in deeper layers. This insight into the limitations of deep attention inspired the development of Inheritune.

Experiments on different sizes of GPT-2 models, using the OpenWebText and FineWeb_Edu datasets, show that models trained with Inheritune, despite having significantly fewer layers, can match or even surpass the performance of larger models. They also outperform several baseline methods, such as stacking and knowledge distillation.

strengths:
(+) The proposed training technique is simple and clear.
(+) The authors empirically investigate attention degeneration in standard LLM settings, focusing on rank collapse in attention matrices within the deeper layers of models, such as GPT-2.

weaknesses:
(-) An analysis of whether attention degradation is resolved in the target model trained with Inheritune is missing and should be included.

(-) The adequacy of training for the reference model in Table 1 is uncertain. Given the hyperparameter setting, the GPT2-large (770M) model was trained with only 1.6 billion tokens, using a batch size of 16K tokens and a total of 100K steps. This training setup appears insufficient for comprehensive model training. Thus, it is crucial to compare the performance of the target model against a reference model that has undergone adequate training. 

(-) Since the Inheritune target model begins with the weights of the pretrained model, it has a much lower initial loss at step 0, giving it a head start compared to models trained from scratch. As a result, it’s difficult to determine whether models using Inheritune truly lead to better generalization and convergence. 

(-) The baselines used for comparison need to be updated to reflect advancements in training techniques. Some of the latest research on half-width models, stacking, and distillation that would serve as recommended baselines are as follows:
•	Xiaoqi Jiao et el., TinyBERT: Distilling BERT for Natural Language Understanding, 2020
•	Sheng Shen et al., Staged Training for Transformer Language Models, 2022
•	Peihao Wang et el., LEARNING TO GROW PRETRAINED MODELS FOR EFFICIENT TRANSFORMER TRAINING, 2023

(-) There seems to be a typo: line 323, "layers 0-17" should be "layers 9-17."

questions:
1. Why are the blanks for downstream task performance in Table 1 not filled in?
2. According to the training details of the GPT-2 models (Supplement C.1), it appears that GPT-2 medium was trained on more data than GPT-2 xLarge. Specifically, GPT-2 medium was trained with a batch size of 50K tokens for 100K steps (totaling 5 billion tokens), while GPT-2 xLarge used a batch size of only 16K tokens for the same number of steps (totaling 1.6 billion tokens). Is this a typo?
3. Have you considered initializing the layers added during the growth phase using methods such as random initialization or copying previous layers of the target model? This is because the layers added during the growth phase closely resemble the "lazy layers" described in the paper.
4. It would be helpful to provide a clear criterion for identifying a specific layer as a lazy layer (e.g., based on the rank, mass, or other properties of the attention matrix).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732778322273

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
We sincerely thank all the reviewers for taking the time to review our work. Below, we summarize the major issues raised by the reviewers and provide our general responses to address them.

We have added a new supplementary section  Section D where we have registered new results and discussions done during the rebuttal phase. Additionally Figure 10 which is not a part of Section D is also added during rebuttal. 

> An analysis of whether attention degradation is resolved in the target model trained with Inheritune is missing and should be included.

The degradation is definitely mitigated compared to the full sized reference model. Please refer Figure 9 and Fig 10 of the updated draft. Also posting the figures here (https://postimg.cc/bGKwCcXc) and (https://postimg.cc/cKKJB1bB).

> The attention degradation is observed in GPT2 models which may be undertrained?

We conducted mass analysis (same way as discussed in Sec 2) for OpenLLaMA 3B (26 layers), 7B (32 layers) and 13B (40 layers) models. Refer Fig 16 in updated draft. Posting the Figure here as well (https://postimg.cc/0bbpy2Q6). We observe that 90% of the total mass of the attention matrices are concentrated in fewer columns i.e. in less than 5% of the total columns (i.e 100 for a 100x100 attention matrix) for 3B and 7B models. Next many of attention matrices in later layers are single column for the 13B model.

> Lazy layers are similar to Attention Sinks.

The term "attention sink" [1] refers to the phenomenon where the first token in a sequence receives disproportionately high attention scores compared to other tokens in the attention maps. While there is some connection, as we have also observed that many attention matrices are not only rank-1 but also single-column (with all attention scores concentrated on the first token), this connection has not been explicitly established in [1] with respect to rank-1 behavior or poor training of later layers.

In contrast, as illustrated in Figure 1 of our paper, we compute the maximum rank (plotted on the y-axis) of all attention matrices within a layer. We define a layer as "lazy" only if the maximum rank of its attention matrices is 1, as discussed in Section 2. For instance, consider a layer where only 2 out of 5 attention heads exhibit attention sink behavior. This does not make the layer lazy, as attention is computed as a concatenation of activations across all heads. A truly lazy layer, however, has all 5 out of 5 attention heads fully degenerated, with their attention matrices being rank-1. We provide evidence that such lazy layers are indicative of poorly trained layers.

Attention sink is a general phenomenon and it occurs in attention matrices of lazy and non-lazy layers (refer current draft Figure 17) (also posted here: https://postimg.cc/FY7CPBF5). A lazy layer has all the attention matrices as attention sinks. However attention matrices of non-lazy layers may or may not have attention sinks. Pruning heads i.e. structural pruning is out of scope of this work and we defer this as a future work.


> Limited Model diversity in experiments.

We have demonstrated that Inheritune is effective across various scales of GPT-2 models, including GPT-2 Medium (355M), GPT-2 Large (770M), and GPT-2 XL (1.5B), all trained for 100K steps on the OpenWebText dataset. Additionally, we trained GPT-2 Medium* and GPT-2 Large models on the fineweb_edu dataset to further validate our approach. Moreover, we showed the efficacy of Inheritune on the OpenLLaMA-1.5B model variant, which was trained on 1B tokens. These results establish the broad applicability of Inheritune across diverse model architectures and datasets. Unfortunately, due to constraints in our academic compute budget, we were unable to train the models for longer durations.

> Does the lazy layer hypothesis holds when the analysis is performed using a different data?

We performed additional experiments. In Fig 1 we have utilized vanilla GPT2 models trained with OpenwebText (OWT) and analysis was conducted using val set of OWT (refer Section 2). Next we took the same models trained with OWT but now we perform rank and mass analysis using **new data Fineweb_edu**( process remains the same). In New Figure (https://postimg.cc/4nPdR3yc) one can clearly see lazy layers in vanilla GPT2 medium , large and Xlarge models. Our analysis is quite robust. We will update the draft accordingly.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733165731796

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Discussions during Rebuttal

comment:
We thank the reviewers for their feedback. 

- **Reviewers Bigz and FevH** (2/4) actively participated in the rebuttal discussions and subsequently increased their evaluation scores. Both reviewers find this work promising and impactful.

- **Reviewer W6ma** acknowledged the strong motivation behind this work. Concerns regarding the lack of model diversity (W1) were addressed during the rebuttal process.

### Key Updates

1. **Extended Analysis to Larger Models**  
   We expanded the attention degeneration analysis to larger models, including OpenLLaMA 3B, 7B, and 13B, as shown in Figure 16. (Discussion with Reviewer Bigz)

2. **Robustness to Data Proven**  
   Through additional analysis, we demonstrated that our attention degeneration analysis is robust across datasets. Refer to Figure 18 ([link](https://postimg.cc/4nPdR3yc)), which will be updated in the final draft. (Discussion with Reviewer FevH)

3. **Mitigation of Attention Degeneration via Inheritune**  
   We provided clear evidence that Inheritune effectively mitigates attention degeneration. See Figures 9 and 10. (Discussion with Reviewer Zpxz)

4. **Correlation with Generalization**  
   We established that layers with higher average maximum ranks of attention matrices correlate with better generalization of target model. Refer to Figure 15.  (Discussion with Reviewer FevH)

Overall, none of the reviewers have concerns about the novelty of this method, and the majority of reviewers (3/4) had their concerns about the generality of the attention degeneration analysis and its mitigation through Inheritune addressed during the rebuttal.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes a novel training paradigm aimed at improving the efficiency of language models by reducing their size while maintaining or enhancing their performance. The key idea is to use a technique termed "Inheritune," which involves selectively fine-tuning smaller models to inherit attention patterns from larger pre-trained models. This approach leads to smaller models that are more computationally efficient and attentive, achieving competitive results on multiple benchmarks. Experimental results are provided to demonstrate the effectiveness of the method, with comparisons to baseline techniques such as knowledge distillation and other model compression strategies.

The paper introduces "Inheritune," which is an innovative approach to training smaller language models by transferring attention patterns from larger models. This demonstrates creativity and relevance to the field. It tries to address a critical problem in NLP, reducing the computational cost of large-scale language models. However, the paper lacks a strong theoretical foundation to explain why "Inheritune" works better than existing methods like knowledge distillation. This limits confidence in its broader applicability. The experimental evaluation is not comprehensive enough, as it does not include comparisons with all relevant state-of-the-art baselines or tests on sufficiently diverse datasets. With these concerns, I recommend rejection for this submission. I hope the authors can improve the paper and submit it to a future venue.