Your Transformer is Secretly Linear (2024)

Abstract: This paper reveals a novel linear characteristic exclusive to transformer decoders, including models like GPT, LLaMA, OPT, BLOOM and others. We analyze embedding transformations between sequential layers, uncovering an almost perfect linear relationship (Procrustes similarity score of 0.99). However, linearity decreases when the residual component is removed, due to a consistently low transformer layer output norm. Our experiments show that pruning or linearly approximating some of the layers does not impact loss or model performance significantly. Moreover, we introduce a cosine-similarity-based regularization in our pretraining experiments on smaller models, aimed at reducing layer linearity. This regularization not only improves performance metrics on benchmarks like Tiny Stories and SuperGLUE but 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.

## introduction transformers have revolutionized the field of natural language processing, offering unprecedented advances in a wide range of applications @xcite . however, despite their widespread adoption and success, the complex work of these models remains an area of active research @xcite . 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. based on this insight, we introduce several new contributions to the field: 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. ## related work research on evaluating and leveraging sparsity for model pruning has become one of the most significant topics within the machine learning community. @xcite 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 @xcite paper is based on molchanov's research. previous work @xcite 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 @xcite , the authors present a technique for pruning llms to high degrees of sparsity without modifying the remaining weights. unlike sparsegpt @xcite , 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 @xcite 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 @xcite 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 @xcite @xcite . primarily, in "logit lens" (nostalgebraist, 2020) and subsequently in @xcite , 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 @xcite 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 @xcite @xcite . decoder-based models are shown to achieve a high level of anisotropy, especially in their middle layers, and have low intrinsic dimensionality @xcite . 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 @xcite was inspired by the previous work @xcite 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 @xcite 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 @xcite and loraprune @xcite 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. ## 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. ## 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 @xcite and tiny-textbooks @xcite , chosen for their diverse and rich content, which has been proven to be suitable for fast training of the small models @xcite and architecture experiments @xcite . we introduce specific loss terms to adjust the relations between embeddings within transformer layers: 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. ## exploiting linearity for pruning 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. ## conclusion 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.