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

arXiv: https://arxiv.org/abs/2510.01165 · HTML: https://mlsanity.com/paper/2510.01165/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 22 of its 36 references that are in the map.
- First posted on arXiv 1 Oct 2025.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)
- Retrieval-augmented generation / [Agentic RAG, agents and memory](https://mlsanity.com/rag/agentic-rag-agents-and-memory/)

## Builds on (most cited first, 22 of 22)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Self-Instruct: Aligning Language Models with Self-Generated Instructions](https://mlsanity.com/paper/2212.10560/) (2022)
- [Retrieval-Augmented Generation for Large Language Models: A Survey](https://mlsanity.com/paper/2312.10997/) (2023)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [What Makes Good In-Context Examples for GPT-3?](https://mlsanity.com/paper/2101.06804/) (2021)
- [Automatic Chain of Thought Prompting in Large Language Models](https://mlsanity.com/paper/2210.03493/) (2022)
- [MathQA: Towards Interpretable Math Word Problem Solving with Operation-Based Formalisms](https://mlsanity.com/paper/1905.13319/) (2019)
- [Larger language models do in-context learning differently](https://mlsanity.com/paper/2303.03846/) (2023)
- [Reinforcement Learning for Reasoning in Large Language Models with One Training Example](https://mlsanity.com/paper/2504.20571/) (2025)
- [Scaling Sentence Embeddings with Large Language Models](https://mlsanity.com/paper/2307.16645/) (2023)
- [GRPO-LEAD: A Difficulty-Aware Reinforcement Learning Approach for Concise Mathematical Reasoning in Language Models](https://mlsanity.com/paper/2504.09696/) (2025)
- [How Many Demonstrations Do You Need for In-context Learning?](https://mlsanity.com/paper/2303.08119/) (2023)
- [Fewer is More: Boosting LLM Reasoning with Reinforced Context Pruning](https://mlsanity.com/paper/2312.08901/) (2023)
- [Adaptive Group Policy Optimization: Towards Stable Training and Token-Efficient Reasoning](https://mlsanity.com/paper/2503.15952/) (2025)
- [Measuring Sentences Similarity: A Survey](https://mlsanity.com/paper/1910.03940/) (2019)
- [Assessing Out-of-Domain Language Model Performance from Few Examples](https://mlsanity.com/paper/2210.06725/) (2022)
- [Auto-Demo Prompting: Leveraging Generated Outputs as Demonstrations for Enhanced Batch Prompting](https://mlsanity.com/paper/2410.01724/) (2024)
- [DoubleDipper: Improving Long-Context LLMs via Context Recycling](https://mlsanity.com/paper/2406.13632/) (2024)
- [Post-training an LLM for RAG? Train on Self-Generated Demonstrations](https://mlsanity.com/paper/2502.10596/) (2025)

## Abstract (arXiv)

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.