You Only Cache Once: Decoder-Decoder Architectures for Language Models

要約

大規模言語モデル向けに、キーと値のペアを 1 回だけキャッシュするデコーダ-デコーダ アーキテクチャ YOCO を導入します。
これは 2 つのコンポーネント、つまりセルフデコーダの上にスタックされたクロスデコーダで構成されます。
セルフデコーダは、クロスアテンションを介してクロスデコーダによって再利用されるグローバル キー/バリュー (KV) キャッシュを効率的にエンコードします。
YOCO は 1 回しかキャッシュしませんが、モデル全体はデコーダ専用の Transformer のように動作します。
この設計により、GPU メモリの需要が大幅に削減されますが、グローバル アテンション機能は維持されます。
さらに、計算フローにより、最終出力を変更せずにプレフィルから早期終了が可能になり、プレフィル ステージが大幅に高速化されます。
実験結果は、モデルのサイズとトレーニング トークンの数をスケールアップするさまざまな設定において、YOCO が Transformer と比較して良好なパフォーマンスを達成することを示しています。
また、ほぼ完璧な針回収精度で YOCO を 1M のコンテキスト長まで拡張します。
プロファイリングの結果は、YOCO がコンテキストの長さとモデルのサイズ全体で推論メモリ、プリフィル レイテンシ、スループットを桁違いに改善していることを示しています。
コードは https://aka.ms/YOCO で入手できます。

要約(オリジナル)

We introduce a decoder-decoder architecture, YOCO, for large language models, which only caches key-value pairs once. It consists of two components, i.e., a cross-decoder stacked upon a self-decoder. The self-decoder efficiently encodes global key-value (KV) caches that are reused by the cross-decoder via cross-attention. The overall model behaves like a decoder-only Transformer, although YOCO only caches once. The design substantially reduces GPU memory demands, yet retains global attention capability. Additionally, the computation flow enables prefilling to early exit without changing the final output, thereby significantly speeding up the prefill stage. Experimental results demonstrate that YOCO achieves favorable performance compared to Transformer in various settings of scaling up model size and number of training tokens. We also extend YOCO to 1M context length with near-perfect needle retrieval accuracy. The profiling results show that YOCO improves inference memory, prefill latency, and throughput by orders of magnitude across context lengths and model sizes. Code is available at https://aka.ms/YOCO.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google