Representatieverschuiving: Het verenigen van tokencompressie met FlashAttention

Representation Shift: Unifying Token Compression with FlashAttention

August 1, 2025

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

cs.AI

Samenvatting

Transformers hebben opmerkelijke successen geboekt op het gebied van visie, taal en video. Echter, de toenemende complexiteit van taken heeft geleid tot grotere modellen en meer tokens, wat de kwadratische kosten van self-attention en de overhead van GPU-geheugentoegang verhoogt. Om de rekenkosten van self-attention te verminderen, hebben eerdere onderzoeken tokencompressietechnieken voorgesteld die redundante of minder informatieve tokens verwijderen. Tegelijkertijd zijn gefuseerde aandachtskernels zoals FlashAttention ontwikkeld om de geheugenoverhead te verminderen door de constructie van aandachtmaps en de bijbehorende I/O naar HBM te vermijden. Dit maakt het echter onverenigbaar met de meeste trainingsvrije tokencompressiemethoden, die afhankelijk zijn van aandachtmaps om de belangrijkheid van tokens te bepalen. Hier stellen we Representation Shift voor, een trainingsvrije, model-agnostische metriek die de mate van verandering in de representatie van elk token meet. Dit integreert naadloos tokencompressie met FlashAttention, zonder aandachtmaps of hertraining. Onze methode generaliseert verder dan Transformers naar CNN's en state space-modellen. Uitgebreide experimenten tonen aan dat Representation Shift effectieve tokencompressie mogelijk maakt die compatibel is met FlashAttention, wat aanzienlijke snelheidswinsten oplevert van respectievelijk 5,5% en 4,4% in video-tekstretrieval en video QA. Code is beschikbaar op 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.