💻 computer science

Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation

本論文は、生成の初期段階で特定され、追加の学習を必要とせずに生成プロセス全体で再利用される、各アテンションヘッドの固有の局所的または大域的なアテンションパターンに基づいて可変のキー・バリューキャッシュ予算を動的に割り当てることで、メモリ効率とスループットを向上させる自己回帰型画像生成のための新たなフレームワーク「HeadKV」を提案する。

原著者： Guotao Liang, Baoquan Zhang, Zhiyuan Wen, Yunming Ye

公開日 2026-05-21

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Guotao Liang, Baoquan Zhang, Zhiyuan Wen, Yunming Ye

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

巨大で複雑な壁画を描こうとしているが、作業台は非常に小さいと想像してください。新しい筆致を加えるたびに、絵の文脈を失わないよう、これまでに行ったすべての筆致の参照を保持しなければなりません。やがて、これらの参照で作業台が散らかって手が動かなくなり、描画は極端に遅くなります。

これがまさに自己回帰的画像生成で起きていることです。これらのAIモデルは、画像を小さな断片（トークン）ごとに生成します。新しい断片が既存の断片と調和していることを保証するため、コンピュータはすでに生成したすべての断片を記憶（キャッシュ）しなければなりません。高解像度の画像の場合、この「記憶の作業台」は巨大化し、コンピュータのリソースをすべて消費して、生成を痛烈に遅くします。

旧来の方法：画一的な対応

以前、研究者たちは古い参照を捨てて作業台を片付けようと試みました。しかし、描画プロセスのすべての部分に対して同じ方法で行っていました。AI内のすべての「脳細胞」（アテンションヘッド）が同じ量の履歴を必要だと仮定していたのです。

この論文の著者たちは、これは次のような状況だと気づきました：次の交差点を知るだけでよい人に対して都市全体の地図を与え、国全体をナビゲートしようとしている人に対して小さな街路標識を与えるようなものです。非効率的です。

新しい発見：2 種類の「脳細胞」

これらのAIモデルがどのように思考するかを詳しく調べることで、著者たちはAIの「脳細胞」が実際には2 つの明確な性格タイプに分類されることを発見しました。

「局所的」な隣人：一部の脳細胞は、自分たちのすぐ近くで起きていることだけを気にします。壁の中の単一のレンガを見ている人のように、自分の仕事をするためには、すぐに触れているレンガだけを見れば十分です。
「全球的」な建築家：他の脳細胞は全体像を気にします。建物全体を見ている建築家のように、屋根が崩れないようにするためには、部屋の反対側からの詳細を記憶する必要があります。

解決策：HeadKV（賢い作業台管理者）

著者たちはHeadKVと呼ばれる新しいシステムを作成しました。すべての脳細胞を同じように扱うのではなく、HeadKV は誰が作業しているかによって異なる量のスペースを割り当てる、賢い作業台管理者のように機能します。

「局所的」な隣人の場合：HeadKV は「あなたは最後の数項目だけを見ればよい」と言います。最近の履歴の小さなスライドウィンドウを保持し、残りはすぐに捨てます。これにより、スペースを大幅に節約できます。
「全球的」な建築家の場合：HeadKV は「あなたは全体像を記憶する必要がある」と言います。より大きな履歴を保持しますが、Stratified Token Eviction（階層化トークン排除）と呼ばれる特別なトリックを使用します。

「Stratified Token Eviction」のトリック：
本棚を片付けていると想像してください。「最も重要な」本だけを拾って選べば、2020 年代の本がより大きくて人気があるため、1990 年代の本をすべて捨ててしまうかもしれません。しかし、文脈のために1990 年代の本は依然として必要です！

HeadKV は履歴を「近隣」と「遠方」の2 つの山に分割します。どちらの山からもいくつか重要な本を保持することを保証します。これにより、AI は（羽の質感のような）即座の詳細に集中しているだけで、（鳥全体の形や空の色のような）遠くで重要な詳細を忘れることがなくなります。

学習なしで行う方法

通常、どの脳細胞がどれであるかをコンピュータに教えるには、何年もの追加学習が必要です。HeadKV は巧妙です。その場でこれを判断します。

新しい人に出会うようなものだと考えてください。彼らが「局所的」な思考者か「全球的」な思考者かを知るために、彼らの人生の全歴史を知る必要はありません。最初の数分間彼らを見るだけで十分です。もし彼らが今起きていることだけを話せば、彼らは「局所的」な思考者です。もし遠くからのことを参照し始めれば、彼らは「全球的」です。

