표현 전환: 토큰 압축과 FlashAttention의 통합

초록

트랜스포머(Transformers)는 비전, 언어, 비디오 분야에서 놀라운 성공을 거두었습니다. 그러나 작업의 복잡성이 증가함에 따라 더 큰 모델과 더 많은 토큰이 필요해지면서, 셀프 어텐션(self-attention)의 이차 비용과 GPU 메모리 접근 오버헤드가 증가하고 있습니다. 셀프 어텐션의 계산 비용을 줄이기 위해, 기존 연구에서는 중복되거나 덜 중요한 토큰을 제거하는 토큰 압축 기법을 제안했습니다. 한편, FlashAttention과 같은 융합 어텐션 커널(fused attention kernels)은 어텐션 맵(attention map) 구축과 HBM(High Bandwidth Memory)으로의 I/O를 피함으로써 메모리 오버헤드를 완화하기 위해 개발되었습니다. 그러나 이는 대부분의 학습 없이 수행되는 토큰 압축 방법과 호환되지 않는데, 이러한 방법들은 토큰의 중요도를 결정하기 위해 어텐션 맵에 의존하기 때문입니다. 여기서 우리는 각 토큰의 표현 변화 정도를 측정하는 학습 없이 적용 가능한 모델-불특정(agnostic) 메트릭인 Representation Shift를 제안합니다. 이는 어텐션 맵이나 재학습 없이도 FlashAttention과 토큰 압축을 원활하게 통합합니다. 우리의 방법은 트랜스포머를 넘어 CNN(Convolutional Neural Networks)과 상태 공간 모델(state space models)에도 일반화됩니다. 광범위한 실험을 통해 Representation Shift가 FlashAttention과 호환되는 효과적인 토큰 압축을 가능하게 하여, 비디오-텍스트 검색과 비디오 QA(Question Answering)에서 각각 최대 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.

표현 전환: 토큰 압축과 FlashAttention의 통합

Representation Shift: Unifying Token Compression with FlashAttention

초록

Support