You Do Not Fully Utilize Transformer's Representation Capacity
Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky, Viacheslav Sinii, Daniil Gavrilov
2025-02-19
Summary
This paper talks about a new way to improve how Transformer models, which are a type of AI used in language processing, use their internal memory. The researchers found that these models weren't using their full potential and created a method called Layer-Integrated Memory (LIMe) to fix this.
What's the problem?
Current Transformer models only use information from the layer immediately before the current one when processing data. This leads to a problem called 'representation collapse' where the model doesn't use all the information it could, which makes it less effective than it could be.
What's the solution?
The researchers developed Layer-Integrated Memory (LIMe), which allows the model to access information from earlier layers, not just the one right before it. This method doesn't increase the overall memory the model uses but helps it make better use of the information it has already processed. They tested LIMe on various types of Transformer models and different tasks, showing that it consistently improved performance.
Why it matters?
This matters because Transformer models are used in many AI applications, from language translation to generating text. By making these models more efficient and effective without needing more computational power, LIMe could lead to better performance in a wide range of AI tasks. This could mean more accurate translations, more natural language generation, and improvements in other areas where AI processes language or data sequences.
Abstract
In contrast to RNNs, which compress previous tokens into a single hidden state, Transformers can attend to all previous tokens directly. However, standard Transformers only use representations from the immediately preceding layer. In this paper, we show that this design choice causes representation collapse and leads to suboptimal performance. To address this issue, we introduce Layer-Integrated Memory (LIMe), a simple yet powerful approach that preserves the model's overall memory footprint while expanding its representational capacity by allowing access to hidden states from earlier layers. Through extensive experiments across various architectures and different lookup mechanisms, we demonstrate consistent performance improvements on a wide range of tasks. Moreover, our analysis of the learned representation dynamics and our exploration of depthwise circuits reveal how LIMe integrates information across layers, pointing to promising directions for future research.