On the Role of Attention Masks and LayerNorm in Transformers

TL;DR

This paper analyzes how attention masks and LayerNorm influence rank collapse in transformers, revealing that these components can slow or prevent collapse, thereby enhancing model expressivity and dynamics.

Contribution

It provides the first comprehensive analysis of rank collapse considering attention masks and LayerNorm, showing their roles in preventing collapse and increasing expressivity.

Findings

Sparse or local attention slows rank collapse.

LayerNorm can prevent exponential rank collapse with proper value matrices.

Transformers with LayerNorm have diverse equilibrium states.

Abstract

Self-attention is the key mechanism of transformers, which are the essential building blocks of modern foundation models. Recent studies have shown that pure self-attention suffers from an increasing degree of rank collapse as depth increases, limiting model expressivity and further utilization of model depth. The existing literature on rank collapse, however, has mostly overlooked other critical components in transformers that may alleviate the rank collapse issue. In this paper, we provide a general analysis of rank collapse under self-attention, taking into account the effects of attention masks and layer normalization (LayerNorm). In particular, we find that although pure masked attention still suffers from exponential collapse to a rank one subspace, sparse or local masked attention can provably slow down the collapse rate. In the case of self-attention with LayerNorm, we first show that for certain classes of value matrices, collapse to a rank one subspace still happens exponentially. However, through construction of nontrivial counterexamples, we then establish that with proper choice of value matrices, a general class of sequences may not converge to a rank one subspace, and the self-attention dynamics with LayerNorm can simultaneously possess a rich set of equilibria with any possible rank between one and full. Our result refutes the previous hypothesis that LayerNorm plays no role in the rank collapse of self-attention and suggests that self-attention with LayerNorm constitutes a much more expressive, versatile nonlinear dynamical system than what was originally thought.