(2024)

Abstract: This paper reveals a novel linear characteristic exclusive to transformer decoders, including models such as GPT, LLaMA, OPT, BLOOM and others. We analyze embedding transformations between sequential layers, uncovering a near-perfect linear relationship (Procrustes similarity score of 0.99). However, linearity decreases when the residual component is removed due to a consistently low output norm of the transformer layer. Our experiments show that removing or linearly approximating some of the most linear blocks of transformers does not affect significantly the loss or model performance. Moreover, in our pretraining experiments on smaller models we introduce a cosine-similarity-based regularization, aimed at reducing layer linearity. This regularization improves performance metrics on benchmarks like Tiny Stories and SuperGLUE and as well successfully decreases the linearity of the models. This study challenges the existing understanding of transformer architectures, suggesting that their operation may be more linear than previously assumed. 1

heading: 1 Introduction
text: Transformers have revolutionized the field of natural language processing, offering unprecedented advances in a wide range of applications (Islam et al., 2023). However, despite their widespread adoption and success, the complex work of these models remains an area of active research (Lin et al., 2021). One aspect that has received less attention is the inherent linearity of intermediate embedding transformations within these architectures. In this study, we embark on an in-depth analysis of the linearity properties of transformers, specifically focusing on decoders, and explore its implications during the pretraining and fine-tuning phases.
1https://github.com/AIRI-Institute/ LLM-Microscope
Our investigation reveals a surprising discovery: the embedding transformations between sequential layers in transformer decoders exhibit almost linear properties. This observation is quantified using Procrustes similarity analysis, demonstrating a near-perfect linearity score of 0.99. Such a discovery not only challenges the traditional understanding of transformer architectures but also opens new opportunities for model optimization and efficiency.
Based on this insight, we introduce several new contributions to the field:
• Extensive analysis of the linearity properties of transformer decoders and its dynamics at the pretraining and fine-tuning stages.
• The development of new algorithms for depth pruning of transformer decoders, allowing to remove the most linear layers without a significant loss in performance.
• A novel distillation technique that involves pruning, replacing certain layers with linear approximations, and then distilling layer-wise embeddings to preserve model performance.
• Introducing a new regularization approach for pretraining based on the cosine similarity, designed to decrease the layer linearity. This method not only enhances the performance of transformer models on benchmark datasets such as SuperGLUE and TinyStories Eldan and Li (2023), but also improves the expressiveness of embeddings, as evidenced by linear probing tasks.
With our findings, we are paving the way for more computationally efficient transformer architectures without sacrificing their effectiveness, thereby addressing one of the critical challenges in deploying these models.
5376

