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

TL;DR

This paper introduces SRAS, a lightweight reinforcement learning-based document selector designed for edge-native RAG systems, achieving high accuracy with minimal latency and resource usage.

Contribution

SRAS is the first ultra-lightweight RL-based document selector optimized for low-latency, resource-constrained edge environments, combining hybrid rewards for improved performance.

Findings

SRAS maintains <1s latency on CPU with a 0.76MB model size.

Outperforms supervised and random selectors on synthetic QA benchmarks.

Achieves BERTScore F1 of 0.8546 on SQuAD v2 without domain-specific tuning.

Abstract

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.