Transformers need glasses! Information over-squashing in language tasks
This addresses a fundamental limitation in Transformer-based LLMs that affects their reliability in tasks requiring precise token sensitivity, such as counting or copying, which is incremental as it builds on known over-squashing phenomena in neural networks.
The paper identifies a representational collapse in decoder-only Transformers, where distinct input sequences can yield arbitrarily close final token representations, exacerbated by low-precision formats, leading to errors in tasks like counting or copying. It provides theoretical analysis and empirical evidence on contemporary LLMs, suggesting simple solutions to mitigate these issues.
We study how information propagates in decoder-only Transformers, which are the architectural backbone of most existing frontier large language models (LLMs). We rely on a theoretical signal propagation analysis -- specifically, we analyse the representations of the last token in the final layer of the Transformer, as this is the representation used for next-token prediction. Our analysis reveals a representational collapse phenomenon: we prove that certain distinct sequences of inputs to the Transformer can yield arbitrarily close representations in the final token. This effect is exacerbated by the low-precision floating-point formats frequently used in modern LLMs. As a result, the model is provably unable to respond to these sequences in different ways -- leading to errors in, e.g., tasks involving counting or copying. Further, we show that decoder-only Transformer language models can lose sensitivity to specific tokens in the input, which relates to the well-known phenomenon of over-squashing in graph neural networks. We provide empirical evidence supporting our claims on contemporary LLMs. Our theory also points to simple solutions towards ameliorating these issues.