💬 NLP

RIMS: Preference Optimization via Smoothed Multi-pair Aggregation for Small-Scale LLM Retrieval-Augmented Generation

The paper proposes RIMS, a three-stage preference optimization framework for small-scale language models in retrieval-augmented generation that utilizes self-generated synthetic data and a differentiable soft aggregation mechanism to effectively leverage multiple preference pairs, thereby outperforming state-of-the-art baselines in noisy retrieval conditions.

Original authors: Pei Tian, Zihan Dong, Tianci Liu, Linjun Zhang, Haoyu Wang

Published 2026-07-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Pei Tian, Zihan Dong, Tianci Liu, Linjun Zhang, Haoyu Wang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but tiny, robot how to answer tricky questions. This robot is a "Small-Scale Language Model" (SLM). Think of it like a brilliant high school student who has read a few books but hasn't memorized the entire library. When you ask it a hard question, it gets nervous and might make things up. To help, you give it a "cheat sheet" of search results from the internet, a technique called Retrieval-Augmented Generation (RAG).

However, there's a catch. The internet is messy. Your cheat sheet might contain a mix of true facts, confusing rumors, and outright lies. A big, super-powerful robot (a Large Language Model) might be able to spot the lies and ignore them. But our tiny robot? It often gets confused by the noise, believing the wrong information and giving a wrong answer. Scientists have tried to fix this by teaching the robot to choose between different answers, a method called "preference optimization." But the old way of teaching was like a strict coach who only looked at the single worst mistake the robot made and ignored everything else. It turns out, ignoring the other clues makes the robot even more confused.

This paper introduces a new training method called RIMS (Retrieval-Augmented Generation via Smoothed Multi-pair Aggregation). Instead of acting like a strict coach who picks just one "worst" example to fix, RIMS acts like a wise mentor who looks at all the examples the robot generated. It gently blends the good and bad clues together into a single, smooth lesson. This helps the tiny robot learn to spot the noise in its cheat sheet much better, even when the information is messy. The researchers tested this on four different difficult question-answering games and found that their method helped the small robots get significantly more answers right than the old, strict methods.

The Problem: The Tiny Robot and the Noisy Library

Let's dive into the story. We have these small, efficient AI models (SLMs) that are great for running on phones or small computers because they don't need massive amounts of power. But they have a small "brain capacity." They don't know as much as the giant models. To fix this, we connect them to a search engine (RAG) so they can look up facts.

The problem is that search engines aren't perfect. Sometimes they bring back documents that are totally irrelevant or even contradictory. When a tiny AI tries to read through a pile of documents that includes both the truth and a bunch of nonsense, it often gets overwhelmed. It might grab the wrong fact and confidently state it as truth.

The Old Way: The "Hardest Mistake" Coach

To teach these AIs to be better, researchers use a technique called preference optimization. Imagine you show the AI ten different answers to a question. Some are good, some are bad. The goal is to teach the AI to prefer the good ones and reject the bad ones.

Previously, a popular method (called RoseRAG) acted like a very strict coach. When the AI generated ten answers, the coach would look at them and say, "Okay, I see one answer that is really bad and one that is really good. I am going to ignore the other eight. We will only train on this single pair of 'best' and 'worst'."

The paper argues that this is a waste of information. By throwing away the other eight answers, the coach is ignoring valuable clues about why the other answers were better or worse. It's like a student taking a test, getting a question wrong, and the teacher only correcting that one specific mistake while ignoring the fact that the student also got three other questions slightly wrong. The student misses the bigger picture.

The New Way: RIMS and the "Smooth Blend"

The authors propose RIMS, which changes the coaching style completely. Instead of picking just one "winner" and one "loser," RIMS looks at all the answers the AI generated.

Here is the magic trick: RIMS uses a mathematical tool called LogSumExp (LSE) to create a "smooth blend." Imagine you have a bowl of soup with ten different ingredients (the ten answers). The old method would scoop out just the saltiest spoonful and the blandest spoonful and throw the rest away. RIMS, however, takes the whole bowl and stirs it together into a perfect, balanced flavor.

