When Attention Collapses: How Degenerate Layers in LLMs Enable Smaller, Stronger Models
Authors: Sunny Sanyal
Ravid Shwartz-Ziv
Alex Dimakis
sujay sanghavi
Venue: Accepted by TMLR ()

Abstract: Large Language Models (LLMs) are known for their performance, but we uncover a significant structural inefficiency: a phenomenon we term attention collapse. In many pre-trained decoder-style LLMs, the attention matrices in deeper layers degenerate, collapsing to near rank-one structures. These underutilized layers, which we call lazy layers, are redundant and impair model efficiency. To address this, we introduce Inheritune, a simple yet powerful training recipe designed to build smaller, stronger language models. Inheritune initializes a compact model by inheriting the potent early layers from a larger pre-trained model and then progressively trains and expands it. Our experiments on various models, including the GPT-2 family, demonstrate that models trained with Inheritune can match or even surpass the performance of their larger counterparts, despite having significantly fewer layers. This work presents a novel path toward model compression by design, enabling the creation of compact, yet highly performant language models.

Published in Transactions on Machine Learning Research (02/2026)
When Attention Collapses: How Degenerate Layers in LLMs Enable Smaller, Stronger Models
Sunny Sanyal*
The University of Texas at Austin
Ravid Shwartz-Ziv
New York University
Alexandros G. Dimakis
UC Berkeley
Sujay Sanghavi
The University of Texas at Austin
Reviewed on OpenReview: https: // // operreveew..ee.... net/ forum? id- 22Qn0bUoPf
Abstract
Large Language Models (LLMs) are known for their performance, but we uncover significant structural inefficiency: phenomenon we term attention collapse. In many pre-trained decoder-style LLMs, the attention matrices in deeper layers degenerate, collapsing to near rank-one structures. These underutilized layers, which we call lazy layers, are redundant and impair model efficiency. To address this, we introduce Inheritune, a simple yet powerful training recipe designed build smaller, stronger language models. Inheritune initializes compact model by inheriting the potent early layers from a larger pre-trained model and then progressively trains and expands it. Our experiments the GPT-2 family, demonstrate that models trained with Inheritune can match or even surpass the performance of their larger counterparts, despite having significantly fewer layers. This work presents a novel path toward model compression by design, enabling the creation of compact, yet highly performant language models. Code is available at https:///ithhu. com conm//anyallmmyll//nnn Inheritune
various models, including
1 Introduction
Large Language Models (LLMs) are composed of stacks of decoder-style transformer blocks Vaswani el. al. 2017 As the model grows in size, the model capacity and performance typically improve (Kaplan al.
2020:
Hoffmann et al. 2022 A substantial fraction of the total parameters is devoted towards adding more transformer blocks to increase the depth. Each block or layer in the stack refines the representations learned by the previous blocks, allowing the model to develop a nuanced understanding of the input data.
A transformer block primarily consists of a self-attention module and a feed-forward network (FFN, also referred to as an MLP). Among these, the causal self-attention mechanism (hereafter referred as attention) weighted
is arguably the most critical component. It. enables the model to combine token embeddings lincar sum of attention scores, effectively capturing long-range dependencics and contextual relationships within text. However, as models become deeper, they often exhibit phenomenon known as allention degeneration, characterized by collapse in the rank of the attention matrices. While prior studies have analyzed rank collapse in simplified transformer settings (Dong et al. 2021 Noci et al. 2022 He et al. 2023) this phenomenon has not, to our knowledge, been systematically explored in standard decoder-only LLMs.
Corresponding author: sanyal.smmy@@uteeas...m

---

Published in Transactions on Machine Learning Research (02/2020)
30
22
Bank
Max Bank
24
Layes
Layers
Layers
(a) GPT-2 Medium
(b) GPT-2 Large
(c) GPT-2 XLarge
Figure 1: In decoder-style LLMs, attention matrices in deeper layers often degenerate to near rank-1, limiting their ability to learn meaningful representations. We compute MaxRank (averaged over N 100 randomly selected sequences each with T = 100 tokens) for each layer using the OpenWebText validation set. Our rank analysis of 24-layer GPT-2 Medium, 36-layer GPT-2 Large, and 48-layer GPT-2 XTarge models reveals that attention matrices in many deeper layers collapse to near rank-1.
formal discussion of attention degeneration is provided in Section 2 In this paper, we conduct a detailed empirical analysis of attention degeneration in the GPT-2 family of LLMs (Radford et al. 2019), including GPT-2 Medium (355M), GPT-2 Large (770M), and GPT-2 XLarge (1.5B). Our analysis reveals that many deeper layers these models exhibit prodominantly rank-1 attention matrices across most attention heads within a layer. This suggests that the attention mechanism loses its discriminative ability among tokens and instead performs near-uniform averaging across the sequence. We refer to layers in which all attention matrices degenerate to near rank-1 as lazy layers (a more formal definition is provided in Definition G.1 In the supplementary material, we further extend this analysis to billion-sized LLMs, including LLaMA-3 8B (refer Figure 11), Falcon-7B (refer Figure 21 OLMo- 1o,M Cerebras GPT 2.7B and LLaMA-3 3B (refer Figure 22 to highlight that attention collapse persists even in several modern architectures.
Motivated by the new finding through our novel analysis we aim to develop performant small base language models (LMs) utilizing weights from inefficient larger base LMs without losing pre-train performance (measured by trainn/vlliattin loss). A base LM is decoder-style model trained solely for next-token prediction without additional enhancements like ion luning or reinforcement learning with human feedback (RLHF). Our proposal is straightforward, we start by initializing our smaller LM (target) using the first few blocks from a large pre-trained LM (reference). We then train the target model for a specified number of steps. Afler this initial training, we incrementally grow the target model by adding more blocks, continuing the training process until it matches or surpasses the pre-train validation loss (also val loss) of the reference model. During the growth phase, the newly added blocks can be initialized with leay layers of the reference LM. We refer to this simple yet effective training approach as Inheritune.
In summary, our key contributions are as follows:
1. Novel analysis of attention degeneration in standard decoder LLMs. We empirically investigate atlention degeneration in standard decoder style LLMs. Our analysis shows that rank- collapse in attention matrices, revealing a significant structural inefficiency in the attention mechanism of standard LLMs in deeper layers (refer Figure 1 This degeneration gives rise to what we refer as lazy layers.
2. Introduction of our training recipe Inheritune. Building our analysis we observe that deep LMS often fail to fully utilize their effective depth. To address this inefficiency, we propose Inheritune- a simple yet effective training recipe for developing smaller LMs without losing pre- training performance. This method involves inheriting few early blocks from a much larger reference model and progressively growing and training the smaller model. We validate the effectiveness of Inheritune through extensive experiments on GPT-2 XLarge (1.5B), GPT-2 Large (770M), and GPT-2 Medium (355M) models, trained primarily on the OpenWebText dataset and additionally on Web data.
2

---

Published in Transactions on Machine Learning Research (02/2026)
3. Evaluation against multiple baselines. Models trained using Inheritune consistently outperform wide range of baselines, including much larger models trained from scratch (see Figure and Figure 7 as well as same sized models trained from scratch for twice as many steps (extended training; see Figure Gp. We further compare against warm-started basclines (models initialized with pre-trained weights rather than random initialization
(Ash Adams 2020)) as shown Table
2 Attention Collapse and the Emergence of Lazy Layers in LLMs
Preliminaries: A vanilla transformer-based model consists of I transformer blocks (layers). The model operates input sequence X R2' where T denotes the sequence length (number of tokens), and represents the embedding dimension or model hidden size. The output of cach layer I is denoted as x E R2'xe_ Each transformer block primarily consists of two modules: a self-attention block and a feod-forward network (FFN). The self-attention mechanism enables the model to weight the relevance of different tokens in the sequence relative to each other. Specifically, for a single attention head, the attention computation is defined as equation
[KI VTx
Attention Q, K,) softmax
V
(1)
Attention matrix: A(X)
where the queries Q = XWQ. keys K = XWK, and values V XWv are linear transformations of the input X. Here, Wq, Wx E IR dddx and Wy E IRddd are the weight matrices for the queries, keys, and values, respectively. Typically, de = du = To where is the number of allention heads. In this single-head scenario, we set dN d d.
The attention matrix A(X) E TTx captures the pairwise attention scores between all token positions in the sequence. The softmax is applied row-wise. The attention matrix A(X) is then used to compute weighted sum of the value vectors. Attention rank collapse refers to the phenomenon where the attention matrices A(X) of individual heads in many layers of transformer-based language models lose their expressive capacity, converging towards lower elfective rank structures. Specifically, the effective rank of attention matrices significantly reduces, often approaching rank-1, limiting the model's ability to meaningfully differentiate between token interactions across positions in the sequence. Previous research by Dong et al. (2021) and He et al. (2023) has shown that in self-attention networks (SANs) without residual connections and feed-forward networks (FFNs), the rank of altention matrix converges to rank-1 doubly exponentially with respect to the depth of the model. This phenomenon, known as rank collapse of attention matrices, results in loss of expressive power the attention mechanism attends to all tokens uniformnly. Noci et al. (2022) showed that even with residual connections (without layernorm) attention matrices can still lose rank in deeper layers if the residual connections are not scaled by 1/VT. Interestingly they also linked the rank collapse to vanishing gradients of the keys and queries in deeper layers which affects the overall trainability of the transformer based models. However, these findings do not directly apply to the standard LLMs, as transformer blocks in these models include residual connections, layernorms and FFNs, which are expected mitigate both rank collapse and the vanishing gradient problem.
Approximate Rank Computation of Attention Matrices To assess the presence and severity of rank collapse within standard decoder style transformer architectures (e.g., GPT-2, LLaMA etc.), we utilize singular value decomposition (SVD) for each attention matrix A(X)= = UEVT, where Σ is a diagonal matrix containing singular values σ1 2 02 2 ... 2 or 2 0. The approximate rank (referred to as rank hereafter) of an attention matrix, paramcterized by a variance threshold r, is formally computed as:
min k kE (1,2 .T)
Li=1
27
Σ o3
where 7 E (0. 1) represents the proportion of variance that must be captured by the top le singular values. A lower value of k* indicates ronger rank collapse. In this work, we set = 0.90.
3

---

Published in Transactions on Machine Learning Research (02/2026)
In Figure D we present the layer-wist analysis of rank of GPT-2 models. For this analysis, we computed A(X) using N = 100 sequences selected at random from the validation set of OpenWebText with 4M tokens, each with a sequence length of T' = 100 tokens across all attention heads within each layer. We then define the average approximate rank for each head and layer as Rank(h.)) = N E Kn,,.- Subscquently, we aggregate this metric per layer by taking the maximum rank across heads: MaxRank = max, (Ra As demonstrated in Figure [ MaxRank reveals that many deeper layers exhibit attention matrices that are predominantly near rank-1. We highlight that this rank collapse occurs in GPT-2 Medium, Large, and XLarge models, which are widely used modern LLMs thereby extending the limited findings of Dong et al. (2021] and Noci et al. 2022 We further visualize attention matrices from representative potent and lazy layers of GPT-2 XLarge in Figure 12 Overall, the degeneration of attention matrices in deeper layers provides quantitative evidence for the existence of lazy layers. Specifically, layers exhibit a near-complete rank collapse of attention matrices across all heads, suggesting potentially reduced representational capacity and less effective token mixing in these layers.
observe that some deeper
We provide an extended discussion of attention collapse in Appendix B We analyze five modern LLMs for attention collapse (also refer Appendix G We further assess the robustness of our findings by conducting collapse analyses on different datasets and by using alternative metric, namely the mass of attention matrices. Finally, we perform an ablation over range of values of t (refer Appendix 4.3 to study the sensitivity of the rank-based analysis with respect to
2.1 The Functional Ineffectiveness of Lazy Layers
Having identified lazy layers, we investigate their practical utility: Do these strucburally degenerated layers retain transferable knowledge, or are they functionally impaired? Our experiments suggest the latter.
Max Rank -
3.6
35
Rand int. (lyeerm))
Intt first four byerotayers--4)
middle leyers/leyene-t))
four myersibyersu.
Validation LoSs
Max Rank
3.4
3.3
5
8
10
12
20
40
60
80
100
# Layers
Steps (in
(a) Rank analysis of GPT-2 Small (12 layer).
(b) Performance of layer GPT-2 small variants.
Figure 2: Higher-rank (potent) layers transfer better. (Left, ap Layer-wise MaxRank ( of a pre-trained 12I. GPT-2 Small. (Right, b) Validation loss of 4 layer variants initialized with potent layers (AvgRank & 8.4 9.5) vs. lazy layers (AvgRank ss 1.2) and random weights, after training for 100K steps. Models initialized with lazy layers mirrors the model with random initialization. Training curves are smoothed for visual clarity.
In the first set of experiments, we trained a vanilla GPT-2 small (125M) model with 12 layers for 100K steps on the OpenWebText dataset. We then performed the rank analysis described earlier, with results presented in Figure 2 Specifically,
aggregated the approximate ranks over groups of contiguous layers using AvgRank- = Σt= MaxRank where L is the number of layers in each group. Subsequently, we trained three GPT-2 small v for 100K steps, each initialized with a different contiguous block of four layers from the trained vanilla GPT-2 small model: (a) layers 1-4, with AvgRank = 8.40; (b) layers 5-8, with AvgRank = 9.48; and (c) layers 12, with AvgRank- = 1.22. The last model is initialized with lazy layers. For comparison, we also trained another GPT-2 small variant with random initialization for 100K stops. All
A variant shares the same configurations the reference model but has fewer layers.
4

---

Published in Transactions on Machine Learning Research (02/2026)
Ubyers
Hit [byers- dayers-18)
5.0
Rand Init 12) first feui Isyyrs
5.0
4.5
Validation Loss
4.5
Valldaton
4.0
4.0
3.5
3.5
3.0
3.0
Steps (in
Steps KI
(a) GPT-2 Medium (12 layer) variants.
(b) GPT-2 Large (18-layer) variants.
Figure 3: Initializing 12-layer and 18-layer variants of GPT-2 Medium and GPT-2 Large with deeper lazy layers exhibiting degenerated attention results in performance comparable to random initialization. In contrast, initializing with carly (high-rank) potent layers leads to substantially better convergence and generalization. Training curves are smoothed for visual clarity.
models were trained on the TpenWebTTxtt dataset. As shown in Figure 2 the model initialized with layers from the vanilla GPT-2 small model having higher AvgRank demonstrated the best performance (i.e., lowest final validation loss). Additionally, we observed that the model initialized with lezy layers performed very similarly to the model with random initialization suggesting that laay layers contain minimal transferable knowledge. The results are also summarized in Table 4
For the second set of experiments we utilized larger models namely GPT-2 Medium and GPT-2 Large both similarly trained for 100K steps using OpenWebText. Here we initialized a 12-layer GPT-2 Medium variant and an 18-layer variant of GPT-2 Large using lazy layers extracted from pre-trained 24-layer GPT-2 Medium and 36-layer GPT-2 Large models. We then trained these GPT-2 variants on the same dataset for 10K steps. For comparison, we conducted two bascline experiments where the GPT-2 variants were initialized cither with the first half of transformer layers (potent layers with high AvgRank) and with random initialization. As shown in Figure 3 models initialized with lozy layers demonstrate poor transferability, performing similarly to models with random initialization. This provides additional evidence that lazy layers with fully degenerated attention, fails to learn meaningful representations.
Theoretical Analysis. Finally, we analyze the implications of atlention rank collapse model training from theoretical perspective (refer Section H Our key insight is that rank-collapsed attention head(s) impede learning by inducing vanishing gradients, effectively suppressing updates to the associated Wq and Wk.
3 Inheritune: Our Proposed Training Recipe
This section provides a detailed description of our method, key implementation considerations, and how it addresses the inefficiencies present in current architectures.
As previously established, we have identified the problem of attention degencration and its connection lazy layers, highlighting specific inefficiencies in pre-trained LLMs. In this work, we transform this challenge into an opportunity to create smaller base language models, which we refer to as target models MMga that achieve comparable performance with similar or lower validation loss compared to their larger, loss efficient. counterparts, which we term reference models Mref-
Our proposed solution builds on two key insights: (1) the early layers of deep LLMs contain a higher concentration of potent layers with high AvgRank values, making them suitable for effective model initialization, and (2) lazy lagers can be identified, removed, or utilized in smaller numbers, then subsequently re-trained improve overall model capacity.
5

---

Published in Transactions on Machine Learning Research (02/2026)
Algorithm Inheritune: Training Recipe for Small Language Models
Require:
Reference model Myef with L layers, datasets Purain and Vaal, steps T 1: Copy embedding layer and LM head from Mref to Mtgt
2: Select early contiguous layers from Myet with high AvgRank
Initialize Migt with selected layers between embeddings and LM head
Train Megt Pusin for T steps
while Mtgt performance < Mref performance val do Grow Megt by inheriting additional layers
7: Train Megt for T steps
8: end while
9: return Optimized model Megs
L=24
L=12
L=14
Select the first out of L layers
L=23
layers
Lazy
L=11
L=13
L=15
L=22
L=12
L=14
Number of layers
L-3
L=3
L=4
L=2
L=2
L=2
Ln3
Potent layers
L=1
L=1
L=1
L=2
Round
Round 2 Myy
Round M
XeW NUBI
Myef
M
Initialize with |=12-
Initialize I=14
Initialize with 1=16
Figure 4: Overview of the Inheritune training recipe using 24-Layer GPT-2 Medium model example. A smaller target model is initialized using early layers from a larger, pre-trained reference model. The target model goes multiple rounds of training while inheriting contiguous layers until it matches/out.peroone the reference model. The intensity of the red color in layers correlates with MaxRank
We split the dataset into a training set Durain and a validation subset D Next, we assume that there exists pre-trained reference model Mret comprising L. layers, represented by Wref = (Wa, W1,..., Wr-1) trained with Durain for T steps. We want to train smaller model MMes with the same or better validation loss (lower is beller) compared to its larger counterpart M
Setup:
We now present Inheritune, our proposed training recipe for efficiently developing small base language models (LMs). Inheritune operates on the principle of zero-shot initialization and progressive growth. The Inheritune process consists of three main steps, which we present below and formalize Algorithm
Inherit: Initialize Megt with the first out of L layers of Mret, including prediction head, and token embeddings.
Train: Train Migt for T steps Dirain and evaluate on D
Grow: needed, increase Meg size by adding the next few contiguous layers and repeat steps 1-2 until desired performance achieved.

---

Published in Transactions on Machine Learning Research (02/2026)
3.25
Bryaennnnn uadd Dut (Egarrraa
3.75
Netinall Sayner OLB 1ay
34
Hal degit (ayer-Q4)
3.50
3.2
Velsation Lewe
3.23
3.0
5.25
2.8
2.75
2.6 Steps
Staps
(a) GPT-2 Medium
(b) GPT-2 Large
(c) GPT-2 XLarge
Figure 5: Models derived using Inheritune converge faster and match the final validation loss of the full-sized model, despite having much fewer layers. Comparison of Inheritune trained models (24-layer GPT-2 XLarge variant, 18ylryer GPT-2 Large variant, 16-layer GPT-2 Medium variant) against their full-sized counterparts and same sized variants trained from scratch. All models are trained for 100K steps using OpenWebTest data.
With method now formally described, we turn to empirical validation. In the following sections, we present comprehensive results demonstrating Inheritune's effectiveness across various scenarios, including different model sizes and data regimes. In addition, we conducted an in-depth ablation study to analyze the impact of initialization on performance, providing insights into the adaptability of our approach.
4
Experiments
We evaluate Inheritune through a comprehensive set of experiments using several GPT-2 models: a 48-layer GPT-2 XLarge (1.5B), a 36-layer GPT-2 Large (770M), a 32-layer GPT-2 Large* (668M), and a 24-layer GPT- Medium (355M) (Radford et al. used in our experiments.
2019). Table provides detailed specifications of all model configurations
We use two training datasets: OpenWebText Gokaslan & Cohen 2019 with 10B tokens and Fen (education subset) (Penedo et al.] 2024) also with 10B tokens. Our experimental setup closely follows prior work (Liu et al. 2023] Sanyal et al.) 2024).. For models trained on OpenWebText,, we report validation loss (log perplexity), while for models trained on FineWeb, we report training loss (also log perplexity).
Additionally, for models trained on Fib, we conduct zero-shot downstream evaluations using the (Gao et al., 2024) across five standard benchmarks: ARC-Easy (ARCE; Clark let al.) 2018), LAMBADA (Paperno et al. 2016), SciQ (Welbl et al. 2017), HellaSwag (Zellers et al. 2019 and PIQA (Bsk et al. 2020). Finally, we perform detailed ablation study by initializing each submodule within a transformer block in isolation and training it for 100K steps to identify which component contributes most to performance. All training curves are smoothed for visual clarity.
We provide experimental details of our proposed training recipe Inheritune using a GPT-2 Medium model as an example; similar procedure was applied to train other models. A visualization of the training recipe is presented Figure 4 Our recipe for applying Inheritune involves the following steps.
We train a vanilla 24-layer GPT-2 Medium model (reference model) on Duraiu for 100K steps and evaluate its validation loss ( log perplexity) D This establishes benchmark validation loss.
Reference Model:
2. Model initialization: We initialize an 12-layer model (1 -t= L/2) using the reference model.
Training and Evaluation: We train the 12-layer model on Duran for T=100K steps and evaluate its validation loss.
7

---

Published in Transactions on Machine Learning Research (02/2026)
H-ll ligie 1o0
W4.000 Un00
Velcetoon
Jalcusaonn 54
speedup
2x speedup
2x s speedup dup
Steps
(a) GPT-2 Medium
(b) GPTT2 Large
(c) GPT-2 XLarge
Figure 6: Models trained with Inheritune match the validation loss of same sized models trained from scratch for twice as many steps. We compare Inheritune-trained models (24-layer GPT-2 XLarge, 18-leyer GPT-2 Large, and 16-layer GPT-2 Medium) against their same-sizee counterparts trained from scratch for twice the number of steps. Inheritune models are trained for 100K steps, while the bascline models (trained from scratch) are trained for 200K steps, on OpenWebText data.
4.00
Rull model (byers-32)
Fall model (byyers-24
Partial depth (layers-15)
partial depth (layers16))
(ayers:16)
3.75
3.8
Trining Loss
Taining Less
3.50
3.00
0
20
40
60
80
100
20
40
60
100
Steps
Steps
(a) GPT-2 Large'
(b) GPT-2 Medium
Figure 7: Models derived using Inheritune without data repetition converge faster and match the final validation loss of the full-sized model despite using lesser layers. Additionally, the model trained using Inheritune demonstrates data efficiency, achieving a lower validation loss in fewer steps compared to its full-sized and half-sized counterparts.
4. Iterative Refinement: If the smaller model's performance is inferior to the reference model, then incrementally increase its size by adding additional layers and repeat steps 2-3 until we achieve parity with the reference model's validation loss.
We choose 1 = L/2 as the starting point and increase the model size by two layers in each round across all our experiments, aiming to minimize the number of training rounds. In principle, Inheritune should generalize to other hyperparameter choices as well.
Baseline-I, We compare GPT-2 model variants (i.e., models with fewer layers than their vanilla configura- tions) trained using Inheritune against the following baselines:
Larger reference models with more layers, trained from scratch (random initialization).
2. Same-sized models with the same number of layers, trained from scratch.
Extended training baselines, same-sized models trained from scratch for twice as many steps compared to models trained with Inheritune.
8

---

Published in Transactions on Machine Learning Research (02/2026)
Models
Recipe
Layers
ARCE
PIQA
SciQ
Hellaswag
Lambada
Average
rand init
24
51.05
61.81
74.8
30.79
20.28
47.74
GPT-2 Medium
rand init
16
49.92
61.92
73.3
29.56
19.54
46.84
Ours
51.26
61.81
73.8
30.55
23
48.08
rand init
32
52.48
64.58
75.3
32.65
22.2
49.44
GPT-2 Large
rand init
16
50.34
63.11
75
30.86
21.56
48. 17
Ours
16
52.9
63.55
76.1
32.14
24.06
49.75
Table 1: Models trained with Inheritune achieve comparable average zero-shot downstream performance to their larger reference models and surpass same-sized counterparts trained from scratch. Downstream evaluations are performed on models pre-trained with the FineWeb dataset. (see Figure 7. Performance is measured using accuracy (acc) and normalized accuracy (acc-norm) metrics, following the Open LLM Leaderboard protocol Becching ct al. (2023). We have highlightod the best average scores in bold.
Models
Layers
Recipe
Steps
Pre-train Val loss (1)
24
Half-width
100K
3.04
16
Stacking
100K
2.84
GPTT2 Medium
16
Hybrid-stacking
100K
2.83
16
Ours
100K
2.81
36
Half-width
100K
3.06
18
Stacking
100K
2.87
GPT-2 Large
18
Hybrid-stacking
100K
2.89
18
Ours
100K
2.80
48
Half-width
100K
2.77
24
Stacking
100K
2.65
GPT-2 XLarge
24
Hybrid-stacking
100K
2.64
24
Ours
100K
2.64
Table 2: Inheritune outperforms warm-started basclines (Basclinc-II), Comparison of pro-training validation loss for GPT-2 XLarge, GPT-2 Large, and GPT-2 Medium variants trained with Stacking, Hybrid Stacking, Half-Widtha and Inheritune recipes. All baselines warm-started, ie., initialized with pre-trained weights rather than random initialization. Across model scales, Inheritune consistently achieves lower validation loss than the majority of warm-started baselines. The lowest validation loss (lower is better) is highlighted in bold.
Baseline-II. Additionally, we compare Inheritune against warm-started baselines (Ash & Adams) 2020 Warm starting in neural network training refers initializing a model's parameters with weights from previously trained model, rather than starting from random initialization (a "cold start"). This setup enables fair comparison with our approach, which also leverages prior learned representations. The warm-started basclines include stacking (Gong et al. 2019 J. Reddi et al. 2023). hybrid stacking (which initializes model layers directly from a larger reference model), and half-width (which relains all layers but reduces both the hidden dimension and number of attention heads by half, using weights from the reference model). Finally, we briefly compare Inheritune with knowledge distillation (Hinton al.] 2015); results are provided in the supplementary material (see Figure 17). Detailed descriptions of all baselines are presented in Section D of the supplementary material.

---

Published in Transactions on Machine Learning Research (02/2026)
4.1 Results and Discussions
Models trained with Inheritune outperform both larger and same sized models trained from scratch. We present our main results in Figure 5 The 24-layer, 18-layer, and 16-layer variants derived using Inheritune from the vanilla 48-layer GPT-2 XLagge, 36-layer GPT-2 Large, and 24-layer GPT-2 Medium, respectively, achieve comparable or lower validation losses than both their full-sized counterparts and same- sized models trained from scratch, when trained for the same number of steps (100K). Our GPT-2 XLarge and GPT-2 Large variants require single round of Inheritune training, while the GPT-2 Medium variant. undergoes three rounds with 12-, 14-, and 16-layer configurations. Furthermore, as shown in Figure B models trained with Inheritune reach the same validation loss as samc-sized models trained from scratch in approximately half the number of training steps. Moreover, for the GPT-2 Medium and Large variants, Inheritune achieves a strictly lower loss floor that same-sized models fail to reach even when trained for twice as many steps. A tabular summary of these results is provided in Appendix Table 5
We conducted additional training experiments, using a high-quality training data namely Fineweb. We trained a custom 32-layer GPT-2 Larget (668M) and 24-layer GPT-2 Medium (355M) reference model from scratch. Next, we trained two 16-layer variants: one derived from GPT-2 Large" and the other from GPT-2 Medium, using their respective reference models following Inheritune. For comparison, we also trained 16-layer baseline models from scratch. All models were trained for 100K steps, and training loss was used to evaluate pro-training performance. We observe thematically consistent results: as shown in Figure 7 the 16-layer variants trained with Inheritune consistently match the performance of their full-sized counterparts and outperform same-sized baselines, both in terms of training loss and zero-shot downstream evaluation. Downstream results are provided in Table Model configurations and training hyper-parametors are detailed in the supplementary material (refer Section
Models trained with Inheritune outperform all warm started baselines. In Table 2 compare GPT-2 XLarge, GPT-2 Large, and GPT-2 Medium variants trained with Inheritune against same-sized variants trained with stacking, hybrid stacking, and half-width basclines. The half-width bascline performs poorly, revealing the limitations of naive width reduction. While stacking and hybrid stacking demonstrate reasonable performance, they still fall short compared to Inheritune. Across all cases, Inheritune consistently outperforms these basclines, highlighting its effectiveness as an initialization strategy, with a single exception in the GPT-2 XLarge case where it matches one baseline. For a detailed view of the training curves across all methods, refer to the training curves in supplementary Figure 18
4.2 Inheritune Mitigates Attention Collapse
We attribute the success of Inheritune, to its ability to miligate atlention collapse, thereby leading to fewer loay layers after training, In Figures g we juxtapose the attention rank patterns of the vanilla and inteittune-trineed GPT-2 Medium. Notably, none of the GPT-2 Medium variants exhibit lazy layers. A similar analysis is conducted with GPT-2 XLarge (refer supplementary Figure 13
The corresponding attention patterns for GPT-2 Medium, shown in Figure 9 further corroborate our observation. The attention patterns for both a vanilla 24-layer model trained from scratch and a 16-layer model trained using our proposed method, Inheritune. Note just for the sake of better visualization we visualized full attention and not causal attention, in practice GPT-2 models compute causal attention. We computed these attention matrices using randomly selected strings from the validation set of OpenWebText and took 40 tokens averaged over 3 runs. In the 24-layer model trained from scratch (top row of Figure g) we observe clear progression in attention patterns. The early layers (L4 and L7) exhibit dense structured attention patterns. In contrast, the deeper layers (L200 and L22) display more uniform patterns, indicating a loss of focus (allention). This uniformity is a hallmark of lazy layers, where the allention mechanism loses its ability to selectively focus on specific relevant tokens. In contrast, our 16-layer model trained with Inheritune (bottom row) demonstrates more focused and effective attention patterns, even its later layers (I.11 and L15). This striking difference suggests that our method makes the model more attentive and addresses attention collapse, potentially leading to more efficient. models in compact size.
10

---

Published in Transactions on Machine Learning Research (02/2026)
(Lay (Layers-66)
Full Model (Layers-24)
Rank
40
Max Rank
30
20
10
0
10 15
20 25
Layers
Figure 8: Rank collapse in deeper layers and its mitigation through Inheritune. The maximum (max) rank across all attention heads for each layer is plotted, following the methodology in Figure [ Analysis of a 24-layer GPT-2 medium model reveals rank-1 attention matrices in later layers (those beyond the halfway point), indicating rank collapse. Specifically, out of the last 12 later layers exhibit rank-1 attention matrices (mean rank across all the 100 runs). Our 16-layer GPT-2 medium variant, trained with Inheritune, demonstrates improved rank across all layers, highlighting the effectiveness of our approach. Notably, none of the later layers in our 16-layer variant exhibit rank-1 attention matrices.
TA H2
L, H3
1,20 H2
T,22 HM
Vanilla 24 layer GPT-2 Medium
L4 H2
L7 H3
L11 H2
L15 H1
GPT-2 Medium 16 layer variant (Ours)
Figure 9: Inheritune preserves effective attention patterns. Comparison of attention patterns across layers (T.) and heads (H) in two GPT-2 medium models: (top) vanilla 24-layer model trained from scratch, (bottom) 16-layer variant trained with Inheritune. Attention maps are averaged over three randomly selected strings, with 40 tokens each from the validation. Darker colors indicate higher attention scores. Inheritune maintains focused attention even in deeper layers, contrasting with the uniform patterns in the vanilla model's 's deeper layers.
4.3 Ablations and Limitations
We conducted extensive experiments to better understand which sub-module initializations within a trans- former block lead to improved generalization (in terms of validation loss) and faster convergence. For these ablations, we fixed the model 16-layer GPT-2 Medium variant and explored three different sub-module initializations using weights from 24-layer GPT-2 Medium reference model. We initialize the transformer
11

---

Published in Transactions on Machine Learning Research (02/2026)
Layers
Initialization
Steps
Pre-train Val loss (4)
16
Attention
100K
2.84
16
MLP
100K
2.85
16
Attention and MLP"
100K
2.80
16
Ours
100K
2.81
Table 3: Impact of initializing varions sub-modules within a transformer block. We compare validation loss of a 16-layer GPT-2 Medium variant when different sets of sub-modules are initialized with weights from the first 16 layers of a 24-layer GPT-2 Medium reference model. Submodules without the marker also include layernorm weights. All models are trained with the OpenWebText dataset. Key findings:
(1) Inheritune initialization and Attention and MLP initialization result in similar performance improvements;
(2) layernorm initialization shows minimal impact. The training curves of corresponding models are presented The lowest and our corresponding validation losses (lower is better) are highlighted in bold.
in Figure 19
blocks with 1) Attention block (key, query, value, and projection weights) along with the layernorn 2) Attention and MLP (FFN) weights without the layer-norm weights, and 3) MLP block weights along with the layernorm. We note that Inheritune performs initialization by inheriting the entire layer weight i.e. attention, MLP along with the layernorm weights.
As shown in Table 3 models initialized with both Attention and MLP weights achieve the best performance, irrespective of the LayerNorm initialization. A detailed validation loss versus training steps plot is provided in the supplementary Figure 19 These results suggest that jointly initializing the Attention and MLP submodules offers clear advantage over initializing either component alone. Interestingly, we also find that initializing only the Attention or only the MLP weights yields comparable improvements in both convergence speed and final validation loss.
Sensitivity Analysis for attention rank computation and lazy layers. Next, we analyze the sensitivity of the approximate rank to the variance threshold т. Following the rank computation methodology described in Section 2 we perform rank analysis on the full GPT-2 Medium and GPT-2 Large models using a randomly sampled subset of the OpenWebText's validation set, varying т E {0.8, 0.85, 0.9, 0.953. As shown in Figure 10 the MaxRank(I) (y-axis) remains highly stable for the lazy layers across all values of T for both the models.
Limitations of our work. Our analysis primarily focuses on pre-layernorm (Pre-LN) architoctures. By default, we initialize training at 1 = L/2, i.e. from the midpoint of the model where early layers have already developed strong representations to minimize the total number of training rounds. Figure 20 presents ablation study showing training curves across three rounds performed following our proposed method. Notably, by the third round (R = 3), models trained with Inheritune match the validation loss of their reference counterpart (Full model). Despite its effectiveness, Inheritune is computationally expensive, as it may require multiple rounds of training during the growth phase. Lastly, our current analysis focuses exclusively on attention submodules; extending this framework to foedforward (MLP) layers remains an important dircction for future research.
5 Related Works
Attention degeneration has been studied in the past through the lens of attention rank collapse Dong ct al. (2021] leading to representation collapse, and attention entropy collapse Zhai et al. (2023) leading training instability. This also has been studied is a theoretical setup for transformer models by Noci let al. (2022): : Barbero al. (2024); Wu et al. (2024). Recently He et al.) (2023) address rank collapse in self-attention networks (SANs) without residual connections or layer norms, using two different model initialization techniques that enable faithful signal propagation -ie., Σg of 4(x4) does not collapse in deeper layers. However, this approach significantly slows down training. Noci et al.
2022 proposes scaling residual
2In GPT-2 models layernorm blocks parameterized.
12

---

Published in Transactions on Machine Learning Research (02/2026)
LAyeR
#
7=0.80
(b) T=0.85
7=0.90
(d) 7=0.95
(e) 0-0.8
(l) 0085
(g) 009.0
(h) 7=0.55
Figure 10: Lazy layers remain robust under variations of To study the sensitivity of attention rank to the varionce-ratio threshold т E {0.8, 0.85, 0.9, 0.95], we visualize MaxRank(I) as function of layer index for (a d) the GPT-2 Medium full model with 24 layers (top row) and (e h) the GPT-2 Large full model with 36 layers (bottom row).
connections by 1/VL, while Barbero al. (2024) suggest that adding additional tokens to already long sequences of repeated tokens can help mitigate collapse. In contrast to prior works, address attention degeneration by developing smaller models that climinate structural inefficiencies and training these models to match the performance of their larger, inefficient counterparts.
LLM training recipes and model initialization. The stacking method Gong et al. 2019 J. Reddi et al.
2023 employs a stage-wise training strategy that uses weights from initial layers to initialize later layers has been shown to be effective for LLM training both empirically Gong et al. (2019): J. Reddi et al. (2023):
Du et al. (2024) and theoretically Agarwal et al. (2024). Knowledge distillation Hinton et al. (2015) has been very successful in training small LMs. In some cases Ture et al. (2020); Sanh et al. (2019) the smaller student model is also initialized with teacher layers, though this is often done without clear explanation intuition. Recent works in model initialization, such as Trockman & Kolter (2023), have studied synthetic attention patterns for initialization, primarily in vision sottings. However, such methods have limited succoss in language models. Xu et al. (2024) use weight initialization for faster fine-tuning of vision models. In contrast, our proposed recipe focuses on creating smaller model by eliminating specific structural inefficiency in lazy layers. This distinction sets our work apart in terms of both objective and methodology.
6 Conclusion
In this work, we identify a structural inefficiency in deep decoder-style LLMs, which we term attention collapse, where attention matrices in deeper layers often degenerate into near rank-one structures, rendering these layers ineffective. These ineffective layers referred to lazy layers contribute little to the model's representational power. To address this, we introduce Inheritune, multi-stage training recipe that initializes a smaller model using a few potent early layers from a larger pre-trained model and then progressively trains and expands through multiple rounds. Our experiments demonstrate that models trained with Inheritune can match or even surpass the performance of their larger counterparts despite having significantly fewer layers. By mitigating attention collapse, our approach produces compact and highly performant models, offering new path toward designing smaller, more attentive architectures from the ground up.
13

---

Published in Transactions on Machine Learning Research (02/2026)
Acknowledgments
This research was supported by NSF Grant 2217069 and by UT Austin's Center for Generative Al and the Machine Learning Lab. We thank Nived Rajaraman for helpful discussions and the anonymous reviewers for their valuable feedback.
References
Naman Agarwal, Pranjal Awasthi, Satyen Kale, and Eric Zhao. Stacking as accelerated gradient descent. arXiv preprint arXiv: 2403.04978, 2024.
Jordan Ash and Ryan Adams. On warm-starting neural network training. Advances neural information processing systems, 33:3884 3894, 2020.
Federico Barbero, Andrea Banino, Steven Kaplurowski, Dharshan Kumaran, João G. M. Araújo, Alex Vitvitskyi, Razvan Pascamu, and Petar Veličković. Transformers need glasses! information over-squashing in language tasks. arXiv preprint arXiv: 2406.04267, 2024.
Edward Beeching, Clémentine Fourrier, Nathan Habib, Sheon Han, Nathan Lambert, Nazneen Rajani, Omar Sanscviero, Lewis Tunstall, and Thomas Wolf. Open Ilm leaderboard. https::/huggingfacee cotpsccescccscccsee...ee 11m leaderboard, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. In The Thirty-Fouth AAAI Conference Artificial Intelligence,
AAAT 2020, The Thirty-Seoond Innovative Applications of Artificial Intelligence Conference, TAAT 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY, USA, February 7-12, 2020, pp. 7432-7439. AAAI Press, 2020. doi: 10.1609/AAA. V34105.64 URL
https://dol org/ 10. 1609/ 1..69//aai......... aaai. v34i05. 6239
Christopher Clark, Kenton Lee, Ming-Wor Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Jill Burstein, Christy Do- ran, and Thamar Solorio (ods.), Procoedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistios: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pp. 2924-2936 Association for Computational Linguisties, 2019. doi: 10.18653 N0.18853/V11M0 1111/N9911000 URL http::///0ol.orr/1....8////////
h19-1300
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge. ArXiv, URL abs/1803.054577 2018.
Together Computer. Redpajama-data: An open source recipe to reproduce llama training dataset, 2023.
-atta
Yihe Dong, Jean- Jcan-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. International Conference on Machine Learning, 2021.
Wenyu Du, Tongxu Luo, Zihan Qiu, Zeyu Huang, Yikang Shen, Reynold Cheng, Yike Guo, and Jie Fu. Stacking your transformers: A closer look at model growth for efficient 11m pre-training. arXiv preprint arXiv: 15319, 2024.
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPoli, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noachh, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailcy Schoelkopf, Aviya Skowron, Lintang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. framework for few-shot language model evaluation, 07 2024. URT. httts::/zennd:: http:e////aasooooorrrrrrr//////p .rg/reeordss//2660022
Xinyang Geng and Hao Liu. Openllama: An open reproduction of Ilama, May 2023. URL https: comr/ppenIm-reeearcch/ooe 1lama
14

---

Published in Transactions on Machine Learning Research (02/2026)
Aaron Gokaslan and Vanya Cohen. OpenWebTextCorpus 2019.
Openwebtext
corpus.
http:///syyLion077.gitbbb io/
Linyuan Gong, Di He, Zhuohan Li, Tao Qin, Liwei Wang, and Tieyan Liu. Efficient training of BERT by progressively stacking. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceodings of Machine Learning Rescarch, pp. 2337 2346. PMLR, 09 15 Jun 2019. URL https::/prrceeddign//e oot.c/eeiiina.alrrprrrr////////p
mlr www.ppe////s/vvsvl....
html
Bobby Hc, James Martens, Guodong Zhang, Alcksandar Botev, Andrew Brock, Samuel L Smith, and Yee Whye Teh. Deep transformers without shortcuts: Modifying self-attention for faithful signal propagation. In The
Eleventh International Conference Learning Representations, 2023. URL https:///ppeneetee..nee
forumPid=MPPraUQgKjj
Gooffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge a neural network. arXiv preprint
arXiv: 1503.02531, 2015.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15552 2022.
Sashank J. Reddi, Sobhan Miryoosefi, Stefani Karp, Shankar Krishnan, Satyen Kale, Seungyeon Kim, and Sanjiv Kumar. Efficient training of language models using few-shot learning. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Rescarch, pp. 14553-14568 PMLR, 23-29 Jul 2023. URL httts::/proceeddngg... .ml press/v202/j-reddii22. html
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint
Ntl::20010083611 2020.
Hong Liu, Zhiyuan Li, David Hall, Percy Liang, and Tengyu Ma. Sophia: A scalable stochastic second-order optimizer for language model pre-training. arXiv preprint arXiv: 2305.14342 2023.
Xinyin Ma, Gongfan Fang, and Xinchao Wang. Llm-pruner: On the structural pruning of large language models. NEURIPS, 2023.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models, 2016. URL https::/arxiv..rgg wttp://.rrg/bbb///aaooaee abwp/a/......../aaa/ 07843
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27198- 28811,, 2022.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Q. N. Pham, R. Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and R. Fernández. The lambada dataset: Word prediction requiring a broad discourse context. Annual Moeting Of The Association For Computational Linguistics, 2016. doi:
00.18633vv1/116 114441144
Guilherme Penedo, Hynek Kydlíček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data at scale, 2024. URL htttps://arxiv.oo www..org/////.eeeee.eo.ooeeeee/e 2406/
Alec Radford, Jeffrey WU, Rewon Child, David Luan, Dario Amodei, and Tly Sutskever. Language models are unsupervised multitask learners. OpenAl Blog, 1(8):9, 2019.
Colin Raffel, Noam Shazoer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yançi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. arXiv e-prints, 2019.
15

---

Published in Transactions on Machine Learning Research (02/2026)
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In The Thirty-Fourth AAAT Conference on Artificial Intelligence AAAI 2020, The Thirty-Second Innovative Applications of Artificial Intelligence Conference, LAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020, New York, NY. CSA, February 7-12, 2020, pp. 8732 8740. AAAT Press, 2020. doi: 10. 10.1609/ AAA. http:::/ddoooor///// 10. 1609/660//00 1609/aaai......aaaaaaaaaaaaa.aaa aaai. v34105. 6399
URL
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, distilled version of bert: smaller, faster, cheaper and lighter. NEURIPS, 2019.
Sumny Sanyal, Atula Tejaswi Neerkaje, Jean Kaddour, Abhishek Kumar, and sujay sanghavi. Early weight averaging meets high learning rates for LLM pre-training. In First Conference on Language Modeling, 2024. URL https://openreviien attp:://pparriiieee.e//rroomomme http://ww...e.rrmmmmmoooaaaoot.
Asher Trockman and J. Z. Kolter. Mimetic initialization of self-atlention layers. International Conference on Machine Learning, 2023. doi: 10.48550/arXiv 2000000000..
Iulia Ture, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. ICLR, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. NEURIPS, 2017.
Johannes Welbl, Nelson F. Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. ArXiv, 0tee/7777066209 2017.
Xinyi Wu, Amir Ajorlou, Yifei Wang, Stefanie Jegelka, and Ali Jadbabaie. On the role of attention masks and layernorm transformers. Advances in Neural Information Processing Systems, 37:14774-14809, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, 2024.
URL httts::/ wwt.:forunurnnee.. Nd5N7
Zhiqiu Xu, Yanjie Chen, Kirill Vishniakov, Yida Yin, Zhiqiang Shen, Trevor Darrell, Lingjie Liu, and Zhuang Liu. Weight selection for model initialization. In The Twelfth International Conference on Learning Representations, 2024. URL https::/openrreiieet net hett/forrrridddddyrrhhhiii
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: 1905.07830, 2019.
Shuangfei Zhai, Tatiana Likhomanenko, Etai Littwin, Dan Busbridge, Jason Ramapuram, Yizhe Zhang, Jiatao Gu, and Joshua M. Susskind. Stabilizing transformer training by preventing attention entropy collapse. In Andreas Krause, Emma. Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 40770-40803. PMLR, 23-29 Jul 2023. URL mlr. pwwps//ess2e///p 20.000/22haaaa htmml /tt.://add.aaa......rrree///////
16

---

Published in Transactions on Machine Learning Research (02/2026)
Supplementary Materials
Contents
Frequently Asked Questions
Extended Discussion on Attention Collapse
Understanding Attention Degradation using Attention Mass Analysis
Baselines
Supplementary Experiments
Model Configurations and Training Details
G
Additional Experiments and Discussions
Theoretical Analysis
A Frequently Asked Questions
A.1 Is your method not just pruning?
Inheritune is a stage-wise efficient training recipe that addresses structural issue in decoder-style transformer collapse whhcc which we consistently observed across multiple models.
Unlike pruning, Inheritune includes a growth phase where the model is expanded until it outperforms the reference model (refer Algorithm [ and Figure 20).. Pruning doesn't always require re-training (Ma al.) 2023), whereas Inheritune may need multiple rounds of re-training. To the best of our knowledge, method has explicitly studied resolved attention collapse in LLMs. Our method has closer proximity to efficient training recipes employing model initialization
pruning
(J. Reddi et al. 2023 Du el. al. 2024 warm
starting Ash & Adams 2020).
A.2 Is the comparison with baseline models trained from scratch unfair since Inheritune uses weights from pre-trained models?
In Bascline-II (refer Section 4 we have also compared our method with warm started basclines which also uses pre-trained model weights for model initialization for fair comparisons.
For Baseline-I, we include much larger reference models as well as same-sized models trained for twice as many steps. Remarkably, Inheritune still outperforms both. We believe these findings are novel and reveal a new axis for scaling.
A.3 The attention collapse analysis is not holistic?
TL;DR. We analyzed the phenomenon of attention collapse across four datasets and four different model architectures, with model sizes ranging from millions billions of parameters.
We evaluated attention patterns and analyzed the phenomenon of attention collapse using four datasets: OpenWebText, Web, RedPajama, and C4. Our analysis used two complementary metrics namely, approximate rank and approximate mass to quantify the structure of attention matrices. We conducted altention collapse analysis on a range of models, including GPT-2 (Medium 355M, Large 770M, XLarge 1.5B), LLaMA-3 (3B and 8B models ), OLMo 1B, Cerebras-GPT 2.7B, Falcon 7B and LLaMA-1 (OpenLLaMA 3B, 7B, 13B) which features notably different architecture. The evaluation was performed on 10K tokens (100 samples 100 tokens each), and we also provide visualizations of the resulting attention patterns.
17

---

Published in Transactions on Machine Learning Research (02/2026)
A.4 What is the connection between Attention collapse and Attention sinks?
The term Attention sink (Xiao el al. 2024 refers to phenomenon where specilic token in sequence receives disproportionatel high attention scores compared to other tokens in the attention map.
In our analysis of Attention collapse, we also observed sink-like bohavior for certain tokens across all attention maps (see Figure and Figure 12). However, unlike lypical attention sinks, we found that beyond the sink token, no meaningful attention structure remains: all other tokens receive nearly uniform attention scores. We further connect this behavior to the emergence of lazy layers. Therefore our analysis has unique insights compared to attention sinks.
B Extended Discussion Attention Collapse
B.1
Attention Collapse in LLaMA-3 models
We conducted rank analysis contemporary LLaMA-3 base model with 8B parameters. We compute Rank(h.)) 1 En=1 h, where N 100 sequences are sampled from subset of the C4 dataset (Raffel et al.
2019 As shown in Figure 11 we observe that nearly 50% of the atlention heads (500 out of 1024
across all layers) are close to near rank-1, highlighted in red. This presents an interesting case: in very large modern architectures such as LLaMA-3 8B, while there may not be entire lazy layers, a substantial number of heads within many layers exhibit degeneracy. different pattern of attention collapse compared to GPT-2 models can be attributed to the architectural differences between these models.
26
25
24
23
22
21
3 14
15
13
12
11
10
Figure 11: Rank analysis of LLaMA-3 8B reveals that nearly half of the attention heads exhibit rank collapse. We analyze the LLaMA-3 8B model, which contains 32 heads per layer (32 32), using the rank metric defined in Section 2 The results are visualized as a heatmap of head index vs. layer index. Potent (non-collapsed) heads are shown in varying shades of green, where higher intensity indicates higher rank, while rank-collapsed heads (near rank-1) are highlighted in red. Approximately 50% of all allention heads exhibit rank collapse, indicating widespread degeneracy.
B.2 Attention Pattern Visualization of Potent and Lazy Layers
Following the analysis in Section 2 and Figure [ we present the attention patterns two representative layers from the GPT-2 XL model: lazy layer (Layer 30) and a potent layer (Layer 8), as shown Figure 12 The attention patterns of these layers exhibit distinctly different behaviors. particular, the lazy layer demonstrates clear collapse, where attention concentrates almost exclusively the first token. Next,
18

---

Published in Transactions on Machine Learning Research (02/2026)
L8 H16
L8 H17
L8 H18
H19
H20
Potent layer of vanilla GPT-2 XLarge.
L30 H16
L30 H17
L30 H18
L30 H19
L30 H20
A lazy layer of vanilla GPT-2 XLarge.
Figure 12: Visualization of attention patterns in lazy and non-lazy layers of vanilla GPT-2 XLarge model with 48 layers. The top displays attention patterns for various heads (H) in layer (L) 8, while the bottom row shows patterns for layer (L) 30. For visual clarity, we display the full attention maps; however, attention in GPT-2 models is inherently cansal.
following the discussions in Section 4.2 we visualize the rank layer relationship for the GPT-2 XLarge model, juxtaposing a vanilla model with a Inheritune-trained model containing half many layers. Although both models achieve similar validation loss, the Inheritune-trainet model exhibits significantly fewer lazy layers compared to the vanilla counterpart.
Ours (Layers-24)
60
Full Model (Layersmt8)
50
Max Rank
30
20
10
Layers
Figure 13: Rank collapse worsens for larger LLMs, Inheritune helps to mitigate rank collapse. The maximum (max) rank across all attention heads for cach layer is plotted, following the methodology in Figure D We analyze 48-layer GPT-2 XLarge model which reveals rank-1 attention matrices in later layers (those beyond the halfway point), indicating rank collapse. Specifically, 22 out of the last 24 later layers exhibit rank-1 I attention matrices (mean rank across all the 100 runs). Next, Our 24-layer GPT-2 XLarge variant, trained with Inheritune, demonstrates improved rank across many layers, highlighting the effectiveness of our approach. Notably, 2 out of 12 of the later layers in our 24-layer variant. exhibit. rank-1 1 attention matrices.
19

---

Published in Transactions on Machine Learning Research (02/2026)
GP72 240
25
SP12 Langet ayers-3d)
25
Aug coummr with 3000
Pe coMMnn with
15
1.0
(a) CPT-2 Medium
(b) GPTT2 Large
(c) GPT-2 XLarge
Figure 14: In decoder-style LLMs, attention matrices in deeper layers often degenerate to near single column matrices, which is a special case of near rank-1. We compute AvgMass()) (averaged over N = 100 randomly selected sequences each with T= 100 tokens) for each layer using the OpenWebText validation set. Our mass analysis of 24-layer GPT-2 Medium, 36-layer GPT-2 Large, and 48-layer GPT-2 XLarge models (L-layer, H:hidden size) reveals that attention matrices in many deeper layers collapse single column matrices average.
C Understanding Attention Degradation using Attention Mass Analysis
In this paper, we have analyzed the attention degradation phenomenon primarily using a single metric-rank of the attention matrices (see Section 2 In this section, we aim to explore another thematically related metric to further investigate the nature of attention degradation.
We further investigated the dominant structure of the rank-1 atlention matrices and observed that, on an average, many of these matrices have their mass concentrated in a single column. This intrinsic structure can be viewed as a special case of rank-1 attention matrices. To quantify this, we computed the proportion of the matrix mass contributed by each column of A(X) by computing [| 12 where A. denotes the j-th A(X)
column of A(X), || is the fz-norm of that column, and ||A(X)] is the Frobenius norm of A(X).
Next we determine the minimal number of columns required to capture 7) proportion of the total mass, formally computed as;
m Σ
| J=1 [A(X)
m min m E 11, ... T}
2
Here 17 (0.1) 1) represents the cumulative column mass threshold. In this work, we set 1) = 0.90. A lower value of m² implies stronger concentration of the allention matrix mass within fewer columns, reinforcing the phenomenon attention collapse. This analysis provides additional quantitative evidence highlighting the reduced representational capability of attention matrices in deeper transformer layers, further supporting the identification of leay layers.
In Figure 14 we present the layer-wise analysis of the attention matrix mass concentration in GPTT2 models. For this analysis, (similar to the rank analysis), we computed A(X) using N 100 sequences selected at random from the validation set of OpenWebText (4.4M tokens), each with sequence length of T 100 tokens across all attention heads within each layer. We define the average minimal column count required to capture 90% of the attention matrix mass for each head and layer as: (h,1) = Σ ",,,, Subseeuuuull (k'.), We
we aggregate this metric per layer by taking the average across all heads: AvgMass = Σ - observe that many of the rank-collapsed atlention matrices in deeper layers exhibit single-colurnn attention structures, as measured by the AvgMass(() criterion. As shown in Figure 15 we performed a mass analysis on contemporary parammter OpenLLaMA models (Geng Liu 2023) and observed similar pattern of attention degradation in the deeper layers. This provides concrete evidence that the phenomenon persists across broad range of architectures and also at the billion-parameter scale.
20

---

Published in Transactions on Machine Learning Research (02/2026)
15
10
20
10
15
Layers
Layers
Layers
(a) OpenLLaMA-3B
(b) OpenLLaMA-7B
(c) OpenLLaMA 13B
Figure 15: The overall of tokens, tends to concentrate in fewer columns. This phenomenon becomes increasingly pronounced as the model size grows. We computed attention matrices using 100 tokens from random subset of RedPajama with 1B tokens. Next, we performed 100 runs and plotted the mean and standard deviation of the mass as a function of layers for our mass analysis, respectively. We followed the same Pre-trained checkpoints of OpenLLaMA-333, OpenLLaMA-7B,, and OpenLLaMA-133B (Geng & Liu 2023) trained on 1T tokens from the RedPajama dataset [Computer 2023 were utilized. Overall, we observed that 90 % of the total mass of the attention matrices resides in lewer columns, with many attention matricos in the OpenLLaMA-1133 model being single-column. This observation aligns closely with our analysis in Figure]
of attention matrices in billion-scale LLMs, pre-trained trillions
procedure as discussed in Section
20
16
25
10
15
20
25
32
10
25
(a) Rank analysis of GPTT2 Medium
(b) Rank analysis of GPT-2 Large
(c) Rank analysis of GPT-2 XLarge
25
12
cotmm 62% moae
20
60
15
Arg. cotumnsswt
les. salamss
5
15
20
25
10
15
20
25
30
35
1.0
20
30
Layers
Layers
Layers
(d) Mass analysis of GPT-2 Medium
(e) Mass analysis of GPT-2 Large
(f) Mass analysis of GPT-2 XLarge
Figure 16: In standard decoder-style LLMs, attention matrices in deeper layers often degenerate into single-column matrices, leading to layers with fully degenerated attention that fail to learn meaningful representations. All models were trained on the OpenWehTexx dataset, and both rank and mass analyses were conducted using the FineWeb validation set, following the same procedure described in Figure I This further demonstrates the robustness of our analysis, as we reach the same conclusion using different evaluation datasets.
21

---

Published in Transactions on Machine Learning Research (02/2020)
C.1 Data Robustness of Attention Degeneration Analysis
In Section 2 (Figure] 1 we performed a rank analysis on three pre-trained GPT-2 models Medium, Large, and XLarge using the validation set of the OpenWebText dataset, whose training split was originally used for pre-training these models. Here, we evaluate the data robustness of our analysis by repeating the same procodure on a validation set from FineWeb, a newer and distinct dataset. Except for the dataset substitution, all experimental steps remain identical to those described earlier. The results in Figure [6 consistently show that attention tends to lose rank, particularly in deeper layers, often collapsing into near single-column structures across all models. These findings further reinforce the robustness and generality of our observations.
Moreover, for the LLaMA-3 model (Figure [1] and the OpenLLaMA models (Figure 15), used publicly available model weights and conducted our analyses on off-tho-shelf datascts that were not part of the models' original training corpora.
D Baselines
We compare Inheritune against several baselines. While some baseline methods are illustrated using GPT-2 Large or medium (for the knowledge distillation baseline) as an example, the same methodology is consistently applicd across all model variants.
Baselines trained from scratch (random initializations) We compare our Inheritune-deriedd model against much larger GPT-2 reference models trained from scratch for the same number steps and similar- sized models trained from scratch for both the same and double the number of training steps.
Baselines trained with warm started training methods. Here we compare our model derived using Inheritune, to similar sized models trained with various model initialization and effcient training techniques which requires model to be initialized with trained weights such as stacking, hybrid stacking, and half-width. We explain these bascline training recipes using GPT-2 Large and its variants as an example and apply the same process for other models.
Stacking Gong et al. [2019] J. Reddi et al. (2023) is a model initialization and efficient (stage-wise) training recipe. We train 9-layer GPT-2 Large variant from scratch for 100K steps, then expanded the model 18 layers by copying the weights from layers 0-8 to layers 9-17. Finally we re-trained this new 18-layer GPT-2 Large variant, using stacking initialization for an additional 100K steps.
Hybrid stacking: Hybrid stacking is stacking but utilizes large pre-trained reference model for initialization instead of using its own pre-trained weights. We took the weights of layers 0-8 from the reference 36-layer GPT-2 Large model and expanded it to a 18-layer model by copying the weights to layers 0-17. We then trained this new 18-layer GPT-2 variant for 100K steps.
Half width: We initialized the baseline GPT-2 Large variant across the width dimension and preserved the entire depth. We copied the weights of the first half the attention heads (0-9) and MLPs of the GPT-2 Large reference model into baseline GPT-2 variant with half the width but all layers.
Baselines trained with Knowledge Distillation As a baseline, we first apply logit-based knowledge distillation Hinton et al. (2015) to train a 16-layer GPT-2 Medium variant (student) initialized randomly. For the second baseline, we use a DistillBERT-style approach [anh et al. (2019), where the student model (-11 layers are initialized with every alternate block of its teacher, and the remaining blocks are initialized using layers 18, 19, 20, and 21 of the teacher. Both baselines are trained for 50K steps, using vanilla 24-layer GPT-2 Medium model as the teacher (our reference model).
E Supplementary Experiments
We provide supplementary tables and plots corresponding to the results discussed the main paper, along with additional experiments, in this section. The final validation losses shown in Figure are presented in Table 4. Similarly, the final validation losses from the training curves in Figure 5 and Figure summarized
22

---

Published in Transactions on Machine Learning Research (02/2026)
Layers
Initialization
AvgRank
Pre-train Val Loss (4)
4
rand
N/A
3.25
4 1-4 layers from vanilla GPT2
8.40
3.22
4 5-8 layers from vanilla GPT2
9.48
3.19
4 9-12 layers (lazy layers) from GPT2
1.22
3.23
Table 4: Impact of initialization strategies on GPT-2 small variants. We analyzed the rank characteristics of vanilla. GPT2-small model (125M, 12 layers) trained OpenWebText for 100K steps. Four-layer GPT2-small variants were initialized using the first layers [1 4/, middle 4 layers [5 8), last 4 layers [9-12], or with random initialization, and then trained for 100K steps on OpenWebText. Models initialized with the last 4 layers performed similarly to random initialization, while those initialized with layers exhibiting higher average max ranks achieved the best validation loss, regardless proximity to the embedding layer. The training plots and rank analysis are provided Figure 2
3.60
Va
teacher irit
Vanilla
3.45
Ours
Validation Loss
3.30
3.15
10
20
30
40
50
Steps (In
Figure 17: A 16-layer GPT-2 Medium variant derived using Inheritune converges faster and generalizes better than a same-sized model trained with logit-based distillation baselines.
in Table We also include training curves for models trained using various warm-started basclines (i.c., models initialized with learned weights) compared against our method, these results correspond to Table 2 and are visualized in Figure I8 Finally, Figure 19 presents the training curves for models used in the ablation study discussed in Section
Distillation Inheritune. We conducted additional experiment to compare Inheritune with knowledge distillation as bascline. Specifically, trained GPT-2 Medium variants with 16 layers under three different settings. First, performed logit-based distillation Hinton et al. (2015), transferring knowledge from a 24-layer vanilla GPT-2 Medium (teacher) to 16-layer student model. Second, we applied DistilBERT-styyl distillation Sanh et al. (2019), where the student initialized with the teacher's layers. Finally, we trained a 16-layer GPT-2 Medium model from scratch using vanilla training. Across all comparisons, the model trained with our Inheritune recipe outperformed both distilled variants, achieving faster convergence and substantially better generalization after 50K training steps. We defer thorough investigation of the relationship between Inheritune and distillation-based approaches to future work. The training configurations are provided in Sention
F Architectural and Training Details
F.1 GPT-2 Model configurations
For our main experiments, we focus on three sizes of GPT-2 models Radford (0119): GPT-2 XLarge with 1.5B parameters, GPT-2 Large with 770M parameters, and GPT-2 Medium with 355M parameters. We
23

---

Published in Transactions on Machine Learning Research (02/2026)
Models
Layers
Initialization
Steps
Pre-train Val loss (L)
24
rand init
100K
2.81
16
rand init
100K
2.86
16
rand init
200K
2.83
GPT-2 Medium
Ours
100K
2.87
14
Ours
100K
2.84
Final Model
16
Ours
100K
2.81
rand init
100K
2.85
rand init
100K
2.97
GPT-2 Large
18
rand init
200K
2.84
18
Ours
100K
2.80
48
rand init
100K
2.65
24
rand init
100K
2.69
GPT-2 XLarge
24
rand init
200K
2.62
Ours
100K
2.64
Table 5: Inheritune achieves superior performance with reduced model size. Comparison of Inheritune-traine models (24-layer GPT-2 XLarge, 18-layer GPT-2 Large, and 16-layer GPT-2 Medium) against full-sized countorparts and extended training basclincs. The training steps of two diffcrent basclines are reported in the table, we use validation loss on the OpenWebText validation set. Note: GPT-2 Large and XLarge uses one round of Inheritune; GPT-2 Medium uses three rounds. The lowest and our corresponding validation losses (lower is better) are highlighted in bold.
3.6
3.6
3.6
blackng ny (ogeswah)
stac bipewwat
3.4
3.4
Oura 1ayee
3.4
(overa-441
3.2
3.2
60
103
Steps
Sheps
(a) CPT-2 Medium
(b) GPT-2 Large
(c) GPT-2 XLarge
Figure 18: Models derived using Inheritune outperform three warm-started baselines (Baseline- II) in terms of final validation loss. Our models demonstrate better convergence and generalization compared to all baselines. All the models are trained with OpenWebText for 100K steps. The curves are smoothed for visual clarity.
developed several variants of these models by adjusting the number of layers, i.c., reducing the depth for vanilla models for to be trained with Inheritune and baseline mathods. In one baseline namely, the half-width variant we modified both the hidden size (and consequently, the number of attention heads) in addition to the depth, as shown in Figure 2 The key architectural configurations of the reference, proposed, and baseline models discussed in this paper are summarized in Table
F.2 Training details of GPT-2 models
All GPT-2 models used this study (unless otherwise stated) were pro-trained on the OpenWebText datasot, which contains approximately 10B tokens. We employed a dataloader that samples tokens with replacement, meaning that the tokens used for training not necessarily unique, following the approach of Liu al.
24

---

Published in Transactions on Machine Learning Research (02/2026)
3.3
3.2
MLP Iw? LNO
Saildation Loss
3.1
3.0
2.9
2.8
2.7
0
20
40 60 80 100
Steps (in K)
Figure 19: Full training curves of 16-layer GPT-2 variants trained during ablations. We analyze Inheritune approach while initializing some specific sub-modules transformer blocks. Here, we initialize each transformer block of a 16-layer GPT-2 Medium variant with three different configurations. First, we separately initialize attention and MLPs (FFNs) submodules; second, we initialize the attention and MLP weights while randomly initializing the layer norms. Finally, we perform Inheritune initialize only the attention and MLP weights with all the respective layer norms.
Full (layers-24)
3.6
R1 (layers-
Ours-R22 (layerom
Ours-R3 (layers- 16)
Validation Loss
3.4
3,2
2.8
20
6D 80
100
Steps (in K)
Figure 20: Training curves for the 24-layer GPT-2 Medium model (full model) and three rounds training following Inheritune recipe (to grow the model). We present the training trajectories for all GPT-2 Medium variants trained using the Inheritune recipe. The final model obtained after the third round (R =) 3) with L = 16 layers matches the final validation loss of the full model. All models are trained for 100K steps on the OpenWeb Text dataset.
25

---

Published in Transactions on Machine Learning Research (02/2026)
Model Family
Type
Layers
Hidden Size
Heads
Notes
Reference
48
1600
25
Original architecture
GPT-2 XTarge (1.5B)
Variants
24
1600
25
Reduced-depth
Reference
36
1280
20
Original architecture Reduced depth
GPT-2 Large (770M)
Variants
18
1280
20
Reference
32
1280
20
Custom architecture Reduced depth
CPT-2 Large" (668M)
Variants
16
1280
20
Reference
24
1024
16
Original architecture Reduced depth
GPT-2 Medium (355M)
Variants
16
1024
Reference
768
Original architecture Reduced depth
GPT-2 Small (125M)
Variants
4
768
12
Table 6: Overview of all GPT-2 models used in this study and their architectural configurations. GPT-2 models are Pre-LN based architectures. The model configurations employed for the stacking and hybrid stacking baselines are identical to those of our variants. For the half-width baseline, we used GPT-2 variants with half the hidden size and number of attention heads.
For evaluating the pro-trained models, we used the validation split of the same datasct, which contains 4.4M tokens. The sole exception this setup is the GPT-2 models trained on the FineWeb edu with 10B tokens (Figure 7 where used unique tokens for training by employing dataloader where sample without replacement.
22223
We employed the AdamW optimizer with 81 = 0.90 and B2 0.95. All GPT-2 models were trained on a single NVIDIA A100 GPU (40 GB memory) with gradient accumulation. For the GPT-2 XLarge and its variants, we utilized an NVIDIA H100 GPU. Most hyperparameters were adapted from Liu et al. (023 with key details discussed this section.
dyper-paramele details of GPT-2 Medium and variants.
Batch size: 394K tokens
Learning rate: 3 x 10-4
Warmup steps: 2K
Scheduler type: cosine decayed to 1 10
Weight decay: 0.1
Gradient clipping value:
Total training steps: 100K
Hyper-parameter details of GPT-2 Large and variants.
Batch size: 128K tokens
Learning rate: 2 x 10-4
Warmup steps: 2K
Scheduler type: cosine decayed to 1 10-5
Weight decay: 0.1
Gradient clipping value:
Total training steps: 100K
26

---

Published in Transactions on Machine Learning Research (02/2026)
Hyper-parameter details of GPT-2 XLarge and variants.
Batch size: 128K tokens
Learning rate: 1.5 x 10-4
Warmup steps: 2K
Scheduler lype: cosine decayed to 1 x 10 10-5
Weight decay: 0.1
Gradient clipping value:
Total training steps: 100K
Hyper-parameter details of knowledge distillation training.
We use the below loss as our distillation based training loss. The validation loss the student loss.
Total loss a student loss + (1 - aa) distillation loss
Model: 16-layer GPT-2 Medium variants
a: 0.6
Batch size: 394K tokens
Learning rate: 3 x 10
Warmup steps: 2K
Scheduler lype: cosine decay to To of max learning rate
Weight decay: 0.1
Gradient clipping value:
Total training steps: 50K
G Additional Experiments and Discussions
G.1 Additional Downstream Evaluation
We extend our evaluation to models trained OpenWeb as described in Section (see Table 7 For this analysis, we focus on the largest model considered in this work, GPT-2 XLarge, along with its variants. Dewnstream evaluation is performed on all datasets listed in Section 4 with the addition of Winogrande (Sakaguchi et al. 2020), BoolQ (Clark et al. 2019),, and Wikitext (Merity et al. 2016). Note that the LAMBADA dataset is evaluated under two scttings: missing-wond word prediction (accuracy) and language modeling (perplexity), and therefore reported twice. All evaluations are implemented using the widely adopted bm-eual-harness.
27

---

Published in Transactions on Machine Learning Research (02/2026)
Task
Full model
Ours
Accuracy-based tasks (t)
ARC-E
50.38 1.03
51.22 1.03
PIQA
66.70 1.10
66.87 1.10
SeiQ
77.00 1.33
79.20 1.28
HellaSwag
33.65 0.47
84.20 + 0.47
LAMBADA
39.90 + 0.68
43,30 + 0.69
WinoGrande
51.93 + 1.40
53.28 1 1.40
BoolQ
57.86 + 0.86
60.40 + 0.86
Average
53.92
55.50
Perplexity-based tasks (1)
Wikitext
25.46
25.52
LAMBADA
20.24
16.51
Average 22.85
21.01
Table 7: Downstream evaluation of GPT-2 XLarge (1.5B) trained from scratch vs. 24-layer model trained with Inheritune (Ours). We evaluate both models on 7 accuracy-basel tasks (higher is better) and 2 perplexity-based tasks (lower is better). All models are trained on OpenWebText. Despite using half the depth, the Inheritune model performs favorably compared to the full-sized counterpart. Best average scores are highlighted in bold.
G.2 Additional Evidence of Attention Collapse in Modern LLMs
Following the discussion in Section B.1 we extend our analysis to several additional open-weight base LLMs (see Figures 21 and 22 including Palcon-7B3 OTMo- Cerebras-GPT-22.77 and LLAMA-3-3R
Many of these models incorporate modern architectural components, including Grouped Query Attention (GQA), RoPE positional embeddings, and RMS normalization applied both before and after the attention module and trained with billions-trillions tokens. We visualized the attention heads of cach model using heatmaps of head index versus layer index. Across all models, we observe a prodominance of near rank-1 1 atlention heads, indicating widespread attention collapse.
G.3 Additional Results on the Functional Ineffectiveness of Lazy Layers
Following the discussion in Section 2.1 the poor transferability of lazy layers, conduct additional experiment. We first train a 24-layer GPT-2 Medium model (reference) on OpenWebText for 100K steps. We then train second model of identical size using the same data and training configuration. In this model, the deeper layers (layers 16-24) are initialized with the corresponding weights from the reference model, while the remaining layers are randomly initialized. As shown in Figure 23 reusing lazy layers consistently degrades performance relative to training all layers from scratch. This result complements our findings Section 2.1 and provides further evidence that lazy layers do not encode transferable reusable representations suitable for initialization. We used batch size of 50K tokens with 2 H100 GPUs; all other training details remain unchanged (see Section F.2
http:://agghagaaa.commm////--b
28

---

Published in Transactions on Machine Learning Research (02/2020)
31
Hledd
Figure 21: Rank analysis of the Falcon-7B model reveals widespread attention collapse. We analyzed Falcon-7B (version 1), which comprises 31 layers with 71 attention heads per layer, using the rank metric defined in Section The results are visualized head index the other. Potent (non-collapsed) heads are shown in varying shades of green, with higher intensity indicating higher rank, while rank-collapsed (near rank-1) heads are highlighted in red. Overall, approximately 75% of attention heads exhibit rank collapse, indicating substantial degeneracy across the model.
healmap with layer index
axis and
Moad B600X
Moad 91204
(a) OLMo 1B
(b) Cerebras-GP GPT 2.7B
(c) LLaMA-3 3B
Figure 22: Rank analysis of large open-weight LLMs reveals widespread attention collapse. We analyze three base open-weight language models using the rank metric defined in Section 2 (a) OLMo (1B) exhibits rank collapse in approximately 58% of attention heads; (b) Cerebras-GPa GPT (2. 72. shows similar level of collapse at roughly 58%; and (c) LLaMA-3 (3B) exhibits rank collapse in about 50% of attention heads. Each heatmap plots attention head index versus layer index. Potent (non-collapsed) heads are shown in varying shades of green, with higher intensity indicating higher rank, while rank-collapsed (near rank-1) heads are highlighted in rod.
G.4 Additional Training Results in the Extended Training Regime
In this section we present supplementary result where we have extended the training steps from 100K (as discussed in Section 4 to 200K steps to gauge the potential Inheritune holds for longer training runs. As shown in Figure 24 GPT-2 medium (24 layer full model) is compared against 18 layer GPT-2 medium variant trained following Inheritune recipe. We have inherited = 18 layers based on the previously known best. configuration for GPT-2 medium variant where a 16 layer model matches the performance its full sized counterpart. We observe that the Inheritune model with layers achieves a lower final validation loss (2.86) than the 24-layer bascline (2.87), demonstrating that Inheritune continues to provide benefits even with extended training. We used batch size of 50K tokens with H100 GPUs; all other training details remain unchanged (see Section F.2
29

---

Published in Transactions on Machine Learning Research (02/2026)
4.00
Rand (layers-24)
af nine lazy layars (4a)ers=24)
3.75
Validation Loss
3.50
3.25
3,00
20
40
60
80
100
Steps (in
Figure 23: Initializing deeper layers of GPT-2 Medium with lazy layers degrades performance. We first train a full GPT-2 Medium model from scratch on OpenWebTest for 100K steps. We then construct a second model in which layers 16-24 are initialized using the corresponding lazy layers layers from the trained reference model, while layers 15 are randomly initialized. Both models are trained using the same optimization and data settings. The model whose deeper layers are initialized with lazy layers performs significantly worse than its counterpart trained entirely from scratch.
3.60
Full model (layers-24) 0ayers=18)
3.45
Validation Loss
3.30
3.15
3.00
2.85
25
50
75
100
125
150
175
200
Steps (In
Figure 24: Inheritune continues to provide gains under extended training. Comparison between a full 24-layer GPT-2 Medium model and an 18-layer GPT-2 Medium variant trained using the Inheritune recipe, both trained on OpenWebTText for 200K steps. The Inheritune model achieves a lower final validation loss (2.66) than the full 24-layer baseline (2.87).
G.5 Discussion on Lazy Layers
In this section, provide formal charactcrization of lazy layers through the rank-based criterion.
Definition G.1 (Lazy Layer (Rank-based)). Let (Xa J.V be collection of N sequences, each consisting of T tokens. For each layer 1 E (1,.... L} and head E {1, (1....,Fl H), let An,ha e TxT denote the attention matrix produced by head in layer on input sequence X,.
Let (A(,,,,,) 2 .... (AA,n))) 2 be the singular values of An,... For a threshold T e (0.1), we define the approximate rank as
Σ o'n,,ii 27 )
(T) = min k e (1....,TT
T}
Σ o, 2
(2)
We choose via ablation study; throughout the paper 0.90 (see Section 4.3
30

---

Published in Transactions on Machine Learning Research (02/2026)
Define the head-wise aggregated rank and layer-wise aggregated rank as
Ra (T))
MaxRank(1)
max he (1.....))
Rank Cank(h,4) (h,1).
(3)
For group of L layers (where I denotes the number of contiguous layers in the group), define
AvgRank L
Σ MaxRank(I).
(4)
Lazy Layer. layer I a Lazy Layer MaxRank(() = 1.
Lazy Layers. A group of contiguous blocks of layers is termed Lazy if [AvgRank] 1.
H Theoretical Analysis
H.1 Causal Attention Rank-1 Collapse and Vanishing Gradients
We isolate simple mechanism in decoder-style (causal) self-attention: when head becomes highly con- centrated on single sink position across time (e-g., the first/BOr token), the resulting attention matrix is approximately rank-1, and the corresponding gradients through the softmax saturate, yielding vanishing updates to Wq and WK.
Sctup. We inherit notations discussed in Section 2 Let X E Tx denote sequence of T tokens. Define Q = XXW, == XWK, V = XWv E RT'xd and masked logits
0
i j, i< j,
Ā QK + M, Mij
55)
-00
with attention weights softmax(A) applied row-wise over unmasked indices, and the final output becomes O AV. Let L denote the training objective.
Definition (attention collapse). Let A TxT be an attention matrix (row-stochastic: Agj 2 and Σ Aij = 1 for each row i). Fix an index j" E fl..... T} that is unmasked for all rows (e.g., j 1 for the BOS token). We say that A c-sink-collapsed to j* if, for every row i {I....,?).
Σ Au S 8. 1#1
(6)
Equivalently, each row placos at least 1 - E probability mass on the same column jr, i.c.,
Atj-
Σ j75
Ag 2 1-8,
Wi (I....;
(7)
Theorem H.1 (Attention collapse to rank-1 sink and vanishing gradients). Consider a single-head causal
self-attention module with A = softmax( in equation o then:
LQKT + M) and O = AV. If A is e-sink-collapen to some j Vi
(i) Rank-1 collapse. Let Ao leJ. (rank-1). Then
| Aollo S SeV2T = (2 < 14 Aolla & 1|4 Aollr & <V2TT
(8)
31

---

Published in Transactions on Machine Learning Research (02/2026)
A is numerically rank-1 where > 0 is small.
(ii) Vanishing gradients through softmax. For each row i, letting ( be the entries of that row (restricted to unmasked indices),
ac = -a
a)) ((
Oc dali)
ac aa 1
|| -
ac oa(
2
(9)
aa(a)
Thus, once rows are near one-hot (e (e=+ the gradient to logits vanishes.
(iii) Vanishing gradients to Wo and WK. Moreover,
ac
2 Vd
ac
ac aw K
11
2
0L || 1Q112
lowe
S
| |
(10)
DA F'
Vd
0A
F
so porameter updates to queries/keye shrink linearly with the collapse level E.
Proof. We prove the three claims sequentially.
(i) Let Au = 1e E RTX, Clearly Au has rank 1. As the Frobenius norm the lo norm aggregated over rows, follows that - | =
E Ei=lla(0) - ez- e Fix a row a E IR" (without loss of generality we omit the superscript i), under e-sink-collapse, Σjjr' aj S e. Since A is an attention matrix (row- stochastic), then Σ}= aj
1 --E 0j 2 1-E. E. Now consider the squared error ||a - ez' 12 = (ag. 1)² + Ejj az. Here, (az- - 1)2 S E² (follows from aj. 2 1 - e), and
and aj 2 0, hence ag.
2
Ejjr az < (Σ aj)
s x². Thus, the squared error ||a - ej. 1131
< 20a. Summing over T rows, gives
T
T
|| - Aol =
Σ ||a(0) - -e--11 -
> 2 = 2T=² + || - Aolle I EV2T,
i=1
i=1
Hence, (2) S | - Aolla $ ||A - A S SeV2T (following Weyl's inequality).
(ii) (Gradients are taken only over unmasked logits; masked positions have zero gradient.) Let g a Ol/ and g A @C/0a(0. For the row-wise softmax, the Jacobian is J diag(a) aa hence ğ Jg. Since
T J = 0, || S tr()). Moreover, tr(J) Σ a ( a;) 1 ||| Under s-sink collapse, aj. 2 E, 80
||a|| z (1 e)² and therefore
||l| S 11 (1-1)))) =)² < 2s.
Thus | || < 2e|lg|12, proving equation 9
(iii) We focus on Wo; the argument for WK analogous. By the chain rule, DWe
X 8C σQ
Let G & ac
BA
Since À = QKTT + M, we have 8Q va
LKK. Bounding norms and using | H|| S || H r, (by
definition)
@wa r 0c 1
Oc 1 |||||||||||
S
|
r S vd
aQ
From (ii), row-wise ||a//a(() |l2 I 2e|laC/da ||,, and therefore | S 2: tion 10
Substituting gives equa- 0A
dc
|p.
Intorpretation and training-time implication. Thcorem H.1 connects rank-1 causal attention collapse (sink behavior) and vanishing gradients: once a head becomes hear-deterministic, the softmax Jacobian saturates and the head receives negligible learning signal to recover. Consequently, collapse becomes sticky (hard to recover from) after entering the collapsed regime, subsequent gradient updates tend to be too small to move the head back to higher rank (non-collapsed) state. Over training, this stickiness can accumulate across layers especially in deeper layers which are more prone to vanishing gradients leading to incrcasingly persistent rank collapsed attention pattorns and progrossively smallor query/key weight updatcs.
32


Reviews:
[{"review_id": "moWMBU6eNQ", "reviewer": "Reviewer_Xk1j", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors identify that 'attention matrix rank collapse' occurs in the deeper layers of LLM architectures. They propose that these collapsed layers are functionally ineffective, and thus initialize shallower models using only the high-rank layers from a pre-trained model. They propose an iterative training algorithm to grow models based on these initial layers. They show that these shallower models can reach the performance of larger pretrained models with this recipe. \n\n\n**Strengths:**\n- The demonstration that models with partial-depth can converge to the same performance as pre-trained 'full-depth' models is interesting and novel to the best of my knowledge. \n- The specific 'inheritune' training recipe which allows for this is also novel and is well described. \n- The paper is written clearly and the experimental details appear complete.\n- The inclusion of some ablations is important and helpful.\n\n**Weaknesses:**\n- The results of the 'initialization' experiments in Section 2.1 appear to have a potential confounding explanation which could invalidate the authors conclusion that 'lazy layers are functionally impaired'. Specifically, since 'lazy' layers are always deeper layers (as the authors have shown), it is very likely that simply transferring earlier layers is the key for strong initialization initialization, and the 'lazy' low-rank property could be irrelevant. This would also seem to be backed up by the red line in Figure 2b surpassing the yellow line which has higher rank. \n\t- In my opinion, one way to partially alleviate this confounder would be to perform the initialization experiments of only deeper layers in a deeper network, leaving the early layers random. However this is still only a partial signal for the functional relevance of the layers and I believe fundamentally something beyond initialization experiments should be done to make this claim. \n- The experimental results in Figures 5 & 6 are a bit misleading since the 'training time' requires not only the training time of full model, but also the training times of the iterative refinement steps. For example, are the 2x speedups reported in Figure 6 completely overshadowed by the fact that the models need to be iteratively retrained more than 2 times?\n- There are no error bars on any plots or tables. This is worrisome for results such as those reported in Table 2 where the performance gain is marginal. \n- Many of the models do not appear fully converged in training (e.g. Fig 2b, 5). The choice of 100k training steps for the base model appears somewhat arbitrary. \n- The comparison of 'inheritune' with the more common alternative 'hybrid-stacking' in Table 2 is challenging. Specifically, the results with the X-large model bring into question the benefit the iterative procedure of inheritune, and if it outweighs the multiplicative increase in training time. \n- The authors appear to ignore the fact that their pre-trained initialization is a huge reason for their faster convergence, rather than some deeper theoretical reason. For example, comments such as the following should be removed:\n\t- \"From a convergence perspective, prior work has linked overparameterization to faster convergence Bengio et al. (2005); Vaswani et al. (2018). Interestingly, we find that smaller models derived using Inheritune converge just as fast as their larger counterparts.\"\n- The results are largely overstated in the discussion section. For example, the authors state: \"While stacking and hybrid stacking demonstrate reasonable performance, they still fall short compared to Inheritune. Across all cases, Inheritune consistently outperforms these baselines, highlighting its effectiveness as an initialization strategy.\" Yet, in Table 2, they show that hybrid-stacking performs equivalently to inheritune for the XL mode(2.64 vs. 2.64) -- this is not 'outperforming'.\n- Figure 8 does not appear to support the authors claim that 'inheritune mitigates attention collapse'. The difference in x-axis scale seems to account for the majority of difference in the plots -- i.e. if you only look at the first 16 layers of Fig 8a, it appears to match Fig 8b. If there is truly a difference, the authors should overlay the plots and quantify this difference precisely. \n- I believe there is a relevant citation that should be included, which also discusses attention rank collapse is causal transformers.\n\t- Wu et al., “On the Role of Attention Masks and LayerNorm in Transformers” (NeurIPS 2024)", "claims_and_evidence_yn": "No", "claims_explanation": "As outlined in the weaknesses listed above, there are multiple claims that are only partially supported by evidence, and there may be confounding factors that could invalidate the conclusions. For example, the claims 'lazy layers are functionally ineffective' and 'inheritune mitigates attention collapse', are not fully supported by the experiments the authors perform. Furthermore, the claim that inheritune yields a 2x speedup in training is largely misleading (as the authors note very subtly in the limitations section) since they must retrain the model multiple times prior to achieving this speedup.", "audience_yn": "Yes", "audience_explanation": "The result that you can achieve the performance of a deeper network with fewer layers is likely of interest to some readers, regardless of training time. The proposed Inheritune algorithm indeed does appear to surpass the performance of comparable alternative (e.g. hybrid stacking) for smaller models (but not for larger models). The authors would do better to emphasize this benefit and not over state their results or conclusions.", "requested_changes": "While the phenomenon of attention collapse is important, the experimental evaluation of the 'inheritune' method has a number of potential confounders and conflicting results with respect to the authors claimed conclusions. I therefore recommend the authors carefully revise their claims (as listed in the weaknesses), or improve their experimental evaluation to increase the evidence for their existing conclusions. These changes would be required for the paper to be ready for acceptance in my view.\n\n\n\n**Questions:**\n- Do the authors disagree with the first noted weakness that layer depth may be a potential confounding factor?\n- Do the authors understand why in Figure 2b, the 'init w/ middle four layers' appears to work better than the 'init w/ first four layers', despite these layers appearing more 'lazy'? This would seem to contradict the authors' main point with this figure. \n - Can the authors justify the choice of 100k training steps for the base model? \n- Can the authors comment on the relation of their findings to the Wu et al. 2024 paper?", "broader_impact_concerns": "N/A", "additional_comments": "**Minor:**\n- Figure 3b says (16-layer) but the caption says (18-layer)\n- The description of the inheritune algorithm is a bit odd. While it makes sense, the authors should avoid using ambiguous phrases such as: \"repeat steps 1-2 until desired performance is achieved.\" Clearly the desired performance is as high as possible, so it would be more appropriate to say something like 'repeat until convergence or a decrease in performance'.\n- The writing in the paper overall is a bit informal and would benefit from both being less verbose and more precise.", "extra_scores": {}, "extra_text": {}}, {"review_id": "xw0aHQEhaL", "reviewer": "Reviewer_q1G7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper identifies attention collapse in decoder-only LLMs, showing that deeper layers often degenerate to near rank-1 attention maps, forming lazy layers that contribute little representational value (e.g., Figures 1–3).\nTo address this, the authors propose InheriTune, a progressive training recipe that initializes a compact model using early high-rank layers from a larger pretrained model and grows it until it matches the reference model's validation loss. Empirically, models trained through InheriTune achieve comparable or better performance than larger models and outperform both same-sized scratch-trained and warm-started baselines (Figures 5–7, Table 2).\n\nStrengths:\n* Clear empirical evidence of collapse\n* Proposal of a simple and practical training method\n* Carried out thorough baselines and analyzed useful ablations\n\nWeaknesses:\n* Limited theoretical explanation\n* Restricted architectural diversity\n* Potential computational overhead from multi-stage training", "claims_and_evidence_yn": "Yes", "claims_explanation": "Rank analyses, transfer experiments, and comparisons with strong baselines all support the central claims. Lazy layers are shown to behave like random initialization, confirming functional degeneration. Inheritune consistently matches or outperforms reference models across datasets and model sizes. The evidence is clear, though theoretical grounding could be deeper.", "audience_yn": "Yes", "audience_explanation": "Understanding structural inefficiencies in transformers is highly relevant to researchers working on model compression and dilution, scaling laws, and efficient LLM training. The proposed method offers a practical alternative to pruning or distillation and is of broad interest.", "requested_changes": "* Clarify and formalize the definition of lazy layers  \nCurrently, layers are considered lazy when their attention matrices exhibit MaxRank(l) ≈ 1 (Figure 1) or AvgRank ≈ 1.2 (Figure 2).\nBetter Provide a precise threshold and analyze sensitivity to τ = 0.90 (used in rank calculation). Show how many layers become “lazy” under alternative τ values (e.g., 0.85, 0.95). This is important because conclusions about depth redundancy depend on the thresholding rule.\n\n* Give deeper diagnostic or theoretical insight into why attention collapse arises  \nThe manuscript references prior theory (Dong et al., Noci et al.), but observed collapse in GPT-2 occurs despite residuals, FFNs, and layernorms. Consider adding gradient-norm depth profiles for Q/K (as suggested in Section 2), or analyzing singular vector alignment over layers.  \nData already shows full collapse in GPT-2 Medium's layers 20–24 (Figure 8a); explain why these layers fail to preserve rank.\n\n* Quantify compute cost of Inheritune relative to baselines.\nInheriTune requires multiple training rounds (e.g., GPT-2 Medium uses three rounds of 100K steps; Figure 19), whereas full baselines train once for 100K or 200K steps (Figures 5–6). Report total FLOPs or wall-clock to determine whether InheriTune truly improves overall efficiency, despite faster convergence of the final configuration.\n\n* Expand evaluation to more modern architectures  \nThe rank analysis on collapse is shown for GPT-2 families and limitedly for LLaMA-3 8B (Figure 10), but architectures with post-LN or parallel-attention (GPT-NeoX, MPT, Falcon) may behave differently. Adding results would strengthen claims of generality.  \nInclude tasks where deeper layers are known to matter (e.g., BoolQ, Winogrande, NarrativeQA), to confirm that removing lazy layers does not degrade reasoning quality.\n\n* Ablate the growth schedule in InheriTune  \nAll experiments use +2 layers per round (Section 4). Test +1, +4, or adaptive increments to validate that results are not schedule-specific. Report how many rounds are required to match reference performance under each schedule.", "broader_impact_concerns": "Improving training efficiency for compact LLMs may accelerate widespread deployment, raising typical concerns around safety, misuse, and access to high-capability models. The Broader Impact section should briefly address risks from making strong models easier to train and distribute.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "mq23JqTicr", "reviewer": "Reviewer_sKfL", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work has a few key contributions. \nAuthor(s) define the term attention collapse and show its prevelance in the deeper layers of decoder LLMs. Secondly, they propose Inheritune where one frist initializes a smaller LM using the first few layers from larger pre-trained LM. Then, the smaller model is trained for a certain number of steps. Lastly, the smaller LM can be incrementally grown by adding more blocks of layer until pre-training loss is good or better than the large LM.\n\nStregnths:\n- Clearly state problem (of attention collapse / lazy layer) and supoprt it with emerical results\n- Propose solution to problem and run a series of experiments to show emperically models trained with Inheritune helps train models 1) good performance and 2) less attention collapse\n- LLMs have been shown to be overparametarized. This work helps address this issue developing smaller models (which is nice for efficiency) that have good performance.\n\nWeaknesses:\n- Most evaluations are based only on the pre-training loss.\n- Some minor typos", "claims_and_evidence_yn": "Yes", "claims_explanation": "Author(s) provide clear and convicing evidence \n- 1) that attention collapse occurs (Fig 1), \n- 2) that attention collapse is indeed a problem (seciont 2.1, Fig 2)\n- 3) that Inheritune helps mitigate this problem of attention collapse (Fig 8)\n- 4) that Interitune helps trains smaller models with similar or better performances than training from scratch (large or same size model and warm started model) (Fig 6 and Table 2)", "audience_yn": "Yes", "audience_explanation": "LLM's are known to be over-parametarized/redundant (i.e. it seems that a large portion of the parameters don't do much and mostly just help with training stability and giving the model more memorization power). I believe having methods for training smaller LLM's that converge faster with similar performance to larger model applicable to TMLR's audience.", "requested_changes": "- Typos\n    - Fig 7 caption: \"Models derived using Inheritune without data repetition converges and matches\" \n        - converges -> converge\n        - matches -> match\n    - Page 12: \"successful in training small LMs in some cases\" -> \"successful in training small LMs. In some cases\" \n\n- It would be nice if in Figure 8(a) and (b), the y-axis was on the same scale (i.e. both 0 - 40 or maybe 0 - 45). Or maybe just put them on the same graph with two different line colors/styles.\n- Just for easy comparison, it would be nice to add in a \"from scracth\"/\"random initialization\" 16 layer model results at the beginning of table 3\n- It would be nice to have some justification about mostly using the pre-training loss (i.e. does pre-training loss always correlate with downstream performance? Why only run downstream taks on one set of pre-training models and not pre-training models trained of OpenWeb)", "broader_impact_concerns": "Following https://jmlr.org/tmlr/ethics.html, I don't believe there are any concerns that would require a Broader Impact Statement.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)