Your Transformer is Secretly Linear

TL;DR

The paper reveals a surprising near-linearity in transformer-decoder embeddings, quantified by a Procrustes similarity of $0.99$ between successive layers. It introduces the Linearity Score and analyzes its dynamics during pretraining and fine-tuning, showing decreasing linearity in pretraining but increases with task-specific fine-tuning. A cosine-similarity-based regularization is proposed to reduce linearity during pretraining, yielding improvements on $TinyStories$ and $SuperGLUE$ and better linear probing, while enabling more efficient representations. Additionally, the authors propose pruning and distillation strategies that leverage linearity to reduce model depth with minimal performance loss, arguing for a more linear-than-expected view of transformer operation and practical gains in efficiency.

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.

Figures (9)

• Figure 1: Linearity profiles for different open source models. Normalized depth is the layer index divided by the total depth.
• Figure 2: Linearity score (averaged across layers) at different pretraining steps of open source models.
• Figure 3: The relationship between transformer block output norm and resulted residual stream embedding norm.
• Figure 4: Linearity score of different layers with and without cosine regularization used at pretraining.
• Figure 5: Linear probing of embeddings from different layers of Mistral-650M pretrained with and without suggested cosine regularization.