Oussama Gabouj

Oussama Gabouj

Co-founder & CTO @ Compresr (YC W26) | Efficient LLMs & context compression | EPFL MSc Data Science

GRAD: Generative Retrieval-Aligned Demonstration Sampler for Efficient Few-Shot Reasoning

Published in Findings of EMNLP 2025, 2025

Abstract

Large Language Models (LLMs) achieve strong performance across diverse tasks, but their effectiveness often depends on the quality of the provided context. Retrieval-Augmented Generation (RAG) enriches prompts with external information, but its reliance on static databases constrains adaptability and can result in irrelevant demonstrations. In this work, we propose a Generative Retrieval-Aligned Demonstrator (GRAD), a dynamic demonstration-based approach where an LLM model is trained to generate input-specific concise demonstrations. By tailoring demonstrations to each input, our method offers better contextual support than traditional RAG approaches.

We demonstrate the superiority of GRAD under budget constraints, where we limit both the number of tokens used per demonstration and the number of tokens used for the final output. Trained solely on a math dataset, GRAD consistently outperforms strong baselines on Qwen2.5-14B across mathematical reasoning and advanced STEM questions, highlighting GRAD’s robust generalization to out-of-distribution (OOD) domains such as physics, chemistry, and computer science. Furthermore, we show that demonstrations generated by trained smaller models can effectively guide larger target models, reducing training costs while maintaining competitive accuracy.

Overall, this work introduces a scalable demonstration generator model presenting the first step toward a dynamic few-shot learning paradigm in resource-constrained settings. We release the code used for the project.

Key Contributions

  • Dynamic Demonstrations: GRAD generates input-specific demonstrations tailored to each query, offering better contextual support than static RAG approaches.
  • Budget-Constrained Performance: Demonstrates superior performance under token constraints for both demonstrations and final outputs.
  • Strong OOD Generalization: Trained on math datasets, GRAD generalizes effectively to physics, chemistry, and computer science domains.
  • Cross-Model Transfer: Smaller trained models can guide larger target models, reducing training costs while maintaining accuracy.
  • Scalable Few-Shot Learning: Introduces a practical framework for dynamic few-shot learning in resource-constrained settings.

Method Notes

Behind the headline result, a few decisions mattered more than the rest:

  • Reward design. The reward combines answer log-probability, accuracy and a demonstration count, with an explicit cap so the model cannot game it by padding demonstrations.
  • PPO vs GRPO. We compared both and kept GRPO for training stability.
  • Memory budget. Training used LoRA to fit inside the memory we had available.
  • GRADi. An SFT warm-started variant of the same recipe.

Publication Details