論文要約：YOCO++: 効率的な LLM 推論のための KV 残差接続を備えた YOCO の強化

本論文は、大規模言語モデル（LLM）の推論におけるメモリボトルネックを解決するための「クロスレイヤー KV 圧縮」手法を改良した**YOCO++**を提案するものです。既存の手法である YOCO の性能劣化を解消しつつ、推論効率を維持・向上させることに成功しています。

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

---

1. 背景と問題定義

Transformer ベースの LLM において、推論時のメモリ消費の主要なボトルネックはキー・バリュー（KV）キャッシュです。これを削減するため、トランスフォーマーの特定の層の KV 情報を他の層で共有・再利用する「クロスレイヤー KV 圧縮」手法（例：YOCO, LCKV, CLA など）が提案されています。

• 既存手法 YOCO の課題:
  • YOCO は中間層の KV を上半分の層で共有することで KV キャッシュを 50% 削減し、推論効率を向上させます。
  • しかし、単純な KV 共有はモデルの表現能力（キャパシティ）を低下させ、無視できない性能劣化を引き起こします。
• 既存の解決策の限界:
  • 複数の層の KV を重み付けして組み合わせる手法（例：FusedKV）は性能を改善しますが、キャッシュ I/O オーバーヘッドが増加し、推論効率を低下させるというトレードオフがありました。

2. 提案手法：YOCO++

YOCO++ は、YOCO の構造を維持しつつ、**KV 間の重み付き残差接続（Weighted Residual Connections）**を導入することで、モデル容量を向上させつつ追加の I/O オーバーヘッドを回避する手法です。

2.1 核心的なアイデア

• 残差接続の KV への適用:
  • 従来の残差接続は「隠れ状態（Hidden States）」間で行われますが、YOCO++ はKV 行列自体に残差接続を導入します。
  • 具体的には、ボトムハーフ（下層）の各層 $i$ において、その層の KV ($K_i, V_i$) と、最下層（Bottom Layer）の KV ($K_1, V_1$) を重み付けして組み合わせます。
• なぜ最下層なのか:
  • 理論的にはすべての過去の層を加えるのが理想ですが、実証研究（Zhou et al., 2025）により、最下層の値（Value）を加えるだけで同程度の効果が得られることが示されています。YOCO++ はこの知見に基づき、計算コストを抑えるため「最下層のみ」を参照対象とします。

2.2 数式と実装

各ボトムハーフ層 $i$ において、キャッシュされる KV は以下のように計算されます：

$$\begin{aligned} K'_i &= \lambda (\alpha^K_{i,1} K_1 + \alpha^K_{i,2} K_i) \\ V'_i &= \lambda (\alpha^V_{i,1} V_1 + \alpha^V_{i,2} V_i) \end{aligned}$$

• $\alpha$ は学習可能な残差重みパラメータです。
• $\lambda$ はスケーリングファクター（固定ハイパーパラメータ、実験では 35 を使用）であり、重みの学習を促進するために導入されました。
• 初期化戦略: 学習開始時は YOCO と同等の状態にするため、$\alpha_{i,1}=0, \alpha_{i,2}=1/\lambda$ に初期化されます。

2.3 効率性の維持

YOCO++ は、以下の点で YOCO と同等の推論効率を維持します：

• メモリ消費: 結合された KV ($K'_i, V'_i$) をキャッシュするため、YOCO がキャッシュする元の KV と同じメモリ量で済みます。
• I/O オーバーヘッド: 推論（デコーディング）段階では、直前のトークンのキャッシュされた結合 KV にアクセスするだけでよく、追加の I/O 負荷は発生しません。
• 計算コスト: 重み付け結合の計算はアテンション機構に比べて無視できるレベルです。

3. 主要な貢献

1. YOCO++ の提案: KV 残差接続を導入し、クロスレイヤー KV 共有による性能劣化を解消した新しいアーキテクチャ。
2. スケーリングファクターの導入: 残差重みの学習を安定化・促進させるための $\lambda$ パラメータの設計。
3. 効率と性能の両立: 追加のキャッシュ I/O オーバーヘッドなしに、標準 Transformer を上回る性能を KV 圧縮率 50% で達成。

4. 実験結果

TinyLlama (1.1B, 22 層) をベースに、SlimPajama データセットで事前学習を行い、各種ベンチマークで評価しました。

• 推論効率:
  • プリフィリング遅延: 標準 Transformer の約半分（YOCO と同等）。
  • デコーディングスループット: 標準 Transformer より高いメモリ効率により、バッチサイズを大きく取れるため、YOCO と同等の高性能を維持。
  • FusedKV は I/O オーバーヘッドによりスループットが低下しましたが、YOCO++ は低下しませんでした。
• モデル性能（ゼロショット精度）:
  • HellaSwag, OBQA, WinoGrande, ARC, PIQA などの常識推論タスクにおいて、YOCO++ は標準 Transformer を上回る結果を達成しました。
  • 平均スコアは 48.99（標準 Transformer: 48.37, YOCO: 47.98）となり、クロスレイヤー KV 圧縮手法の中で SOTA（State-of-the-Art）を記録しました。
• アブレーション研究:
  • スケーリングファクター $\lambda$ を除くと、YOCO と同等の性能低下が見られ、学習の重要性が確認されました。
  • キー（Key）の残差接続を除去すると性能が低下し、Key と Value 両方の残差接続が重要であることが示されました。

5. 結論と意義

YOCO++ は、LLM の推論における KV キャッシュのメモリボトルネックを解決する際、「メモリ削減」と「モデル性能」のトレードオフを打破する画期的な手法です。

• 技術的意義: KV 行列自体に残差接続を導入するという新しい視点を提供し、クロスレイヤー共有の弱点を補完する有効なメカニズムを確立しました。
• 実用性: 追加のハードウェア要件や I/O オーバーヘッドを増やすことなく、既存の YOCO 構造をアップグレード可能であり、大規模モデルの効率的なデプロイに大きく貢献します。

本論文は、50% の KV キャッシュ圧縮率において、標準的な Transformer よりも優れた性能と効率を両立させることを実証し、今後の効率的 LLM 推論研究の重要な指針となっています。