FlashAttention対応のトークン圧縮『Representation Shift』の仕組みと効果
3つの要点
✔️ Representation Shiftを提案し、トークンの表現変化量で重要度を測定する手法を導入
✔️ Attentionマップに依存せず、FlashAttentionやCNN・SSMにも適用可能な汎用性を持つ
✔️ 実験で最大5.5倍の推論速度向上を達成し、精度と効率の両立を実証
Representation Shift: Unifying Token Compression with FlashAttention
written by Joonmyung Choi, Sanghyeok Lee, Byungoh Ko, Eunseo Kim, Jihyung Kil, Hyunwoo J. Kim
(Submitted on 1 Aug 2025)
Comments: International Conference on Computer Vision (ICCV), 2025
Subjects: Computer Vision and Pattern Recognition (cs.CV)
概要
本論文は、トランスフォーマーモデルにおける計算コストの増大という課題に対して、新たなアプローチを提示しています。
近年、自然言語処理や画像・動画理解の分野でトランスフォーマーが広く用いられる一方、自己注意機構は入力トークン数の二乗に比例して計算量が増加するため、大規模化に伴い処理効率が深刻な問題となっています。
従来、この課題には二つの方向性からの解決が試みられてきました。
一つは FlashAttention に代表されるメモリ効率化手法であり、もう一つは トークン圧縮 に基づく計算削減手法です。
しかし、トークン圧縮は通常、注意マップを利用してトークンの重要度を推定するため、FlashAttentionのように注意マップを構築しない仕組みとは両立できませんでした。
そこで著者らは、新しい指標である Representation Shift を提案。これは各トークンが層を通過する際にどれだけ表現が変化したかを測定し、重要度を定義する方法です。
この指標は学習不要かつモデル非依存であり、FlashAttentionとの組み合わせを可能にします。
実験の結果、この手法は従来法と比較して効率と精度の両面で優れた性能を示し、最大5.5倍の推論速度向上を実現しました。
提案手法
提案手法である Representation Shift は、層の入力と出力における各トークンの埋め込み表現の差分を測定することで、そのトークンがモデル内部でどれほど情報を強調されたかを定量化します。
具体的には、MLP層や注意層を通過する前後のベクトル間距離を計算し、その値を重要度スコアとして利用。このとき距離の測定にはL2ノルムが最も安定した性能を示しました。
従来の手法が注意マップに依存するのに対し、本手法は 注意メカニズムに依存せず にトークン重要度を推定できるため、FlashAttentionのように注意マップを構築しない計算方式とも自然に統合可能。
また、この枠組みはTransformerに限らずCNNや状態空間モデル（SSM）にも適用できる汎用性を持ちます。
著者らはさらに、どの層でRepresentation Shiftを測定するか、どの演算（AttentionかMLPか）に基づくかといった設計選択も詳細に検討。
その結果、MLP層での変化量を利用することが最も効果的であることが示されました。
この設計により、トークンの冗長性を除去しつつ、情報損失を最小限に抑えることが可能となります。
実験
著者らは提案手法の有効性を検証するため、画像分類と動画理解の両タスクで大規模な実験を行いました。
まず、動画タスクでは UMT（Unmasked Teacher） を用いたビデオ・テキスト検索やビデオQAを対象に、トークンを層ごとに20%削減する設定で評価しました。
その結果、FlashAttentionと組み合わせたRepresentation Shiftは、既存の注意スコアに基づく手法よりも高速かつ高精度であり、最大で5.5倍のスループット向上を達成しました。
また、単なるモデルの小型化と比べても、より良好な速度・精度トレードオフを示しました。
次に、画像分類では DeiT 系列を用いてImageNetで検証し、FlashAttentionとの組み合わせにより、推論速度を1.2倍向上させつつ、従来の注意ベース手法を超える精度改善を実現。
さらに、ResNetやVision MambaといったCNN・SSMにも適用し、これらの非Transformer系アーキテクチャにおいても有効性が確認されました。
特にResNet-50では行単位のトークンプルーニングによって18%以上の高速化を達成し、精度をほぼ維持しました。
これらの実験は、Representation Shiftが汎用的かつ強力なトークン圧縮基準であることを示しています。
この記事に関するカテゴリー
