Your Transformer is Secretly Linear — part 1/2

Cleaned-text code points [0, 15824). Bibliography omitted here; source.md includes it.

# Your Transformer is Secretly Linear

###### 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 1 1 https://github.com/AIRI-Institute/LLM-Microscope

## 1 Introduction

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.

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.

## 2 Related Work

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 fine-tuning 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 half-sized 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.

## 3 Analysis of Pretrained Architectures

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.

### 3.1 Linearity Score

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 ∈ ℝ n × d X,Y\in\mathbb{R}^{n\times d} represent the centered sets of embeddings, to calculate linearity score we use normalized matrices X ~ = X / ‖ X ‖ 2 \tilde{X}=X/||X||_{2} , Y ~ = Y / ‖ Y ‖ 2 \tilde{Y}=Y/||Y||_{2} (where | | ⋅ | | 2 ||\cdot||_{2} denotes the Frobenius norm of the matrix) and defined

linearity_score := 1 − min A ∈ R d × d ⁡ ‖ X ~ ​ A − Y ~ ‖ 2 2 \textrm{linearity\_score}:=1-\min_{A\in R^{d\times d}}||\tilde{X}A-\tilde{Y}||_{2}^{2}

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 L 2 L_{2} 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 outcomes. 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 L 2 L_{2} 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.

### 3.2 Linearity Dynamics at Pretraining and Fine-tuning

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 reward-modeling 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.

## 4 Improving Linearity with Regularized Pretraining

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. L M ​ S ​ E = λ ​ ∑ ( ‖ e ​ m ​ b i − e ​ m ​ b i − 1 ‖ 2 ) . L_{MSE}=\lambda\sum(\|emb_{i}-emb_{i-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. L c ​ o ​ s ​ i ​ n ​ e = λ ​ ∑ ( 1 − cos ⁡ ( e ​ m ​ b i , e ​ m ​ b i − 1 ) ) . L_{cosine}=\lambda\sum(1-\cos(emb_{i},emb_{i-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.