Transformer Block Coupling and its Correlation with Generalization in LLMs

TL;DR

This paper investigates the internal dynamics of transformer blocks in LLMs, revealing a phenomenon called transformer block coupling that correlates positively with model performance and generalization, supported by empirical analysis and experiments.

Contribution

It introduces the concept of transformer block coupling, analyzes its emergence during training, and links it to improved generalization in both language and vision transformers.

Findings

Coupling correlates positively with model performance.

Coupling and linearity increase during training.

Coupling observed in both LLMs and Vision Transformers.

Abstract

Large Language Models (LLMs) have made significant strides in natural language processing, and a precise understanding of the internal mechanisms driving their success is essential. In this work, we analyze the trajectories of token embeddings as they pass through transformer blocks, linearizing the system along these trajectories through their Jacobian matrices. By examining the relationships between these block Jacobians, we uncover the phenomenon of \textbf{transformer block coupling} in a multitude of LLMs, characterized by the coupling of their top singular vectors across tokens and depth. Our findings reveal that coupling \textit{positively correlates} with model performance, and that this relationship is stronger than with other hyperparameters such as parameter count, model depth, and embedding dimension. We further investigate how these properties emerge during training, observing a progressive development of coupling, increased linearity, and layer-wise exponential growth in token trajectories. Additionally, experiments with Vision Transformers (ViTs) corroborate the emergence of coupling and its relationship with generalization, reinforcing our findings in LLMs. Collectively, these insights offer a novel perspective on token interactions in transformers, opening new directions for studying their mechanisms as well as improving training and generalization.