Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

EmbeddingRWKV: State-Centric Retrieval with Reusable States

Haowen Hou, Jie Yang

TL;DR

This work addresses inefficiency in retrieval-augmented generation by introducing State-Centric Retrieval, which unifies embedding and reranking through reusable matrix-valued states. It introduces EmbeddingRWKV, a RWKV-based backbone that produces both dense embeddings and compact states, and a state-based RWKV-Reranker that reuses precomputed document states to decouple inference cost from document length. Through a domain-aware curriculum and a linear-RNN architecture, the method delivers strong retrieval quality while achieving substantial efficiency gains, including offline speedups of $5.4\times$--$44.8\times$ and memory reductions, with the state cache requiring roughly $32/T$ of Transformer KV memory. Ablation studies show that state representations are essential for reranking, that layer-wise state diversity matters, and that caching KV for Transformer rerankers is impractical at scale. Overall, state-centric modeling provides a practical path to more scalable and cost-efficient retrieval systems and suggests broader applicability to state-based components in intelligent agents.

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

EmbeddingRWKV: State-Centric Retrieval with Reusable States

TL;DR

This work addresses inefficiency in retrieval-augmented generation by introducing State-Centric Retrieval, which unifies embedding and reranking through reusable matrix-valued states. It introduces EmbeddingRWKV, a RWKV-based backbone that produces both dense embeddings and compact states, and a state-based RWKV-Reranker that reuses precomputed document states to decouple inference cost from document length. Through a domain-aware curriculum and a linear-RNN architecture, the method delivers strong retrieval quality while achieving substantial efficiency gains, including offline speedups of -- and memory reductions, with the state cache requiring roughly of Transformer KV memory. Ablation studies show that state representations are essential for reranking, that layer-wise state diversity matters, and that caching KV for Transformer rerankers is impractical at scale. Overall, state-centric modeling provides a practical path to more scalable and cost-efficient retrieval systems and suggests broader applicability to state-based components in intelligent agents.

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--44.8 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}.
Paper Structure (46 sections, 5 equations, 3 figures, 12 tables)

This paper contains 46 sections, 5 equations, 3 figures, 12 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Two-Stage Retrieval for Text Embeddings.
  4. Multi-Stage Training Recipes.
  5. From Transformers to Linear RNNs.
  6. Linear RNNs with Matrix-valued States.
  7. State-Centric Retrieval
  8. Background: Matrix-valued States
  9. State Representation Learning
  10. Model Architecture.
  11. Training and Data Recipe.
  12. Learning Objective.
  13. State-Based Reranking
  14. Model Architecture.
  15. Reranking Training Objective.
  16. ...and 31 more sections

Figures (3)

  • Figure 1: Traditional vs. State-Centric Retrieval. (Top) The traditional two-stage retrieval is fundamentally disjointed, suffering from redundant computation as the reranker re-encodes full document tokens. (Bottom) Our State-Centric Retrieval unifies the two stages into an efficient system via a shared, reusable state. By jointly generating embeddings and compact states, it enables offline state-based reranking that decouples inference cost from document length, yielding 5.4$\times$--44.8$\times$ speedup.
  • Figure 2: (a) State Representation Learning: illustrates the process of "Learning the State." (b) State-Based Reranking: demonstrates the paradigm of "Utilizing the State" by directly reusing the cached state for inference, thereby avoiding redundant re-computation.
  • Figure 3: Performance benchmarking across various sequence lengths. Baseline: mxbai-rerank-large-v2 (1.5B) li2025prorank with FlashAttention-2 vs. Ours: EmbeddingRWKV-1.4B + RWKV-Reranker-1.3B. Top: Throughput comparison (pair/s), where higher is better. Bottom: Peak VRAM usage (GB), where lower is better.