APE: 適応並列エンコーディングを介した高速かつ長いコンテキスト拡張生成

APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding

February 8, 2025

著者: Xinyu Yang, Tianqi Chen, Beidi Chen

cs.AI

要旨

コンテキスト拡張生成（CAG）技術、RAGおよびICLを含む、ユーザークエリに対する応答を生成するために複数のコンテキストを効率的に組み合わせる必要があります。これらのコンテキストを直接シーケンスとして入力することは、各リクエストごとに組み合わせられた複数のコンテキストを再エンコードすることにより、かなりの計算負荷を導入します。この課題に対処するために、私たちは並列エンコーディングの有望な可能性を探求し、各コンテキストのKV状態を独立して事前計算およびキャッシュすることで、直接キャッシュされた状態を推論中にロードすることを可能にし、位置の再利用を通じてより多くのコンテキストを収容します。ただし、注意分布の不整合により、直接並列エンコーディングを適用すると、大幅な性能低下が発生します。効果的かつ効率的なCAGを実現するために、Adaptive Parallel Encoding（APE）を提案します。これにより、共有プレフィックス、注意温度、およびスケーリングファクターが導入され、並列エンコーディングの分布を直列エンコーディングと整合させます。RAGおよびICLタスクの結果は、APEが同じ入力を使用してシーケンシャルエンコーディングの性能を98％および93％維持しながら、それぞれ3.6％および7.9％の性能向上を達成することを示しています。また、APEは多数のコンテキストを効果的に並列エンコードし、エンドツーエンドで128K長のコンテキストに対して28倍のプリフィリング時間を削減することで、4.5倍の高速化を実現できることを効率評価で示しています。

English

Context-augmented generation (CAG) techniques, including RAG and ICL, require the efficient combination of multiple contexts to generate responses to user queries. Directly inputting these contexts as a sequence introduces a considerable computational burden by re-encoding the combined selection of contexts for every request. To address this, we explore the promising potential of parallel encoding to independently pre-compute and cache each context's KV states. This approach enables the direct loading of cached states during inference while accommodating more contexts through position reuse across contexts. However, due to misalignments in attention distribution, directly applying parallel encoding results in a significant performance drop. To enable effective and efficient CAG, we propose Adaptive Parallel Encoding (APE), which brings shared prefix, attention temperature, and scaling factor to align the distribution of parallel encoding with sequential encoding. Results on RAG and ICL tasks demonstrate that APE can preserve 98% and 93% sequential encoding performance using the same inputs while outperforming parallel encoding by 3.6% and 7.9%, respectively. It also scales to many-shot CAG, effectively encoding hundreds of contexts in parallel. Efficiency evaluation shows that APE can achieve an end-to-end 4.5times speedup by reducing 28times prefilling time for a 128K-length context.