This "smooth" mixture creates a single, unified lesson that contains the signal from every answer. It doesn't ignore the easy ones or the hard ones; it weighs them all together. This is called "differentiable soft aggregation." It's "differentiable" because the math allows the AI to learn from the blend smoothly, and "soft" because it doesn't make a harsh, all-or-nothing decision about which answer to focus on.

What They Found

The researchers tested RIMS on four different "multi-hop" question-answering benchmarks. These are like trivia games where you have to connect several pieces of information to get the answer, and the search results are often full of distractions. They tested it on two different small AI models (Qwen2.5-1.5B and Gemma-2-2b).

The results were clear:

Better Scores: RIMS consistently beat the best existing methods (like RoseRAG) in getting the exact right answer (Exact Match) and in matching the right words (F1 score).
Handling Noise: When the researchers added more and more "noise" (irrelevant documents) to the search results, the old methods started to fail. RIMS, however, kept performing well. It proved to be much more robust against messy information.
Theoretical Proof: The authors didn't just guess this would work; they did the math. They proved that their "smooth" method reduces the "variance" (or randomness) in the learning process. In simple terms, the old method was like a shaky hand trying to draw a straight line, while RIMS is a steady hand. They also showed that the "error" in their smoothing method can be controlled and kept very small.

Why It Matters

This paper suggests that for small, efficient AI models, we shouldn't throw away data. Even when an AI generates a bunch of answers, every single one of them holds a clue. By using a "smooth" way to combine all those clues, we can teach these tiny robots to be much smarter and more reliable, even when the information they find is messy. It's a step toward making powerful AI accessible on everyday devices without needing a supercomputer.

Technical Summary: RIMS – Preference Optimization via Smoothed Multi-pair Aggregation for Small-Scale LLM Retrieval-Augmented Generation

Problem Statement

Small-scale language models (SLMs) are increasingly vital for resource-constrained Retrieval-Augmented Generation (RAG) due to their lower inference costs and portability. However, their limited capacity makes them highly sensitive to noisy or spurious retrieved evidence, often leading to hallucinated outputs or an inability to distinguish relevant information from distractors.

Existing methods for improving SLM robustness in RAG settings face significant limitations:

Prompt-based approaches rely on the model's intrinsic ability to filter noise, which often fails for SLMs.
Supervised Fine-Tuning (SFT) is highly sensitive to noise in retrieved contexts and prone to overfitting on imperfect evidence.
Preference Optimization (PO) methods offer a more robust alternative by learning from relative quality signals. However, current state-of-the-art approaches like RoseRAG utilize multi-pair preference data inefficiently. RoseRAG selects only the single "hardest" preference pair (via hard argmin/argmax) based on initial model likelihoods and discards all other signals. This hard selection fixes the training signal early, ignoring the evolving policy and the rich information contained in the remaining candidate pairs. Other methods, such as NCA, aggregate negative responses but fail to jointly model positive candidates, resulting in incomplete utilization of the preference structure.

Methodology: RIMS

The authors propose RIMS (Retrieval-Augmented Generation via Smoothed Multi-pair Aggregation), a three-stage framework designed to fully leverage multi-pair preference data for SLMs without relying on proprietary models.

1. Multi-Pair Preference Data Generation

RIMS generates preference data using the target SLM itself. Through rejection sampling with Chain-of-Thought (CoT) prompting, the model generates multiple reasoning rationales. Responses matching the ground truth are retained as preferred ( $y_w$ ), while vanilla RAG outputs (without ground-truth guidance) serve as dispreferred ( $y_l$ ) candidates. This process yields multiple $(y_w, y_l)$ pairs per query, creating a rich dataset of multi-pair preferences.

2. Differentiable Soft Aggregation

The core innovation of RIMS is replacing the non-differentiable hard selection (argmin/argmax) with a differentiable soft aggregation mechanism based on the LogSumExp (LSE) function.

