Spostamento della Rappresentazione: Unificare la Compressione dei Token con FlashAttention

Representation Shift: Unifying Token Compression with FlashAttention

August 1, 2025

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

cs.AI

Abstract

I Transformer hanno dimostrato un successo straordinario nei campi della visione, del linguaggio e del video. Tuttavia, l'aumento della complessità dei compiti ha portato a modelli più grandi e a un maggior numero di token, incrementando il costo quadratico dell'auto-attenzione e l'overhead dell'accesso alla memoria GPU. Per ridurre il costo computazionale dell'auto-attenzione, lavori precedenti hanno proposto tecniche di compressione dei token che eliminano quelli ridondanti o meno informativi. Nel frattempo, kernel di attenzione fusi come FlashAttention sono stati sviluppati per alleviare l'overhead della memoria evitando la costruzione delle mappe di attenzione e il relativo I/O verso la HBM. Questo, tuttavia, li rende incompatibili con la maggior parte dei metodi di compressione dei token senza addestramento, che si basano sulle mappe di attenzione per determinare l'importanza dei token. Qui, proponiamo Representation Shift, una metrica senza addestramento e indipendente dal modello che misura il grado di cambiamento nella rappresentazione di ciascun token. Questo integra perfettamente la compressione dei token con FlashAttention, senza bisogno di mappe di attenzione o riaddestramento. Il nostro metodo si generalizza ulteriormente oltre i Transformer ai CNN e ai modelli a spazio di stati. Esperimenti estesi dimostrano che Representation Shift consente una compressione efficace dei token compatibile con FlashAttention, ottenendo accelerazioni significative fino al 5,5% e al 4,4% nel recupero video-testo e nelle domande e risposte sui video, rispettivamente. Il codice è disponibile all'indirizzo 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.