Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Rajiv Chaitanya Muttur

TL;DR

SRAS introduces a lightweight, reinforcement-learning–based document selector for edge-native RAG pipelines, training a compact policy with PPO under sparse QA-derived rewards. The method uses a hybrid reward combining lexical and semantic signals, $R = \alpha \cdot \text{Relaxed-F1} + (1 - \alpha) \cdot \text{BERTScore}$ with $\alpha = 0.6$, and a compact architecture totaling ~197K parameters (~$0.76$ MB). Key contributions include the hybrid reward design, supervised warmup, curriculum learning, and extensive ablations showing stability and efficiency under CPU-only, latency-constrained conditions. Empirical results on synthetic data and real-world SQuAD v2 demonstrate competitive QA performance with sub-second latency and strong generalization, highlighting practical viability for on-device RAG. These findings advance adaptive, low-resource retrieval for edge-AI applications and open avenues for multi-domain and end-to-end integration in constrained environments.

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.

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

TL;DR

SRAS introduces a lightweight, reinforcement-learning–based document selector for edge-native RAG pipelines, training a compact policy with PPO under sparse QA-derived rewards. The method uses a hybrid reward combining lexical and semantic signals, with , and a compact architecture totaling ~197K parameters (~ MB). Key contributions include the hybrid reward design, supervised warmup, curriculum learning, and extensive ablations showing stability and efficiency under CPU-only, latency-constrained conditions. Empirical results on synthetic data and real-world SQuAD v2 demonstrate competitive QA performance with sub-second latency and strong generalization, highlighting practical viability for on-device RAG. These findings advance adaptive, low-resource retrieval for edge-AI applications and open avenues for multi-domain and end-to-end integration in constrained environments.

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.
Paper Structure (25 sections, 2 equations, 6 figures, 3 tables)

This paper contains 25 sections, 2 equations, 6 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Static and Learned Retrieval in RAG Pipelines
  4. Reinforcement Learning for Retrieval and Ranking
  5. Document Selection under Resource Constraints
  6. Methodology
  7. Pipeline Overview
  8. Hybrid Reward Signal: QA-derived Feedback
  9. SRAS Architecture
  10. Training with PPO under Sparse Rewards
  11. PPO Configuration
  12. Stabilization Techniques
  13. Experimental Setup
  14. Synthetic QA Dataset
  15. Candidate Document Pooling
  16. ...and 10 more sections

Figures (6)

  • Figure 1: End-to-end SRAS-enhanced RAG pipeline. The selector uses embeddings from both SBERT and the QA pair generator. Reward engine and PPO are used only during training.
  • Figure 2: SRAS scoring architecture. Question and document embeddings are projected to a shared space, combined via a $\tanh$ nonlinearity, and scored linearly.
  • Figure 3: SRAS achieves strong QA quality with low latency.
  • Figure 4: Ablation study showing the effect of removing supervised warmup (SW), reward shaping (RS), and curriculum learning (CL).
  • Figure 5: Average PPO reward per epoch. Removing RS stalls learning. SW and CL improve early training stability.
  • ...and 1 more figures