Mudança de Representação: Unificando Compressão de Tokens com FlashAttention

Representation Shift: Unifying Token Compression with FlashAttention

August 1, 2025

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

cs.AI

Resumo

Os Transformers têm demonstrado sucesso notável em visão, linguagem e vídeo. No entanto, o aumento da complexidade das tarefas levou a modelos maiores e mais tokens, elevando o custo quadrático da autoatenção e a sobrecarga de acesso à memória da GPU. Para reduzir o custo computacional da autoatenção, trabalhos anteriores propuseram técnicas de compressão de tokens que descartam tokens redundantes ou menos informativos. Enquanto isso, kernels de atenção fundidos, como o FlashAttention, foram desenvolvidos para aliviar a sobrecarga de memória, evitando a construção de mapas de atenção e seu I/O associado à HBM. Isso, no entanto, torna-o incompatível com a maioria dos métodos de compressão de tokens sem treinamento, que dependem de mapas de atenção para determinar a importância dos tokens. Aqui, propomos o Representation Shift, uma métrica sem treinamento e independente de modelo que mede o grau de mudança na representação de cada token. Isso integra perfeitamente a compressão de tokens com o FlashAttention, sem mapas de atenção ou retreinamento. Nosso método ainda se generaliza além dos Transformers para CNNs e modelos de espaço de estados. Experimentos extensivos mostram que o Representation Shift permite uma compressão eficaz de tokens compatível com o FlashAttention, resultando em acelerações significativas de até 5,5% e 4,4% em recuperação de vídeo-texto e QA de vídeo, respectivamente. O código está disponível em 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.