Only Large Weights (And Not Skip Connections) Can Prevent the Perils of Rank Collapse

TL;DR

This paper demonstrates that large weights are essential in attention mechanisms of large language models to prevent layer collapse, which limits expressiveness, and that quadratic time complexity in attention computations is unavoidable for expressive transformers.

Contribution

It establishes that large weights, rather than skip connections, are necessary to avoid layer collapse and maintain expressiveness in attention-based models.

Findings

Large weights prevent layer collapse in attention models.

Skip connections do not prevent layer collapse with small weights.

Quadratic time complexity in attention is unavoidable for expressive transformers.

Abstract

Attention mechanisms lie at the heart of modern large language models (LLMs). Straightforward algorithms for forward and backward (gradient) computation take quadratic time, and a line of work initiated by [Alman and Song NeurIPS 2023] and [Alman and Song NeurIPS 2024] has shown that quadratic time is necessary unless the model weights are small, in which case almost linear time algorithms are possible. In this paper, we show that large weights are necessary to avoid a strong preclusion to representational strength we call layer collapse, which means that the entire network can be approximated well by a network with only a single layer. Thus, the quadratic running time of attention is unavoidable for expressive transformers. The notion of layer collapse that we introduce is a variant on the notion of rank collapse from the work of [Dong, Cordonnier, and Loukas ICML 2021]. They showed that in Self Attention Networks with small weights and with skip connections, rank collapse must occur. This is typically interpreted as justifying the necessity of skip connections in expressive networks. However, our result shows that even with skip connections, if the weights are small, then layer collapse still occurs. Thus, only large weights, and not skip connections, can prevent these representational weaknesses.