提出APE方法,让大模型更快更长地处理多段上下文生成。
APE: Faster and Longer Context-Augmented Generation via Adaptive Parallel Encoding
- 通过自适应并行编码,预先缓存上下文键值状态
- 在128K上下文下实现4.5倍端到端加速,预填充时间减少28倍
- 支持数百个上下文并行处理,适合长文本问答与多示例学习
上下文增强生成(CAG)技术如RAG和ICL需高效融合多个上下文以响应用户查询。直接将上下文序列输入会因每次请求重复编码而带来巨大计算开销。为此,我们探索并行编码的潜力:独立预计算并缓存每个上下文的键值(KV)状态,推理时直接加载缓存状态,并通过跨上下文位置复用支持更多上下文。然而,由于注意力分布错位,直接使用并行编码导致性能显著下降。为此,我们提出自适应并行编码(APE),引入共享前缀、注意力温度和缩放因子,使并行编码的分布与串行编码对齐。在RAG和ICL任务上,APE使用相同输入可保持98%和93%的串行编码性能,分别优于纯并行编码3.6%和7.9%。同时可扩展至多示例CAG,有效并行编码数百个上下文。效率评估显示,对于128K长度上下文,APE可实现4.5倍端到端加速,预填充时间减少28倍。
原文摘要 · Abstract (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 ($\textbf{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.5$\times$ speedup by reducing 28$\times$ prefilling time for a 128K-length context.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。