SRAS: A Lightweight Reinforcement Learning-based Document Selector for Edge-Native RAG Pipelines
This work addresses computational overheads for edge-native RAG pipelines, offering an incremental improvement by making RL-based selection lightweight and latency-aware.
The paper tackled the problem of inefficient document selection in Retrieval-Augmented Generation (RAG) systems by proposing SRAS, a lightweight reinforcement learning-based selector that achieves a BERTScore F1 of 0.8546 on SQuAD v2 with under 1s latency and a compact 0.76MB policy.
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.