効率的な長文脈生成のためのコンテキスト記憶化

要旨

近年、大規模言語モデル（LLM）の応用において、推論時にモデルの振る舞いを制御するために、長い条件付けプレフィックスに依存するケースが増加している。プレフィックス拡張推論は効果的であるものの、二つの構造的制約を伴う。すなわち、(i) プレフィックスの影響は生成が進むにつれて弱まり、(ii) プレフィックスに対する注意計算のコストはその長さに比例して増大する。既存のアプローチでは、プレフィックスを圧縮して注意機構内に保持するか、勾配ベースの学習を通じてモデルパラメータに内部化する。前者は依然として推論時にプレフィックスに注意を払う必要があり、後者は学習負荷が高く、プレフィックスの更新に適さない。これらの問題に対処するため、我々は注意状態メモリを提案する。これは、プレフィックスとクエリトークン間の注意状態を事前計算した軽量な参照ベースメモリにプレフィックスを外部化する、訓練不要の手法である。LLaMA-3.1-8Bを用いたManyICLBenchにおいて、本手法は1K～8Kのメモリ予算下で文脈内学習を上回る精度を達成し、8Kでは注意レイテンシを1.36倍削減する。また、NBAベンチマークでは、フルアテンションRAGの性能を、そのメモリフットプリントのわずか20%で凌駕する。

English

Modern large language model (LLM) applications increasingly rely on long conditioning prefixes to control model behavior at inference time. While prefix-augmented inference is effective, it incurs two structural limitations: i) the prefix's influence fades as generation proceeds, and ii) attention computation over the prefix scales linearly with its length. Existing approaches either keep the prefix in attention while compressing it, or internalize it into model parameters through gradient-based training. The former still attends to the prefix at inference, while the latter is training-intensive and ill-suited to prefix updates. To address these issues, we propose attention-state memory, a training-free approach that externalizes the prefix into a lightweight, lookup-based memory of precomputed attention states between prefix and query tokens. On ManyICLBench with LLaMA-3.1-8B, our method improves accuracy over in-context learning at 1K-8K memory budgets while reducing attention latency by 1.36x at 8K, and surpasses full-attention RAG performance on NBA benchmark using only 20% of its memory footprint.