GRAD : Generative Retrieval-Aligned Demonstration Sampler for Efficient Few-Shot Reasoning (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: https://github.com/charafkamel/GRAD-demonstration-sampler

## introduction large language models (llms) excel in multi-step reasoning tasks, often operating in a zero-shot setting. techniques such as context augmentation and * equal contribution. to address these challenges, we introduce grad: a generative retrieval-aligned demonstrator designed to combine the adaptability of generation with the efficiency of retrieval. grad dynamically generates task-specific demonstrations under a strict token budget, enabling better performance in both in-distribution (id) and ood tasks. to achieve this, we train grad with reinforcement learning (rl), which is the core mechanism to produce demonstrations, steering target models toward correct and concise generations. additionally, we consider a variant called gradi (grad initialized), where the generator is first initialized with supervised fine-tuning (sft). this auxiliary step provides initial guidance to the demonstrator, after which the training process is identical to the one used with grad. unless stated otherwise, we use grad to denote our method in general, and we will explicitly distinguish between grad and gradi in the results section. we evaluate our approach on diverse reasoning benchmarks (see section 3.6) using models ranging from 3b to 14b parameters and observe that grad variants consistently outperform rag across all evaluated datasets for our largest model. example 2 answer grad figure 1 : example input query from gsm8k with demonstrations and outputs from rag and grad. rag retrieves demonstrations from a static database, whereas grad generates task-specific demonstrations within a token budget. a final output length constraint is applied in both cases; grad produces shorter demonstrations and a more concise final answer. our main contributions are as follows: in-context learning and prompting. few-shot prompting enables llms to reason more effectively by embedding relevant examples into the input. brown et al. (2020) demonstrated that prompting with examples allows large models to generalise without gradient updates. wei et al. (2023) showed that larger models show new capabilities in icl settings, particularly in adapting to unusual input-output pairs that differ from those seen during training. subsequent work by min et al. (2022) and l. zhao and li (2023) emphasized the importance of input format and semantic similarity between prompts and queries. kojima et al. (2022) showed that even minimal prompts like "let's think step by step" significantly change reasoning style. complementing these findings, chen et al. (2023) examined the effect of demonstration quantity, revealing diminishing returns beyond a small number of examples. these studies confirm that prompting structure influences not only model accuracy but also the verbosity of the output. building on this, our method focuses on applying icl by generating high-quality demonstrations, enhancing the model's ability to generate concise reasoning leading to the correct solution. rag. rag improves llm performance by enhancing the model's input with information retrieved from external sources. the original rag framework laid the foundation by combining a retriever and a generator for open-domain tasks. more recently, chen et al. (2024) systematically assessed how different llms benefit from rag, and gao et al. (2023) offered a survey that outlined its challenges and design space. notably, while icl traditionally relied on fixed demonstrations, integrating rag to fetch relevant in-context examples has proven to be more effective. liu et al. (2021) found that selecting semantically similar in-context examples enhances gpt-3's performance, while huang et al. (2023) introduced cot-influx, a method that prunes less informative tokens to include more concise chain-of-thought examples, significantly boosting mathematical reasoning. despite its advantages, rag faces challenges in ood scenarios. finlayson et al. (2025) highlighted that finetuning llms with rag can lead to performance degradation when the training data is ood, causing misalignment between retrieved content and target responses. motivated by these limitations, we propose a dynamic alternative: instead of retrieving, our model learns to actively generate demonstrations relevant to the input. generative demonstration learning. a growing line of research focuses on generating demonstrations dynamically, rather than solely relying on retrieval. methods like self-icl wang et al. (2022) and auto-cot zhang et al. (2022) synthesize task-specific exemplars, guiding models with tailored chains of thought that often enhance performance, especially in ood settings. more recent approaches, such as auto-demo prompting feng et al. (2024), create demonstrations during batch inference by reusing earlier outputs, reducing token overhead while preserving accuracy. similarly, in long-context qa, context recycling can generate effective few-shot examples from the input passage itself cattan et al. (2024). these generative strategies offer fine-grained control over demonstration length and structure, balancing reasoning quality with efficiency in both token usage and compute cost. however, these methods rely on vanilla models prompted to generate demonstrations. in contrast, we propose training a model to query-adaptive, token-budgeted demonstrations that condition the target model's reasoning. reinforcement learning rl. rl allows us to move beyond static fine-tuning by optimizing the model based on outcome-driven feedback. schulman et al. (2017) introduced proximal policy optimization (ppo), offering stable policy optimization. ouyang et al. (2022) demonstrated the effectiveness of rlhf for instruction-following models, while rafailov et al. (2023) proposed dpo to align outputs with preferences without explicit rewards. for mathematical reasoning, recent works such as zhang and zuo (2025) and shao et al. (2024) introduced group relative policy optimization (grpo), which promotes more stable training and improves performance on complex multi-step problems. recent work also explores training efficiency: wang et al. (2025) showed that minimal supervision can drive reasoning gains, and li et al. (2025) proposed agpo to stabilize and optimize training. these findings clearly show that integrating rl in our model is particularly promising for generating highly contextually relevant demonstrations effectively. sentence embedding and dataset similarity. accurately measuring semantic similarity is essential for selecting relevant demonstrations in retrieval-based systems. motivated by this, jiang et al. (2023) investigates the use of large language models for sentence embeddings, showing that in-context learning improves embedding quality without fine-tuning. zhang et al. (2024) proposed prompt engineering techniques like pretended chain of thought and knowledge enhancement. comprehensive surveys by farouk (2019) and ai (2025) provide insights into various approaches for measuring sentence similarity. these key insights motivate our use of similarity scoring in our work: we employ it to efficiently retrieve relevant documents in our rag pipeline and systematically evaluate and rank datasets based on their similarity to the training data. ## conclusion this work introduced generative retrieval-aligned demonstration samplers (grad and gradi) for creating input-aligned demonstrations under strict token budgets, that consistently outperforms traditional rag and zero-shot methods, particularly for ood scenarios. our findings highlight the benefits of dynamic generation over static retrieval, particularly the generalizability of our approach beyond the training domain to ood queries that differ in both task and domain. we also show that our approach is modular and cost-efficient: demonstrations generated by smaller, less expensive models can boost the accuracy of larger models. we also emphasize that our approach does not replace rag. retrieval is still applicable for knowledge-intensive tasks that might require grounding from an external corpus. both grad variants can be best viewed as an alternative that aims to generalize to ood inputs, which complements, rather than replaces, retrieval-based methods.