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

Transformer은 비전, 언어, 영상 분야에서 놀라운 성공을 거두었으나, 작업의 복잡성이 증가함에 따라 모델 크기와 토큰 수가 증가하면서 자기주의(self-attention)의 이차 비용과 GPU 메모리 접근 오버헤드가 증가하는 문제가 발생하고 있다. 자기주의의 계산 비용을 줄이기 위해 기존 연구에서는 중복되거나 정보량이 낮은 토큰을 제거하는 토큰 압축 기법을 제안해왔다. 한편, FlashAttention과 같은 융합형 주의 커널(fused attention kernels)은 주의 맵(attention map)을 생성하지 않고 HBM에 대한 I/O를 피함으로써 메모리 오버헤드를 완화하는 방식으로 개발되었다. 그러나 이러한 접근은 주의 맵을 기반으로 토큰의 중요도를 평가하는 대부분의 훈련 불필요한 토큰 압축 기법과 호환되지 않는다. 본 연구에서는 각 토큰의 표현 변화 정도를 측정하는 훈련 불필요하고 모델 독립적인 지표인 '표현 이동(Representation Shift)'을 제안한다. 이 방법은 주의 맵 없이도 재훈련 없이 토큰 압축과 FlashAttention을 원활하게 통합할 수 있다. 또한 본 방법은 Transformer를 넘어서 CNN과 상태 공간 모델(state space models)에도 일반화 가능하다. 광범위한 실험을 통해 Representation Shift가 FlashAttention과 호환되는 효과적인 토큰 압축을 가능하게 하며, 영상-텍스트 검색 및 영상 질의응답(Video QA)에서 각각 최대 5.5%, 4.4%의 성능 향상을 달성함을 확인하였다. 코드는 https://github.com/mlvlab/Representation-Shift 에 공개되어 있다.