Instead of selecting a single hardest pair, RIMS aggregates all preferred candidates into a unified pseudo-preferred response ( $\tilde{y}_w$ ) and all dispreferred candidates into a pseudo-dispreferred response ( $\tilde{y}_l$ ).
The aggregation uses a temperature parameter $\tau$ $τ$ controlled by an error bound $M$ $M$ .
- For preferred responses, it approximates the minimum likelihood (hardest positive) using negated LSE.
- For dispreferred responses, it approximates the maximum likelihood (hardest negative) using standard LSE.
This operator is differentiable, allowing gradient signals to propagate from all candidate pairs simultaneously. Unlike RoseRAG, which fixes the pair selection once, RIMS's aggregation weights adapt dynamically as the policy evolves during training.

3. Preference Optimization

The smoothed objective is algorithm-agnostic. It can be applied to standard preference optimization algorithms (DPO, CPO, ORPO) by simply replacing the original policy likelihoods in the scoring function with the aggregated pseudo-likelihoods ( $\pi_\theta(\tilde{y}_w|x)$ and $\pi_\theta(\tilde{y}_l|x)$ ).

Theoretical Contributions

The paper provides two key theoretical guarantees:

Controllable Approximation Error: The authors prove that the LSE-based smoothing admits a controllable error bound relative to the exact maximum/minimum. They derive a relationship between the temperature $\tau$ , the top-gap $\delta$ between candidate scores, and the error bound $M$ , allowing practitioners to adaptively control the sharpness of the approximation.
Gradient Stability: Under candidate-pool sampling, the authors demonstrate that smooth aggregation yields provably tighter gradient alignment to the oracle objective than hard selection. Hard selection introduces an intrinsic variance floor due to the discrete randomness of pair selection, whereas smooth aggregation reduces this variance by utilizing a continuous weighted mixture of all candidates.

Experimental Results

RIMS was evaluated on four multi-hop question answering benchmarks (HotPotQA, 2WikiMultiHopQA, MultiHopRAGQA, and MuSiQueQA) using two SLM backbones: Qwen2.5-1.5B-Instruct and Gemma-2-2b-it.

Performance: RIMS consistently outperformed state-of-the-art baselines (including RoseRAG, InstructRAG, and various prompting strategies) across all datasets. The gains were particularly pronounced on challenging multi-hop datasets with noisy retrieval contexts, achieving consistent improvements in Exact Match (EM) and F1 scores.
Robustness to Noise: In experiments varying the number of retrieved documents ( $K$ ), RIMS maintained superior performance as $K$ increased and noise grew, whereas vanilla RAG performance degraded significantly.
Ablation Studies:
- Aggregation Strategy: RIMS (LSE smoothing) significantly outperformed alternative aggregation strategies like simple averaging or fixed-weighted averaging, confirming that the margin-aware structure is preserved while utilizing all candidates.
- Algorithm Agnosticism: The method proved effective when applied to DPO, CPO, and ORPO, with DPO yielding the highest performance.
- Hyperparameters: Sensitivity analysis showed that moderate smoothing sizes and an error bound $M=1$ offered the best balance between optimization smoothness and ranking fidelity.

Significance and Claims

The paper claims that RIMS addresses a critical gap in SLM-based RAG: the underutilization of multi-pair preference data. By moving from hard, static pair selection to a dynamic, differentiable aggregation of all available signals, RIMS enables small models to better distinguish relevant evidence from noise.

The authors emphasize that their approach does not require larger models or proprietary APIs for data generation, making it practical for resource-constrained environments. Theoretically, the work establishes that smooth aggregation provides a more stable training signal with lower variance than hard selection, offering a principled alternative to existing margin-aware methods. The results suggest that structured multi-pair aggregation is essential for robust evidence utilization in noisy retrieval conditions, allowing SLMs to achieve performance levels previously thought to require larger architectures or more complex inference-time strategies.