Changement de représentation : Unification de la compression de tokens avec FlashAttention

Representation Shift: Unifying Token Compression with FlashAttention

August 1, 2025

papers.authors: Joonmyung Choi, Sanghyeok Lee, Byungoh Ko, Eunseo Kim, Jihyung Kil, Hyunwoo J. Kim

cs.AI

papers.abstract

Les Transformers ont démontré un succès remarquable dans les domaines de la vision, du langage et de la vidéo. Cependant, l'augmentation de la complexité des tâches a conduit à des modèles plus volumineux et à un plus grand nombre de tokens, ce qui a accru le coût quadratique de l'auto-attention et la surcharge d'accès à la mémoire GPU. Pour réduire le coût de calcul de l'auto-attention, des travaux antérieurs ont proposé des techniques de compression de tokens qui éliminent les tokens redondants ou moins informatifs. Parallèlement, des noyaux d'attention fusionnés tels que FlashAttention ont été développés pour atténuer la surcharge mémoire en évitant la construction de cartes d'attention et les E/S associées vers la HBM. Cela rend cependant ces méthodes incompatibles avec la plupart des techniques de compression de tokens sans entraînement, qui s'appuient sur les cartes d'attention pour déterminer l'importance des tokens. Nous proposons ici le *Representation Shift*, une métrique sans entraînement et indépendante du modèle, qui mesure le degré de changement dans la représentation de chaque token. Cela permet d'intégrer de manière transparente la compression de tokens avec FlashAttention, sans nécessiter de cartes d'attention ni de réentraînement. Notre méthode se généralise également au-delà des Transformers aux CNN et aux modèles à espace d'états. Des expériences approfondies montrent que le *Representation Shift* permet une compression efficace des tokens compatible avec FlashAttention, offrant des accélérations significatives allant jusqu'à 5,5 % et 4,4 % dans la recherche vidéo-texte et les Q&A vidéo, respectivement. Le code est disponible à l'adresse suivante : https://github.com/mlvlab/Representation-Shift.

English

Transformers have demonstrated remarkable success across vision, language, and video. Yet, increasing task complexity has led to larger models and more tokens, raising the quadratic cost of self-attention and the overhead of GPU memory access. To reduce the computation cost of self-attention, prior work has proposed token compression techniques that drop redundant or less informative tokens. Meanwhile, fused attention kernels such as FlashAttention have been developed to alleviate memory overhead by avoiding attention map construction and its associated I/O to HBM. This, however, makes it incompatible with most training-free token compression methods, which rely on attention maps to determine token importance. Here, we propose Representation Shift, a training-free, model-agnostic metric that measures the degree of change in each token's representation. This seamlessly integrates token compression with FlashAttention, without attention maps or retraining. Our method further generalizes beyond Transformers to CNNs and state space models. Extensive experiments show that Representation Shift enables effective token compression compatible with FlashAttention, yielding significant speedups of up to 5.5% and 4.4% in video-text retrieval and video QA, respectively. Code is available at https://github.com/mlvlab/Representation-Shift.