Сдвиг представления: унификация сжатия токенов с использованием FlashAttention
Representation Shift: Unifying Token Compression with FlashAttention
August 1, 2025
Авторы: Joonmyung Choi, Sanghyeok Lee, Byungoh Ko, Eunseo Kim, Jihyung Kil, Hyunwoo J. Kim
cs.AI
Аннотация
Трансформеры продемонстрировали выдающиеся успехи в задачах, связанных с обработкой изображений, текста и видео. Однако увеличение сложности задач привело к созданию более крупных моделей и увеличению количества токенов, что повысило квадратичную стоимость механизма самовнимания и накладные расходы на доступ к памяти GPU. Для снижения вычислительной стоимости самовнимания в предыдущих работах предлагались методы сжатия токенов, которые удаляют избыточные или менее информативные токены. В то же время были разработаны оптимизированные ядра внимания, такие как FlashAttention, которые уменьшают накладные расходы на память, избегая построения карт внимания и связанных с ними операций ввода-вывода в высокоскоростную память (HBM). Однако это делает их несовместимыми с большинством методов сжатия токенов, не требующих обучения, которые полагаются на карты внимания для определения важности токенов. В данной работе мы предлагаем метрику Representation Shift, которая не требует обучения и является универсальной для различных моделей. Она измеряет степень изменения представления каждого токена, что позволяет интегрировать сжатие токенов с FlashAttention без использования карт внимания или повторного обучения. Наш метод также обобщается за пределы трансформеров, применим к сверточным нейронным сетям (CNN) и моделям пространства состояний. Эксперименты показывают, что Representation Shift обеспечивает эффективное сжатие токенов, совместимое с FlashAttention, что приводит к значительному ускорению до 5,5% и 4,4% в задачах поиска видео по тексту и ответов на вопросы по видео соответственно. Код доступен по адресу 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.