EmbeddingRWKV: State-Centric Retrieval with Reusable States

TL;DR

EmbeddingRWKV introduces a unified, state-centric retrieval approach that leverages reusable states from an RWKV-based language model, significantly improving efficiency and maintaining high retrieval quality in RAG systems.

Contribution

The paper proposes a novel state-centric retrieval paradigm with a unified model, EmbeddingRWKV, enabling efficient retrieval and reranking by reusing compact states and reducing redundant computation.

Findings

Achieves 5.4x to 44.8x speedup in reranking.

Maintains 98.62% of full-model performance with only 25% of layers.

Demonstrates high-quality retrieval and reranking results.

Abstract

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}.