Cambio de Representación: Unificación de la Compresión de Tokens con 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

Resumen

Los transformadores han demostrado un éxito notable en visión, lenguaje y video. Sin embargo, el aumento en la complejidad de las tareas ha llevado a modelos más grandes y a un mayor número de tokens, incrementando el costo cuadrático de la autoatención y la sobrecarga en el acceso a la memoria de la GPU. Para reducir el costo computacional de la autoatención, trabajos previos han propuesto técnicas de compresión de tokens que eliminan aquellos redundantes o menos informativos. Paralelamente, se han desarrollado núcleos de atención fusionados, como FlashAttention, para aliviar la sobrecarga de memoria al evitar la construcción de mapas de atención y su correspondiente E/S hacia la HBM. Esto, sin embargo, los hace incompatibles con la mayoría de los métodos de compresión de tokens sin entrenamiento, que dependen de los mapas de atención para determinar la importancia de los tokens. Aquí, proponemos Representation Shift, una métrica sin entrenamiento y agnóstica al modelo que mide el grado de cambio en la representación de cada token. Esto integra de manera fluida la compresión de tokens con FlashAttention, sin necesidad de mapas de atención o reentrenamiento. Nuestro método se generaliza más allá de los transformadores, aplicándose también a CNNs y modelos de espacio de estados. Experimentos extensivos muestran que Representation Shift permite una compresión efectiva de tokens compatible con FlashAttention, logrando aceleraciones significativas de hasta el 5.5% y 4.4% en la recuperación de texto-video y en preguntas y respuestas sobre video, respectivamente. El código está disponible en 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.