推理加速:RASD 和 REFRAG 深度解析
Update: 2025-09-10
Description
REFRAG (REpresentation For RAG) 和 RASD (Retrieval-Augmented Speculative Decoding) 都是旨在提高大型语言模型 (LLM) 推理效率的方法,但它们关注的方面和实现机制有所不同:
- 核心目标和解决的问题:REFRAG:主要目标是解决RAG应用中长上下文输入带来的显著系统延迟和对键值 (KV) 缓存的大量内存需求,从而提高吞吐量并解决知识丰富与系统效率之间的基本权衡问题。它特别关注首个token生成时间 (TTFT) 的优化。REFRAG提出,RAG上下文通常信息稀疏,且检索到的段落之间的交叉注意力模式呈块对角线状,因此大部分计算是不必要的。
RASD:旨在通过引入检索方法来增强基于模型的推测解码 (speculative decoding),从而加速LLM的推理过程。推测解码通过生成草稿token供目标模型验证,显著提高了LLM推理的计算并行性。RASD的核心在于解决传统草稿模型在域外场景中效率低下以及草稿阶段耗时导致验证步长上限低的问题。 - 实现机制:
REFRAG:通过压缩 (Compress)、感知 (Sense) 和扩展 (Expand) 上下文表示来优化效率。压缩:不直接使用检索到的段落中的所有token作为输入,而是利用预先计算的、压缩的块嵌入 (compressed chunk embeddings) 作为近似表示,并直接将这些嵌入馈送给解码器。这缩短了解码器输入长度,重用了检索过程中的计算,并将注意力计算复杂度从与token数量的平方关系降低到与块数量的平方关系。
感知/选择性压缩:通过一个轻量级的强化学习 (RL) 策略来选择性地确定何时需要完整的块token输入,以及何时低成本的近似块嵌入就足够。RL策略以“下一段落预测困惑度”作为负奖励来指导,以决定哪些块应以原始形式保留。
扩展:引入选择性token压缩功能,通过将重要的上下文块未压缩地扩展来改进答案预测,而将不那么关键的片段进行压缩。
架构:由一个仅解码器基础模型(如LLaMA)和一个轻量级编码器模型(如Roberta)组成。
RASD:通过检索过程 (Retrieval Process)、检索树剪枝 (Retrieval Tree Pruning) 和树融合 (Tree Fusion) 三个主要步骤实现。检索过程:利用检索方法(如PLD或REST)根据当前上下文s'作为查询,从数据存储D中识别延续候选S。PLD通过后缀匹配从输入中检索,REST则从预定义上下文数据库中检索相关token。
检索树剪枝:将检索到的候选序列构建成检索树,并根据草稿模型的输出分布(置信度分数)对其进行剪枝,以选择高质量的检索结果。
树融合:将草稿模型生成的树与检索树进行融合。通过最长前缀匹配算法合并共享相同前缀的分支,创建一个统一的“检索增强草稿token树”,融合了语言模型和知识库的信息。
草稿树验证:目标模型使用树注意力 (tree attention) 在一次前向传递中验证融合后的草稿token树,从而实现加速。RASD可以增强任何使用树注意力的推测解码方法(如EAGLE2)。 - 主要区别总结:
关注点:REFRAG关注的是RAG上下文的表示和处理方式,通过压缩上下文信息来减少计算量和内存占用,特别优化TTFT。
RASD关注的是推测解码中的草稿token生成质量和效率,通过引入外部检索来辅助生成更准确的草稿序列,减少目标模型的验证步骤。
上下文处理方式:REFRAG将整个RAG上下文分块并压缩成嵌入,并通过RL策略动态决定哪些块需要“解压缩”以保留完整信息。
RASD通过检索获取上下文延续的候选序列,这些序列被组织成树形结构,并与草稿模型的生成结果融合,用于推测解码的验证。
核心技术:REFRAG通过改变上下文的输入表示(从token到块嵌入) 来达到效率提升,并利用强化学习进行智能压缩。
RASD通过增强推测解码的草稿生成阶段来实现加速,它依赖于检索算法、树结构操作(剪枝、融合)和树注意力验证。
应用场景侧重:REFRAG在多种长上下文任务中表现出色,包括RAG、多轮对话和长文档摘要,特别是在相同延迟预算下可以处理更多信息。
RASD在文档问答 (DocQA)、摘要 (Summary)、代码生成 (Code) 和域内问答 (In-Domain QA) 等任务中加速推理,并且在知识密集型或草稿模型在域外场景表现不佳时特别有效。
简而言之,REFRAG通过优化RAG上下文的**“输入形式”** 来提高效率,而RASD通过优化推测解码中**“草稿token的生成和验证”** 来加速推理。
Papers:
Comments
In Channel