SRAS: A Lightweight Reinforcement Learning-based Document Selector for Edge-Native RAG Pipelines

📄 arXiv: 2601.01785v1 📥 PDF

作者: Rajiv Chaitanya Muttur

分类: cs.IR, cs.LG

发布日期: 2026-01-05

备注: Presented at ICEdge 2025; nominated for Best Paper Award

💡 一句话要点

提出SRAS：一种轻量级强化学习文档选择器，用于边缘原生RAG流水线。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 检索增强生成 强化学习 文档选择 边缘计算 近端策略优化 轻量级模型 自然语言处理

📋 核心要点

  1. 现有RAG系统依赖固定top-k选择，忽略下游生成质量，计算开销大，不适用于边缘设备。
  2. SRAS利用强化学习训练轻量级文档选择策略，结合Relaxed F1和BERTScore作为混合奖励。
  3. SRAS在CPU上延迟<1秒，在合成和真实数据集上均优于传统方法，无需领域特定调整。

📝 摘要（中文）

检索增强生成(RAG)系统通常依赖于固定的top-k文档选择机制，忽略了下游生成质量并带来了计算开销。我们提出了SRAS（稀疏奖励感知选择器），一种通过强化学习(RL)训练的轻量级文档选择器，用于边缘原生RAG部署。与之前假设大内存和延迟预算的基于RL的检索器不同，SRAS使用近端策略优化(PPO)学习了一个紧凑的策略（约0.76MB），并由结合了Relaxed F1和BERTScore的混合奖励信号引导。我们的方法在严格的token和计算约束下运行，在CPU上保持<1秒的延迟。SRAS在合成QA基准测试中优于监督和随机选择器，并推广到真实世界数据，在SQuAD v2上实现了0.8546的BERTScore F1，无需特定领域的调整。这项工作首次证明了基于RL的文档选择可以变得超轻量级、延迟感知且对设备上的RAG流水线有效。

🔬 方法详解

问题定义：RAG系统中文档检索模块通常采用固定的top-k策略，这种策略无法根据下游生成任务的实际需求进行优化，导致检索到的文档可能包含大量冗余信息，增加了计算负担，并且影响最终生成结果的质量。特别是在边缘设备上部署RAG系统时，计算资源和延迟要求更加严格，传统的top-k策略难以满足需求。

核心思路：SRAS的核心思路是利用强化学习(RL)训练一个轻量级的文档选择器，该选择器能够根据下游生成任务的反馈（奖励信号）动态地选择最相关的文档。通过学习一个紧凑的策略，SRAS能够在计算资源有限的边缘设备上高效地运行，并优化RAG系统的整体性能。

技术框架：SRAS的整体框架包括以下几个主要模块：1) 文档编码器：将文档编码成向量表示；2) 策略网络：根据查询和文档向量，输出选择每个文档的概率；3) 环境：模拟RAG系统的运行环境，接收选择的文档，并生成答案；4) 奖励函数：根据生成的答案质量，计算奖励信号，反馈给策略网络。SRAS使用近端策略优化(PPO)算法来训练策略网络。

关键创新：SRAS最重要的技术创新点在于其轻量级的设计和强化学习的训练方式。与传统的基于规则或监督学习的文档选择方法相比，SRAS能够根据下游任务的反馈进行自适应优化，并且其紧凑的模型大小使其非常适合在边缘设备上部署。此外，SRAS采用混合奖励信号，结合了Relaxed F1和BERTScore，能够更全面地评估生成答案的质量。

关键设计：SRAS的关键设计包括：1) 策略网络采用小型神经网络，参数量控制在0.76MB以内；2) 奖励函数采用Relaxed F1和BERTScore的加权平均，权重根据实验结果进行调整；3) 使用PPO算法进行训练，设置合适的学习率、折扣因子和裁剪参数；4) 在训练过程中，采用token和计算约束，模拟边缘设备的资源限制。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

SRAS在合成QA基准测试中优于监督和随机选择器，并在SQuAD v2数据集上取得了显著的性能提升，BERTScore F1达到0.8546，且无需任何领域特定的调整。该模型大小仅为0.76MB，在CPU上的延迟小于1秒，验证了其在边缘设备上部署的可行性和有效性。

🎯 应用场景

SRAS适用于各种需要在边缘设备上部署的RAG系统，例如智能助手、移动搜索、本地知识库问答等。它能够显著降低计算开销，提高响应速度，并提升生成结果的质量。该研究为边缘计算和人工智能的结合提供了新的思路，有望推动更多智能应用在资源受限的环境中落地。

📄 摘要（原文）

Retrieval-Augmented Generation (RAG) systems often rely on fixed top-k document selection mechanisms that ignore downstream generation quality and impose computational overheads. We propose SRAS (Sparse Reward-Aware Selector), a lightweight document selector trained via reinforcement learning (RL) for edge-native RAG deployment. Unlike prior RL-based retrievers that assume large memory and latency budgets, SRAS learns a compact (~0.76MB) policy using Proximal Policy Optimization (PPO), guided by a hybrid reward signal combining Relaxed F1 and BERTScore. Our method operates under tight token and compute constraints, maintaining <1s latency on CPU. SRAS outperforms supervised and random selectors on a synthetic QA benchmark, and generalizes to real-world data, achieving BERTScore F1 of 0.8546 on SQuAD v2 without domain-specific tuning. This work is the first to demonstrate that RL-based document selection can be made ultra-lightweight, latency-aware, and effective for on-device RAG pipelines.