heading: 2 Related Work
text: Research on evaluating and leveraging sparsity for model pruning has become one of the most significant topics within the machine learning community. Molchanov et al. (2016) explored the sparsity of convolutional neural networks through backpropagation and fine-tuning, laying the groundwork for understanding the potential applications of sparsity in resource-efficient inference. The verification approach utilized in a more recent DejaVu (Borse et al., 2023) paper is based on Molchanov’s research.
Previous work (Kurtic et al., 2023) has addressed the challenges associated with naive sparse finetuning in the context of LLMs. Issues such as training instability, poor recovery, and overfitting have prompted an exploration for alternative approaches. The study introduced SquareHead distillation, a method that consistently addresses the challenges in naive sparse fine-tuning, demonstrating accurate recovery even at high sparsity levels.
In a more recent study WANDA (Sun et al., 2023), the authors present a technique for pruning LLMs to high degrees of sparsity without modifying the remaining weights. Unlike SparseGPT (Frantar and Alistarh, 2023), WANDA seamlessly implements pruning in a single forward pass, leveraging feature norm statistics for efficient pruning. This method achieves noticeable sparsity without the need for a sophisticated iterative weight update procedure, differentiating itself from other pruning techniques.
Contextual sparsity introduced by Borse et al. (2023) involves sparsifying MLP and attention blocks in LLMs to reduce generation latency. The study identifies essential attention heads and MLP
neurons for computation, maintaining performance across in-context learning and language modeling tasks.
Recent work by Ashkboos et al. (2024) shows that LLMs can be sparsified post hoc. Their approach introduces a scheme to replace each weight matrix with a smaller dense matrix, thereby reducing the dimensionality of the networks. Their results show that models of different sizes can be reduced with varying degrees of success. For example, LLAMA-2 70B and OPT 66B can maintain 99% zero-shot accuracy while reducing 25% of the parameters reduced while performing LLM evaluation tasks. In contrast, the smaller Phi-2 is more sensitive to pruning, experiencing a 10% drop compared to its dense version.
The inner structure of transformer models has captured significant attention among researchers (Nostalgebraist, 2020; Xu et al., 2021; Belrose et al., 2023; Din et al., 2023). Primarily, in “logit lens” (Nostalgebraist, 2020) and subsequently in (Belrose et al., 2023), the authors have focused on analyzing how hidden representations evolve across different layers of transformer architecture, aiming to elucidate their impact on final model outputs. Complementing these findings, the Anthropic team’s research into small transformer-based models (Elhage et al., 2021) uncovers a profound linear structure inherent in this architecture. Their work demonstrates the effectiveness of decomposing operations into individual sum components and multiplying chains of matrices, thus highlighting the linear complexity within these sophisticated neural architectures.
Structure-based pruning Topological features that analyze the structure of inner embeddings in
transformer-based models are also useful in LLM pruning and distillation. Previous research examined the intrinsic dimensionality of neural networks to evaluate their capacity and effectiveness in the fine-tuning process (Ansuini et al., 2019; Aghajanyan et al., 2020; Razzhigaev et al., 2023). Decoder-based models are shown to achieve a high level of anisotropy, especially in their middle layers, and have low intrinsic dimensionality (Razzhigaev et al., 2023). Recent popular approaches include low-rank approximation, which replaces or adjusts the weight matrix with the product of two matrices with a smaller inner dimension. This approach typically requires a fine-tuning procedure that adjusts the matrix representations. For example, LoRA (Hu et al., 2021) was inspired by the previous work (Aghajanyan et al., 2020) showing that neural networks can be successively trained in lower-dimensional subspaces. The research also shows that there it is not necessary to update millions of parameters on small fine-tuning datasets. Our results are on par with the results of this research, showing that via fine-tuning, the linearization of models grows steadily.
The Bonsai model (Dery et al., 2024) tends to prune the LLMs relying only on the inference step, while they achieve performance comparable to halfsized semistructured sparsity with WANDA 2:4 and outperforms the LLM-Pruner (Ma et al., 2023) and LoRAPrune (Zhang et al., 2023) on 4 out of 6 evaluation settings in the experiments conducted.
In this paper, we investigate several techniques for pruning LLMs, leveraging the linearity of the decoder-based layers. Our techniques offer efficient yet lightweight methods, maintaining high model performance on the evaluated benchmarks.

heading: 3 Analysis of Pretrained Architectures
text: In our study of the embedding properties of various layers of transformer decoders, we focus on understanding the degree of linearity and smoothness of transformations between sequential layers.

