arXiv:2601.07861cs.CLcs.IR2026-01

用可复用的状态统一检索与重排序,速度提升5.4到44.8倍。

EmbeddingRWKV: State-Centric Retrieval with Reusable States

  • 以状态为桥梁,将嵌入模型与重排序器融合为统一架构。
  • 重排序仅处理查询词元,推理速度提升5.4–44.8倍。
  • 仅用25%层即可保持98.62%性能,适合高效率检索场景。

当前检索增强生成(RAG)系统通常采用两阶段流水线:先用嵌入模型初步检索,再通过重排序器优化结果。然而,该范式因阶段间缺乏信息共享,导致大量重复计算,效率低下。为此,我们提出「状态中心检索」,利用「状态」作为连接嵌入模型与重排序器的桥梁。首先,通过微调基于RWKV的LLM,将其转化为可同时担任嵌入模型与状态主干的EmbeddingRWKV。在此基础上,设计基于状态的重排序器,充分复用预计算状态。重排序时,模型仅处理查询词元,实现推理成本与文档长度解耦,速度提升5.4×至44.8×。此外,发现保留所有中间层状态非必要;采用均匀层选择策略,仅用25%层即可维持98.62%的完整模型性能。大量实验表明,该方法在保证高质量检索与重排序的同时,显著提升系统整体效率。代码已开源。

原文摘要 · Abstract (English)

Current Retrieval-Augmented Generation (RAG) systems typically employ a traditional two-stage pipeline: an embedding model for initial retrieval followed by a reranker for refinement. However, this paradigm suffers from significant inefficiency due to the lack of shared information between stages, leading to substantial redundant computation. To address this limitation, we propose \textbf{State-Centric Retrieval}, a unified retrieval paradigm that utilizes "states" as a bridge to connect embedding models and rerankers. First, we perform state representation learning by fine-tuning an RWKV-based LLM, transforming it into \textbf{EmbeddingRWKV}, a unified model that serves as both an embedding model and a state backbone for extracting compact, reusable states. Building upon these reusable states, we further design a state-based reranker to fully leverage precomputed information. During reranking, the model processes only query tokens, decoupling inference cost from document length and yielding a 5.4$\times$--44.8$\times$ speedup. Furthermore, we observe that retaining all intermediate layer states is unnecessary; with a uniform layer selection strategy, our model maintains 98.62\% of full-model performance using only 25\% of the layers. Extensive experiments demonstrate that State-Centric Retrieval achieves high-quality retrieval and reranking results while significantly enhancing overall system efficiency. Code is available at \href{https://github.com/howard-hou/EmbeddingRWKV}{our GitHub repository}.

检索增强高效推理状态机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。