Non Sfrutti Appieno la Capacità Rappresentativa del Transformer

You Do Not Fully Utilize Transformer's Representation Capacity

February 13, 2025

Autori: Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky, Viacheslav Sinii, Daniil Gavrilov

cs.AI

Abstract

A differenza delle RNN, che comprimono i token precedenti in un unico stato nascosto, i Transformer possono prestare attenzione direttamente a tutti i token precedenti. Tuttavia, i Transformer standard utilizzano solo le rappresentazioni dello strato immediatamente precedente. In questo articolo, dimostriamo che questa scelta progettuale causa un collasso delle rappresentazioni e porta a prestazioni subottimali. Per affrontare questo problema, introduciamo la Layer-Integrated Memory (LIMe), un approccio semplice ma potente che preserva l'impronta di memoria complessiva del modello, espandendo al contempo la sua capacità rappresentativa consentendo l'accesso agli stati nascosti degli strati precedenti. Attraverso esperimenti estesi su varie architetture e diversi meccanismi di ricerca, dimostriamo miglioramenti consistenti delle prestazioni su un'ampia gamma di task. Inoltre, la nostra analisi della dinamica delle rappresentazioni apprese e l'esplorazione dei circuiti in profondità rivelano come LIMe integri le informazioni attraverso gli strati, indicando direzioni promettenti per la ricerca futura.

English

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.