heading: 3.1 Linearity Score
text: To determine the degree of linear dependence of two sets of vectors, we used a metric obtained by generalizing the Procrustes similarity (Gower, 1975) to the case of arbitrary linear transformations.
Let X,Y ∈ Rn×d represent the centered sets of embeddings, to calculate linearity score we use normalized matrices X̃ = X/||X||2, Ỹ = Y/||Y ||2
(where || · ||2 denotes the Frobenius norm of the matrix) and defined
linearity_score := 1− min A∈Rd×d ||X̃A− Ỹ ||22
This is almost the same formula as in Procrustes similarity, the only difference is that, instead of considering the minimum among orthogonal transformations, we use the minimum among all linear transformations to find the optimal mapping in terms of squared errors.
We chose such approach for its robustness in evaluating the linearity of embeddings, especially considering the scale variance across transformer layers. Unlike L2 norm, which lacks scale invariance, Procrustes normalization offers a bounded metric in the range [0,1].
Surprisingly, the linearity scores of layers in all tested transformer decoders were found to be close to 1, indicating a high degree of linearity in embedding transformations (Figure 1).
This phenomenon can be partly explained by the observation that the norm of each block’s contribution to the residual stream is remarkably low (Figure 3). Moreover, when assessing the linearity of the main stream (embeddings w/o residual component) by subtracting the embedding values of each layer from the previous layer, one can notice that the degree of linearity significantly decreases (Figure 1). This suggests that the inherent linearity is not as straightforward as it is initially estimated. Moreover, the low norm contribution of individual blocks resulted in embeddings from adjacent layers being closely aligned in terms of cosine similarity.
One more insight is that the combination of seemingly linear blocks can lead to non-linear out-
comes. Elhage et al. (2022) suggests that complex features can be encoded across components of neural networks, applicable to attention heads in transformers. This indicates that the cumulative effect of linear transformations might enable the encoding of intricate non-linear representations.
Furthermore, our feature triggering regime hypothesis proposes that rare specific features on a few tokens with high non-linearity significantly influence model behavior — in the Figure 9 one can see that some layers of OPT-1.3B have the long tailed distribution of L2 errors, which means that there are still sparse spikes of non-linearity.
Borse et al. (2023) explored how a sparse subset of model parameters can be dynamically activated for efficient inference, supporting the idea that within predominantly linear architectures, certain non-linear interactions are crucial for model functionality.

heading: 3.2 Linearity Dynamics at Pretraining and Fine-tuning
text: Our exploration extends to examining the linearity dynamics of both open-source models with publicly available intermediate checkpoints and our custom models trained on small datasets. Through this analysis, we aimed to understand the dynamics of linearity, especially in the main stream (contextualized embeddings including the residual component), across different stages of model training.
As illustrated in the Figure 2, the analysis reveals a notable trend: as the models undergo pretraining, the linearity of the main stream gradually decreases
on average. This phenomenon is consistently observed in all models examined, indicating a fundamental aspect of transformer-decoder learning dynamics.
In our analysis of the fine-tuning phase across diverse tasks, including those in the SuperGLUE benchmark (Wang et al., 2019) and the rewardmodeling task on the Anthropic-Helpful dataset (Bai et al., 2022), we notice an interesting change. Contrary to the decreasing trend of linearity observed during the pretraining phase, all models under study show an increase in linearity during fine-tuning. This finding indicates that task-specific fine-tuning tends to reinforce and amplify the linear characteristics of transformer models, as shown in Table 1.
In fine-tuning, we train models on three NLI tasks from the SuperGLUE benchmark: MultiRC, BoolQ, and CB, treating them as binary text classification challenges. In the BoolQ task, for instance, we combine the question and the passage into a single text, marking them with "question:" and "passage:" respectively, and consider the binary answer as the classification label.
Reward models trained on text pairs with contrastive loss (Ouyang et al., 2022) demonstrate a similar trend in linearity scores, proving even more stability across different seed values.

