Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention

論文紹介 Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention 2026/06/06 瀧研 小林 柳青

論文概要 タイトル： Forget, Then Recall: Learnable Compression and Selective Unfolding via Gist Sparse Attention 著者： Yuzhen Mao, Michael Y. Li, Emily B. Fox 概要： Gist Tokenを用いて長文コンテキストを圧縮し、必要なチャンクのみを選択的に復元する ことで、計算効率と精度を両立するGist Sparse Attention（GSA）を提案した。 1

既存手法 従来のSelf-Attentionは全トークン対の相互作用を計算するため、計算量が 𝑂 𝑛2 に増加し、 シーケンスが長くなると計算コストが爆発する。 既存の長文処理手法はSparse Attention(重要な部分だけを見る) いつ使う 推論時のみ: H2 𝑂, StreamingLLM, Quest 学習時: NSA(Native Sparse Attention),DSA (Differentiable Sparse Attention), MoBA (Mixture of Block Attention) どう使う Compression(圧縮，Gist token): [token 1...N] → [gist1, gist2, gist3, ...] [token 1...16] [gist1] [token 17...32] [gist2] [token 33...48] [gist3] ... 論文の課題 “Why not turn compression into routing: compress first, then selectively unfold the right details?” 2

既存手法 NSAはtokenを圧縮、Qによるtop-kのblockを選択、現在位 置の近傍トークンをそのまま保持することを同時に行う 欠点：アーキテクチャの改変が必要 DSAは独立したモジュールのindexerによって、 top-kを選択 して、アテンションを計算する 欠点：外部indexに依存 MoBAはMoEの発想をアテンション機構に応用したもの すべてのKVブロックにアテンションするのではなく、各ク エリに対して最も関連性の高いブロックだけを選んでアテン ションする 欠点：Mean Poolingを使うため、End-to-Endの勾配ベース 最適化ができない 3

目標 1.クエリ適応的 2. End-to-End学習可能 3.アーキテクチャ改変不要 4

Gist Sparse Attention（GSA）とは Gist tokenは複数のトークンを1つに要約した、 学習可能なもの g1はチャンク1の概要、g2はチャンク2の概要… queryがすべてのgist tokenにスコアをつけ、 top-kを選ぶ 選ばれたチャンクの元のトークンを gist tokenととも に復元する 選ばれたgist+チャンクのペアに対してのみ アテンションを計算する 5

Gist token g1 →x1,x2 g2 →g1,x3,x4 g3 →g1,g2,x5,g3 Gist Tokenにより過去の文脈を要約・保持し、 x1,x2などの元のトークンを参照せずにグ ローバル情報へアクセス可能とする。 MoBAの平均プーリングと比べると、 Gist Tokenは 勾配で最適化できる gist causal mask NSAはsummary token生成のために追加MLPが必要 で、アテンション構造を変更する GSAは変更不要 6

Selective Unfolding and Hybrid Attention 関連性スコアに基づいて、スコアが最も高い上位k個のチャンクを選択する。 選ばれたチャンクの元のトークンをとgist tokenを通常のアテンションを計算する。 7

Training GSAは標準的なTeacher Forcingに基づく自己回帰言語モデルの目的関数で学習される。 x1 x2 g1 x3 x4 g2 x5 [compressed context] prefix x6 x7 x8 [generation context] suffix 圧縮 g1 g2 prefix x5 x6 x7 suffix x8 Stage 1: Continued Pretraining（必要） モデルはコンテキスト情報をGist Tokenに圧縮する能力を学習する。圧縮対象となるコンテキスト （Prefix）にはGist Tokenが一定間隔で挿入され、定義されたGist Causal Maskに従って処理される。一方、 Suffixは元のPrefixトークンを見ることができず、Gist Tokenのみを参照して予測を行う。そのため、モデ ルは重要な情報をGist Tokenに保存するように学習する。 ここまで、gist tokenはすでに学習可能であり、他の推論時のみのスパースアテンション手法と比 較して既に大幅な改善をもたらしている。 8

Training Stage 2: Selective Finetuning (選択) 学習中にモデルを選択的展開メカニズムにさらすことで、性能をさらに向上させるこ とができる。 Chunk1 → g1 Chunk2 → g2 Chunk3 → g3 Chunk4 → g4 圧縮 0.1 0.8 0.7 0.2 g1 F g2 T g3 T g4 F Suffix Tokenの Attention(Query, Gist) Top-k k=2 g1 g2 g3 g4 g2 → Chunk2 g3 → Chunk3 Unfold 9

Hierarchical Gist-of-Gist Compression(H-GSA) 本のように、「章 → 段落 → 文」、階層で探す x1 …x100 g1 x101 …x200 g2 x201 … … g1000 … g1000 圧縮 g1 g2 g3 … 圧縮 G1 G2 G3 … G100 10

計算量の比較 •計算量が文章長に対して対数的に増える •理論上、無限の長さにも対応可能 チャンクサイズ L が大きいほど削減効果が大きくなりますが、大きすぎると各gist tokenが担う 情報量が増えすぎて圧縮の質が落ちるというトレードオフがある。 11

実験 実験用モデル 12

実験結果 13

実験結果 14

実験結果 15

Selective Rules Comparison SR-only: 選ばれたchunkのraw tokenのみ（gistなし） AG+SR : 選ばれた＋選ばれなかったchunkのgist tokenをすべて保持 SG+SR : 選ばれたchunkのgist＋ raw tokenのみ SG+SRが一番の理由； 固定されたtoken予算の中で、選ばれなかったchunkのgist tokenを除外し、その分選ばれたchunkの 生tokenを展開する方が、アテンション文脈がより情報豊かになるからです。 また選ばれたchunkのgist tokenも残す理由は、chunk内のtoken間のパターンを捉えた補完的な圧縮 情報を提供するためです。 16

Top-k vs. Top-p 固定的なTop-p選択と比較して、Adaptive Top-k選択がすべてのタスクでより高い性能を示した。 17

結論 • GSAは、「忘れて（圧縮）、思い出す（展開）」というシンプルな発想 で、構造変更も外部モジュールも非微分操作も不要なまま、長コンテキ スト処理の効率と精度を両立したフレームワークである。 • LongBenchとRAGベンチマークにおいて、8×〜32×の圧縮率のもとで 既存のgist圧縮手法およびSparse Attention手法を一貫して上回りまし た。 感想 優点「一つのgist tokenが圧縮とルーティングを同時に担う」という設計はまるで、人間の 記憶のように「要点で覚え、必要なときだけ詳しく思い出す」というエレガントな発想である。 それにより、従来は複数の独立したモジュールを必要としていた問題を標準Transformerの枠 組みの中に統一し、シンプルさと実用性を両立させた。 欠点 超長いシーケンスと多様なタスクに関する検証が欠如 18