微分可能なキャッシュ拡張を介した潜在空間における熟考

Deliberation in Latent Space via Differentiable Cache Augmentation

December 23, 2024

著者: Luyang Liu, Jonas Pfeiffer, Jiaxing Wu, Jun Xie, Arthur Szlam

cs.AI

要旨

大規模言語モデル（LLM）が中間推論ステップを生成し、それに注意を払うことで「より考える」ことを可能にする技術は、複雑な問題の解決に有望性を示しています。しかしながら、標準的なアプローチは、応答の直前に離散トークンのシーケンスを生成するため、著しい遅延コストが発生し、最適化が難しいことがあります。本研究では、凍結されたLLMにオフラインコプロセッサを追加することで、モデルのキー・バリュー（kv）キャッシュ上で動作することができることを示します。このコプロセッサは、キャッシュに一連の潜在的な埋め込みを追加し、後続のデコーディングの忠実度を向上させるよう設計されています。我々は、このコプロセッサを、デコーダーの標準事前トレーニングデータ上での言語モデリング損失を使用してトレーニングし、デコーダー自体は凍結したままにします。このアプローチにより、モデルはエンドツーエンドで微分可能な方法で、kvキャッシュに追加の計算を蒸留する方法を学習することができます。デコーダーが変更されていないため、コプロセッサはオフラインおよび非同期で動作し、コプロセッサが利用できない場合や特定のキャッシュが追加の計算を必要としない場合でも、言語モデルは通常通り機能することができます。実験的に示されたように、キャッシュが拡張されると、デコーダーは多数の後続トークンでより低いパープレキシティを達成します。さらに、タスク固有のトレーニングがなくても、実験は、キャッシュの拡張が一貫してパープレキシティを低下させ、推論集約的なタスクの性能を向上させることを示しています。

English

Techniques enabling large language models (LLMs) to "think more" by generating and attending to intermediate reasoning steps have shown promise in solving complex problems. However, the standard approaches generate sequences of discrete tokens immediately before responding, and so they can incur significant latency costs and be challenging to optimize. In this work, we demonstrate that a frozen LLM can be augmented with an offline coprocessor that operates on the model's key-value (kv) cache. This coprocessor augments the cache with a set of latent embeddings designed to improve the fidelity of subsequent decoding. We train this coprocessor using the language modeling loss from the decoder on standard pretraining data, while keeping the decoder itself frozen. This approach enables the model to learn, in an end-to-end differentiable fashion, how to distill additional computation into its kv-cache. Because the decoder remains unchanged, the coprocessor can operate offline and asynchronously, and the language model can function normally if the coprocessor is unavailable or if a given cache is deemed not to require extra computation. We show experimentally that when a cache is augmented, the decoder achieves lower perplexity on numerous subsequent tokens. Furthermore, even without any task-specific training, our experiments demonstrate that cache augmentation consistently reduces perplexity and improves performance across a range of reasoning-intensive tasks.