heading: 4 Improving Linearity with Regularized Pretraining
text: Aiming to understand the impact of linearity on transformer models, we embark on pretraining experiments using the Mistral architecture with model sizes of 150M, and 650M. These models are pretrained on carefully selected clean datasets, TinyStories (Eldan and Li, 2023) and Tiny-textbooks (Li et al., 2023), chosen for their diverse and rich content, which has been proven to be suitable for fast training of the small models (Zhao et al., 2023) and architecture experiments (Sharifnassab et al., 2024).
We introduce specific loss terms to adjust the relations between embeddings within transformer layers:
• MSE regularization term: Experimentation with mean squared error (MSE) loss between embeddings of consecutive layers, designed to minimize the distance between these embeddings, thereby promoting consistency across the layers.
LMSE = λ ∑ (∥embi − embi−1∥2).
• Cosine Similarity regularization term: The application of a cosine-based regularization that encourages contextualized embeddings from sequential layers to align closer to each other, effectively reducing their angular difference to zero on average.
Lcosine = λ ∑ (1− cos(embi, embi−1)).
The most promising results are achieved using a cosine-based approach that encourages the embeddings of sequential layers to converge, effectively making the cosine similarity between them closer to 1 on average. This method shows significant perspectives in the enhancing model performance. We evaluate the effectiveness of our approach through validation using GPT-4 on TinyStories prompts according to the Eldan and Li (2023) methodology, linear probing techniques, and evaluation on SuperGLUE benchmarks. The results are presented in the Table 2 and Table 3. As it can be seen in the Figure 5, linearity scores are lower at each layer of the model after pretraining with such regularization.
To further assess the expressiveness of embeddings across different layers, we conducted linear probing on outputs of all the layers of the Mistral650M model, both pretrained with and without cosine regularization, on the xstorycloze-en task
from SuperGLUE. The results clearly indicate that embeddings from the model pretrained with regularization exhibit better performance compared to those from the standard model (Figure 4).
This contradictory outcome, where the term appears to draw embeddings from neighbouring layers closer together, making them more similar in terms of cosine similarity, has prompted a deeper investigation. Our observations suggest that as embeddings become more similar across layers, the model may compensate for the reduction in variability by amplifying non-linear processing capabilities in the residual stream. Although this hypothesis requires further exploration, it offers a fascinating insight into the adaptive mechanisms of transformer models in response to altered internal dynamics.

heading: 5 Exploiting Linearity for Pruning
text: Leveraging the inherent linearity of transformer layers, we explore a pruning strategy that sequentially removes the most linear layers. This approach allows you to reduce the size of the model slightly by removing just a few layers without significantly compromising performance. Further enhancement of this strategy involves replacing the pruned layers with linear approximation and incorporating a distillation loss (specifically MSE layerwise) to
minimize performance degradation. The training focuses on these linear replacements, fine-tuning them to effectively mimic the original layers’ function. The effectiveness and the impact of these methods are detailed in the Figure 8. We use TinyStories for linear approximation and distillation training stage. As it can be seen in the Figure 7, perplexity is less affected by pruning with linear replacements and following distillation compared to just removing transformer layers.

heading: 6 Conclusion
text: In our study we provide an in-depth exploration of linearity within transformer decoders, revealing their inherent near-linear behavior in various models. We discover that while pretraining tends to increase nonlinearity within layers, fine-tuning on specific tasks can paradoxically reduce it. We propose new pruning and distillation techniques inspired by previous observations, demonstrating
that it is possible to refine and optimize transformer models without compromising their performance. The suggested cosine-based regularization approach during pretraining further contributes to model efficiency and performance on benchmarks such as SuperGLUE and TinyStories, while reducing the linearity of its layers (w/o residual components).
Our study highlights the significant relationship between linearity and performance of transformer decoders, offering strategic guidance for future developments in the efficiency and flexibility of these models.

heading: 7 Limitations
text: Despite the promising advancements presented in this study, it is essential to acknowledge its limitations. Firstly, our analysis predominantly focuses on transformer decoders, thus the generalizability of our findings to encoder-only or encoder-decoder architectures may be limited.
Secondly, the depth pruning and distillation techniques, while being effective in our experiments, were evaluated within a specific set of conditions and models. The scalability of these methods to larger, more complex models or different domains is yet to be fully ascertained.
Moreover, the new regularization approach aimed at pretraining demonstrates potential, yet its effectiveness across a broader spectrum of tasks requires further validation.

heading: 8 Ethics Statement
text: We are committed to ethical principles for AI research, focusing on transparency and responsible
experimentation. Our research, while suggesting efficiency improvements, prompts consideration of implications such as privacy and fairness.

heading: A Error Distribution by Layers
text: In the Figure 9 we present a visualization of L2 error distribution across several layers of OPT-1.3B decoder architecture.

heading: Future Work
text: Our study highlights the significant relationship between linearity and performance of transformer decoders, offering strategic guidance for future developments in the efficiency and flexibility of these models