18日前

表現のシフト:トークン圧縮とFlashAttentionの統一

Joonmyung Choi, Sanghyeok Lee, Byungoh Ko, Eunseo Kim, Jihyung Kil, Hyunwoo J. Kim
表現のシフト:トークン圧縮とFlashAttentionの統一
要約

Transformerは、視覚、自然言語、動画の分野において顕著な成功を収めてきた。しかし、タスクの複雑さが増すにつれ、モデルの規模やトークン数が増大し、自己注意(self-attention)の計算コストが二次関数的に上昇する一方で、GPUメモリへのアクセスオーバーヘッドも増大している。自己注意の計算コストを低減するため、先行研究では冗長または情報量の少ないトークンを削除するトークン圧縮技術が提案されている。一方で、Attentionマップの構築を回避し、HBMへのI/Oを削減することでメモリオーバーヘッドを軽減する、FlashAttentionのような統合型Attentionカーネルも開発されている。しかしながら、こうした手法は、トークンの重要度を評価するためにAttentionマップに依存する多くのトレーニングフリーなトークン圧縮手法と互換性がなく、その適用を制限している。本研究では、トレーニング不要かつモデル非依存の指標として「Representation Shift」を提案する。この指標は、各トークンの表現の変化度を測定することで、トークンの重要度を評価する。この手法はAttentionマップを必要とせず、再訓練も不要であり、FlashAttentionとシームレスに統合可能である。さらに、本手法はTransformerにとどまらず、CNNや状態空間モデル(State Space Models)にも拡張可能である。広範な実験により、Representation ShiftがFlashAttentionと併用可能な効果的なトークン圧縮を実現できることを示した。動画-テキスト検索および動画QAにおいて、それぞれ最大5.5%および4.4%の高速化を達成した。コードは https://github.com/mlvlab/Representation-Shift にて公開されている。