HeadKV は画像生成の開始時にごく短い間AIを観察し、どの脳細胞がどれであるかを決定してから、プロセスの残りの部分に対して適切なメモリルールを適用します。追加の学習は不要で、AI が生成しようとするあらゆる画像で機能します。

結果

この「賢い作業台」アプローチを使用することで、著者たちは以下のことを示しました。

メモリ使用量の削減：必要な元のメモリの1/4、1/6、あるいは1/8 しか保持しなくてもよくなりました。
速度の向上：不要なデータに溺れていないため、AI ははるかに高速に画像を生成します。
品質の維持：画像は、完全で散らかったメモリで作成されたものと同じように美しく見えます。

要するに、HeadKV はAI が無用の情報を溜め込むのを止めさせ、より少ない労力で美しい絵を速く描けるようにします。

技術的サマリー：効率的な自己回帰的画像生成のためのヘッド意識型キー・バリュー圧縮

問題定義

画像を視覚トークンの系列として表現する自己回帰（AR）視覚生成モデルは、モデルサイズとトークン数のスケーリングにより、高解像度かつ写実的な結果を達成してきた。しかし、このスケーリングは、自己回帰的デコーディング中に必要なキー・バリュー（KV）キャッシュにより、莫大なメモリオーバーヘッドと低い推論スループットをもたらす。自然言語処理（NLP）における最近の研究では、KV キャッシュ管理（トークンの選択、排除、量子化など）が探求されてきたが、画像の固有の空間構造により、これらの技術を視覚ドメインに直接適用することは困難である。

既存の視覚固有の手法、例えば LineAR は、トークンが主に空間的に隣接する隣接トークンに注意を向けるという観察に基づき、トークンのサブセットのみを保持することでキャッシュ使用量を削減しようとする。しかし、これらのアプローチは、すべての注意ヘッドに対して固定された均一な予算を割り当てる。この戦略は、注意ヘッド間の異質性を見落とし、メモリ割り当ての最適化を阻害する。具体的には、一部のヘッドは局所的な領域に焦点を当て（小規模な予算を要する）、他のヘッドは物体の構造や背景領域のような長距離依存関係を捉える（より大きな予算を要する）。均一な割り当ては、これらの異なるパターンに適応できず、メモリ浪費または生成品質の低下を招く。

手法：HeadKV

著者は、ヘッド意識型 KV キャッシュ圧縮を実行する、新しいトレーニング不要のフレームワークHeadKVを提案する。核心的な洞察は、AR 画像生成モデルにおける注意ヘッドが局所的と大域的という 2 つの明確なパターンを示し、与えられた入力に対して同じレイヤー内のトークン位置にわたってヘッドの挙動が一貫していることである。

HeadKV パイプラインは、以下の 3 つの主要なステップから構成される：

キャッシュプリフィリング：モデルは条件付き入力をエンコードし、圧縮なしで視覚トークンの初期系列を生成する。この段階では、ヘッドごとの注意分布を推定するために十分なコンテキストが蓄積される。
ヘッドグルーピング（分類）：キャッシュが事前に定義された長さに達すると、フレームワークは閾値ベースの注意蓄積基準を用いて、各注意ヘッドを「局所的」または「大域的」に分類する。
- 各ヘッドについて、現在のトークン位置から逆方向に注意スコアを蓄積し、閾値（ $\tau = 0.9$ ）に達するまで行う。
- 閾値に達するために必要なトークン数が、事前に定義された局所ウィンドウサイズより小さい場合、そのヘッドは局所的と分類される。それ以外の場合は大域的と分類される。
- この分類は生成の初期段階で 1 回だけ実行され、その後のデコーディングプロセス全体で再利用されるため、追加のトレーニングやデータセットレベルの統計は不要である。
ヘッド意識型キャッシュ圧縮：ヘッドタイプに基づいて異なる圧縮戦略が適用される。
- 局所的ヘッド：これらのヘッドは主に近傍のコンテキストに注意を向ける。フレームワークは、スライディングウィンドウ機構を介して最も最近のトークンのみを保持するコンパクトな予算を割り当てる。過去のトークンは完全に省略される。
- 大域的ヘッド：これらのヘッドは長距離依存関係を捉える。線形なメモリ成長を招くことなく情報豊富な遠方のトークンを保持するために、著者は**階層化トークン排除（STE）**戦略を設計する。
  - 階層化トークン排除：候補となる過去のトークンは、現在の位置からの距離に基づいて、近距離トークンと長距離トークンの 2 つの互いに排他的なサブセットに分割される。
  - 重要度スコアは、過去のクエリ特徴を用いて、各サブセットに対して独立して計算される。
  - スコアが最も低いトークンが、各サブセットから排除される（例：近距離から $M/2$ 、長距離から $M/2$ ）。これにより、局所的に支配的なトークンが意味的に関連する遠方のトークンを覆い隠すことを防ぎ、局所的な詳細と大域的な構造の両方をバランスよく保持する。

