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

Rajiv Chaitanya Muttur

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

Information Retrieval 2026-01-06 v1 Machine Learning

Authors: Rajiv Chaitanya Muttur

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.

Keywords

retrieval-augmented generation reinforcement learning sequence alignment

Cite

@article{arxiv.2601.01785,
  title  = {SRAS: A Lightweight Reinforcement Learning-based Document Selector for Edge-Native RAG Pipelines},
  author = {Rajiv Chaitanya Muttur},
  journal= {arXiv preprint arXiv:2601.01785},
  year   = {2026}
}

Comments

Presented at ICEdge 2025; nominated for Best Paper Award

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

Abstract

Keywords

Cite

Comments

Related papers