論文「Compressing Sequences in the Latent Embedding Space: K-Token Merging for Large Language Models」の技術的サマリー

本論文は、大規模言語モデル（LLM）が長い入力プロンプトを処理する際に生じる計算コストとメモリ使用量の増大という課題に対し、**「潜在埋め込み空間（Latent Embedding Space）」における冗長性を活用した新しいトークン圧縮手法「K-Token Merging」**を提案するものです。

既存の手法が主にトークンレベルでの削除や要約（ハード圧縮）や、新しいトークンの学習（ソフト圧縮）に焦点を当てていたのに対し、本手法はトークンそのものの埋め込み表現の非効率性に着目し、連続する K 個のトークン埋め込みを単一の埋め込みベクトルにマージすることで、入力長を大幅に削減しつつ性能を維持することを可能にしました。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細をまとめます。

1. 背景と問題定義

LLM の自己注意機構（Self-Attention）は入力長に対して二次関数的に計算コストとメモリ使用量が増加します（$O(n^2)$）。これを緩和するため、入力トークン数を減らす「トークン圧縮」の研究が進められています。しかし、既存のアプローチには以下の限界がありました。

• ハード圧縮（Hard Prompt Compression）: 不要なトークンを削除または要約する手法。情報密度の高いタスク（コード編集、数式推論など）では、重要なコンテキストが失われ性能が低下しやすい。
• ソフト圧縮（Soft Prompt Compression）: 新しいトークンを学習させたりパラメータを適応させたりする手法。これらはトークンレベルで動作しており、埋め込み空間自体の非効率性（例えば、1 つのトークンを識別するのに必要な理論最小ビット数と、実際の埋め込みベクトルが消費するビット数の巨大なギャップ）を見落としていた。

QWEN 2.5の例を挙げると、語彙サイズ 15 万程度に対し、1 トークンの識別に必要な理論最小ビットは約 18 ビットですが、実際の埋め込み（896 次元×32 ビット）は約 28,672 ビットを消費しています。この巨大な無駄を埋め込み空間で解消する余地があることが示唆されました。

2. 提案手法：K-Token Merging

本手法は、連続する K 個のトークン（K-gram）を、軽量なエンコーダを用いて単一の埋め込みベクトルに圧縮する「潜在空間圧縮フレームワーク」です。

2.1 モデル構造

• プレフィル段階（Prefill Stage）:
  • 入力トークンの連続する K 個のブロック $\{T_{iK+1}, \dots, T_{iK+K}\}$ を、軽量なエンコーダ $f$（3 層の MLP など）に入力します。
  • エンコーダはこれらを単一の圧縮埋め込み $C_i$ に変換します。
  • 初期化戦略: 学習の収束を早めるため、エンコーダの出力を「K 個の元のトークン埋め込みの平均（Mean Pooling）」に近づけるように設計されています。具体的には、平均値をベースとし、小さな MLP の出力を残差として加算する構造を採用し、MLP の重みを初期値 0 に設定することで、学習開始時に安定した意味的起点を提供します。
• 生成段階（Generation Stage）:
  • 生成されるトークンは、圧縮された入力とは異なり、**元の語彙のトークン（Uncompressed tokens）**として出力されます。
  • 生成されたトークンは、圧縮されたプレフィックスに付加され、通常の自己回帰生成が続行されます。
  • この設計により、出力側の語彙サイズは増えず、推論時の辞書爆発（Vocabulary Explosion）を防ぎます。

2.2 学習手法

• LoRA 適応: 圧縮された入力に適応させるため、LLM のパラメータに対して LoRA（Low-Rank Adaptation）を適用して微調整を行います。
• 目的関数: 損失関数は、圧縮されたトークン位置ではなく、**生成された元のトークン（Uncompressed targets）**の位置でのみ計算されます。これにより、モデルは圧縮された埋め込みを解釈しつつ、元の語彙での生成品質を維持することを学習します。

3. 主要な貢献

1. K-Token Merging の提案: 潜在埋め込み空間の冗長性を直接ターゲットとした新しいソフト圧縮フレームワークの導入。
2. 適応学習レシピの提示: 圧縮された入力を扱う LLM の微調整方法（エンコーダと LoRA の共学習、平均初期化戦略など）を確立。
3. 高い圧縮率と性能の両立: 多様なタスクにおいて、最大 75% の入力長削減を達成し、性能低下を最小限に抑えることを実証。

4. 実験結果

3 つの異なるタスク（構造的推論、感情分析、コード編集）で評価を行いました。ベースラインには、ハード圧縮手法（SelectiveContext, LLMLingua2）とソフト圧縮手法（LTSC）を使用しました。

• 評価指標: 精度（Accuracy）またはパープレキシティ（Perplexity）と、長さ削減率（Length Reduction Ratio）の調和平均である「P-L F1 スコア」を使用。
• 結果の概要:
  • Textualized Tree（構造的推論）: 入力長を 75% 削減（K=4）しても、精度は 98.38%（Uncompressed: 99.97%）を維持。P-L F1 スコアは 0.851 で、ベースラインを大幅に上回りました。
  • Amazon Reviews（感情分析）: 75% 削減で精度 91.05% を達成。P-L F1 スコア 0.822 で、最も強力なベースラインより 25.5% 改善。
  • CommitPackFT（コード編集）: 75% 削減でパープレキシティ 1.391（Uncompressed: 1.293）。P-L F1 スコア 0.830 で、コードタスクにおいても優れたトレードオフを実現。
• パレート最適性: どのデータセットにおいても、K-Token Merging は「性能 vs 圧縮率」のパレートフロンティア上に位置し、既存手法よりも優れたバランスを示しました。
• 計算コストの削減: 入力長が出力長よりも遥かに長いシナリオでは、プレフィル段階の計算量が $O(n^2)$ であるため、75% 削減は約 94% の計算量削減（FLOPs）に相当すると推定されています。

5. 意義と将来展望

• 埋め込み空間の非効率性の解消: トークンレベルの操作ではなく、埋め込みベクトルそのものの冗長性を活用することで、従来の手法では達成困難だった高圧縮率と高性能の両立を実現しました。
• 汎用性と拡張性: 語彙サイズを増やすことなく任意の K 値に対応でき、未見の K-gram にも一般化可能です。
• 将来の方向性:
  • 生成中の圧縮: 現在は入力のみを圧縮していますが、生成トークンに対しても圧縮を適用することで、さらにメモリ効率を向上させる可能性があります。
  • 再帰的エンコーディング: 8 トークンを 4 に、さらに 4 を 2 にという階層的な圧縮による効率化。
  • 適応的圧縮: テキストの統計的特性に応じて K 値を動的に変更する適応戦略の検討。

結論

K-Token Merging は、LLM の長文脈処理におけるボトルネックである計算コストを劇的に削減する可能性を秘めた画期的な手法です。埋め込み空間の非効率性を意識的に利用することで、情報密度の高いタスクにおいても高い性能を維持しつつ、最大 75% の入力長削減を実現しました。これは、特に長文コンテキストを必要とするアプリケーションにおいて、LLM の実用性と効率性を飛躍的に高める重要なステップとなります。