主要な貢献

注意の異質性の分析：本論文は AR 画像生成モデルの深い分析を行い、2 つの明確な注意パターン（局所的および大域的）を特定し、固定入力に対してレイヤー内のトークン位置にわたってヘッドの挙動が不変であることを明らかにする。
HeadKV フレームワーク：注意ヘッド間で非対称な予算を割り当てる、新しいトレーニング不要の圧縮フレームワーク。初期段階の分類メカニズムに導かれ、局所的ヘッドにはコンパクトな予算を、大域的ヘッドにはより大きな予算を割り当てる。
階層化トークン排除（STE）：大域的ヘッド向けの特定の戦略であり、トークンを距離で分割することで、情報豊富な長距離依存関係の保持を確保する。これは、しばしば局所的トークンを優先する標準的な重要度スコアリングの限界を克服する。
実証的検証：広範な実験により、HeadKV が複数のモデル（Janus-Pro、Lumina-mGPT、LlamaGen-XL）において、高品質な画像を維持しながら KV キャッシュメモリ使用量を大幅に削減し、推論スループットを向上させることが示された。

結果

著者は、4 つの自己回帰的画像生成バックボーン（Janus-Pro-1B、Janus-Pro-7B、Lumina-mGPT-768、LlamaGen-XL）に対して HeadKV を評価した。

画像品質：HeadKV は、GenEval や DPG などのベンチマークにおいて、フルキャッシュのベースラインと同等のパフォーマンスを達成する。例えば、圧縮率（ $\rho$ ）が 1/8 の Janus-Pro-7B において、HeadKV は GenEval スコア 0.99（フル：0.99）、DPG スコア 83.78（フル：84.20）を維持する。
LineAR との比較：均一な割り当て手法である LineAR と比較して、HeadKV は圧縮率の増加に伴う性能低下が小さい（例： $\rho=1/8$ ）。LlamaGen-XL において、極端な圧縮条件下で HeadKV は 2.75 という優れた FID-50K スコアを達成し、LineAR（2.95）を上回る。
効率性：HeadKV はメモリ使用量を大幅に削減し、スループットを向上させる。
- LlamaGen-XL（ $\rho=1/6$ ）において、HeadKV はメモリを 63.60% 削減し、フル設定に対して3.00 倍の高速化を達成。LineAR（2.93 倍）を上回る。
- Janus-Pro-7B（ $\rho=1/6$ ）において、メモリを 52.08% 削減し、3.68 倍の高速化を達成。
- Lumina-mGPT-768において、グルーピングの実装上のオーバーヘッドによりレイテンシが LineAR よりわずかに高いものの、メモリ大幅削減により 1.35 倍の高速化を達成。

意義と主張

本論文は、HeadKVが追加のトレーニングや外部教師なしを必要とせず、自己回帰的画像生成におけるメモリボトルネックに対する実用的な解決策を提供すると主張する。注意ヘッドの固有の異質性を認識し活用することで、この手法は先行研究の「万能型」キャッシュ割り当てを超えている。

著者は、生成の初期段階で観察される内在的な注意パターンに依存するため、このアプローチは異なる入力やモデルアーキテクチャに対して汎用性があることを強調する。階層化トークン排除戦略は、標準的な排除手法でしばしば失われる長距離の構造的情報を保持することにより、厳しいメモリ予算下でも生成画像の品質を維持するための重要な要素として注目されている。この研究は、均一なトークン保持よりも、適応的でヘッド固有のリソース割り当ての方が、視覚生成にとってより効率的なパラダイムであることを示唆している。