Distilling LLM Agent into Small Models with Retrieval and Code Tools

We sincerely thank the reviewer for the thoughtful comments and positive assessment. We are glad that you found (1) the agent distillation idea well-motivated and well-executed, and (2) appreciated the breadth of our experiments, including the diverse benchmarks and model sizes. Below, we address the questions and concerns in detail.

---

W1. The training method is a pretty straightforward extension of existing distillation methods and may slightly lack novelty. But a simple method that works well is a reasonable contribution.

We appreciate the recognition that a simple method with strong empirical results can be a valuable contribution. Although our framework builds on distillation, it introduces two novel components tailored for agent distillation:

• First-Thought Prefix (FTP): Unlike standard prompting, FTP prepends a first thought generated via CoT prompting to guide the agent's reasoning, improving performance on reasoning-intensive tasks for both LLM (Table 4) and distilled agents (Table 2).
• Self-Consistent Action Generation (SAG): his extends self-consistency to agent settings by sampling diverse action trajectories and selecting valid, executable ones based on environment feedback, thereby enhancing robustness and action quality.

These components, to our knowledge, have not been explored in the context of tool-using agent distillation.

---

Q1. One common pitfall of distillation is overfitting to teacher rollouts. Did the author observe any overfitting in their experiments and how to prevent them?

To assess generalization, we evaluated the distilled agents on six out-of-distribution tasks, some of which are more challenging than the training tasks (e.g., AIME vs. MATH, MuSiQue (3-hop) vs HotpotQA (2-hop)). The consistent performance improvements over CoT-distilled baselines on those tasks suggest that overfitting is limited. We also used LoRA-based fine-tuning, which is known to reduce overfitting compared with full fine-tuning [1].

[1] Biderman et al., LoRA Learns Less and Forgets Less, TMLR 2024

---

Q2. How many rollouts did the authors collect from teacher per prompt?

We collected one teacher rollout per prompt for training. As suggested in prior work on CoT distillation [1], using multiple rollouts may further improve performance, which we plan to explore in future extensions.

[1] Ho et al., Large Language Models Are Reasoning Teachers, ACL 2023

We sincerely thank the reviewer for the feedback and their appreciation of our work on multiple aspects. We appreciate the reviewer's recognition of (1) our clear writing, (2) our complete code release, (3) thorough ablations, (4) simple yet effective techniques, and (5) strong generalization across diverse benchmarks. Below, we address each concern.

---

W1. Evaluation on single model family: Even though it is acknowledged as a limitation, all experiments were conducted on Qwen models

Q1. Could the authors provide additional experiments involving llama, phi, mistral, or other alternatives? It would bring more certainty about the Agent Distillation's strengths.

Thank you for the thoughtful comment. We agree that evaluating our method across multiple model families would strengthen the generality of our findings. We plan to extend our evaluation to other small models, such as LLaMA-3.2-1B-Instruct and Phi-3-mini-instruct, in a future revision.

---

W2. SAG poses some safety risks: Executing code sampled from a stochastic policy could lead to irreversible or unsafe operations.

Q3. Could the authors further discuss whether code execution during SAG is safe? Given the stochasticity of decoding, is there a risk of producing and executing dangerous or irreversible commands from the long tail of the model's distribution?

We appreciate the reviewer raising this important concern. In practice, such risks can be mitigated through sandboxing mechanisms (e.g., Docker, E2B Sandbox), which are widely adopted in agent systems such as OpenHands [1]. Additionally, incorporating a safety guard model to verify code safety prior to execution is another viable future direction. We have clarified this in the discussion section.

[1] Wang et al., OpenHands: An Open Platform for AI Software Developers as Generalist Agents, ICLR 2025

---

Q2. Have the authors thought of any ways of mitigating the downsides of the first-thought prefix (i.e., reduction of retrieval calls)?

One possible solution is to apply reinforcement learning after distillation to encourage effective retrieval behaviors in students. While we did not pursue this direction in this work, we agree it is a promising avenue for future research in training small agents.

We sincerely thank the reviewer for the thoughtful and prompt response.

We fully agree that evaluation on additional small models is important, and we apologize for not including at least preliminary results in our rebuttal. While the process itself is not complex, our experiments have been delayed due to limited GPU availability. We are currently running evaluations on other small models and will share preliminary results within the next two days.

Thank you again for your continued engagement and constructive feedback.

Dear Reviewer NTqR,

We are pleased to share our preliminary evaluation results for Llama-3.2-1B-Instruct:

These results mirror the trends we observed with Qwen-2.5 models: our distilled agent with both FTP and SAG outperforms the distilled CoT baseline by over 10 points on average. The performance gains from applying FTP and SAG are similarly consistent, demonstrating the effectiveness of our methods across different model families.

We have included these findings into the revised manuscript. We are also running evaluations on additional small models, such as Phi-4-mini-instruct (3.8B), and will report those results during the discussion phase.

Please let us know if you have any further questions or comments. We would be happy to address them. Thank you again for your time and valuable feedback.

Dear Reviewer NTqR,

We would like to share additional evaluation results on Phi-4-mini-instruct (3.8B), which further support the generality of our proposed method:

These results are consistent with the trends we observed on Qwen2.5 and Llama3 models, further confirming that our approach is effective across multiple model families. We have updated the revised manuscript to reflect these findings.

If you have any further concerns or questions, please feel free to share them during the discussion phase. We would be more than happy to engage.

Best regards,

The Authors

We thank the reviewer for the feedback and appreciate that they found our work (1) clearly written and (2) evaluated on multiple benchmarks with extensive analysis. We believe some concerns may stem from a lack of clarity about certain aspects of our contributions. In the following, we provide explanation to clarify our novelty and distinguish our work from prior works.

---

W1. No particular novelty in the method -- same distillation loss, but the CoT trajectory is substituted with the agentic one. The prompting and self-consistency in generating code indeed enhance the quality of the trajectories, but have already been studied in the literature, e.g. [6].

The contribution of our work is not based solely on methodological novelty. As acknowledged by other reviewers, our work provides well-motivated methods (NTqR, yrcR), comprehensive and extensive evaluations (qN16, NTqR, yrcR), and a valuable contribution to the field (qN16). We also argue that our method does introduce meaningful contributions beyond simply applying existing techniques, as detailed below:

• First, the First-Thought Prefix (FTP) is not a naive reuse of existing prompting methods. We identified a specific failure mode of LLM agents in reasoning tasks and proposed a novel solution: prepending an initial "first thought" derived from CoT prompting to better steer the agent's reasoning process. This approach results in improved performance for both the teacher (Table 4) and student models (Table 2). To our knowledge, no prior work has applied similar approach to agentic reasoning or explored its advantages for distillation.
• Second, our Self-Consistent Action Generation (SAG) adapts the idea of self-consistency decoding to the agentic setting. Unlike prior methods that rely on ground-truth answers [1] or external LLM judges [2], our approach uses majority voting over sampled actions, based solely on observations from the environment. This adaptation and its demonstrated effectiveness (Table 2) have not been explored in previous work, highlighting its contribution.

[1] Wang et al., Self-Consistency Improves Chain of Thought Reasoning in Language Models, ICLR 2023

[2] Chen et al., Universal Self-Consistency for Large Language Model Generation

---

W2. [1] uses a large-scale LLM to collect a dataset of reasoning chains and tool calls and then applies DPO to train a smaller LLM. The proposed method should be compared to this baseline. Overall, the related works miss a lot of concurrent work, e.g., [4].

We clarify that while our work and Manduzio et al. (2024) [1] share the high-level goal of training small models using trajectories from LLM agents, their objectives, tool settings, training methodology, and evaluation scope differ substantially, as summarized below.

Due to these differences, we believe a direct comparison would not be meaningful. We have cited this work and clarified its distinction in the revision.

---

W3. The main benefit over the original CoT distillation comes from using tools, not the distillation method itself, as can be seen by the higher average score of the teacher model when allowed to use tools, as well as a <= 3% difference in scores between distillation w and w/o FTP+SAG. Thus, it would be fair to compare with some of the tool-learning methods, like TinyAgent [2], ToolkenGPT, and others [3, 4].

The teacher’s performance gain indeed stems from its ability to use tools. However, we argue that this fact underscores the necessity and value of our agent distillation method.

Our central hypothesis is that a student model cannot acquire tool-use capabilities through simple prompting alone. As clearly shown in Table 2, directly prompting the student to use tools yields minimal or even negative gains, where 0.5B model achieves nearly zero accuracy.

This demonstrates that only through distillation of the teacher's agent trajectories can tool-use ability be effectively transferred to the student. In other words, agent distillation is the mechanism that unlocks the potential of tools for small language models (sLMs).

Furthermore, while the 2-3 point absolute gain from FTP+SAG may appear modest, it is highly significant for sLMs, corresponding to their relative gains of 13.83%, 8.87%, and 8.93% for 0.5B, 1.5B and 3B models, respectively.

Therefore, our work should be viewed not as a simple tool-learning method, but as a study that identifies the limits of CoT distillation and pushes the capability boundaries of sLMs by the agent distillation paradigm.

We thank the reviewer for their valuable and insightful feedback. We are glad to hear that the reviewer found our work (1) offers a practical and effective pathway for transferring the agentic capabilities, (2) presents a valuable contribution towards developing powerful on-device agents, (3) is well-structured and the experiments are comprehensive.

To address your concerns, we have clarified our contribution in terms of methods, run additional suggested experiments regarding ablation studies on SAG and multiple runs on AIME.

---

W1. While their application is effective, the core methods are derivative of existing work. The first-thought prefix (FTP) is a well-executed form of prompt engineering, designed to elicit a specific reasoning format. The self-consistent action generation (SAG) is a direct application of the well-established Self-Consistency decoding strategy.

We acknowledge that both the First-Thought Prefix (FTP) and Self-Consistent Action Generation (SAG) are inspired by prior work. However, we would like to clarify that our contribution lies in tailoring and validating these ideas in the novel context of agent distillation.

• FTP addresses the performance drop of language model agents in reasoning tasks. We find prepending the first-thought from CoT in generating agent trajectories improves the performance of LLM agent by +5.07 (Table 4) as well as downstream performance of distilled agent (1.5B) by +1.05 (Table 2). To our knowledge, this problem has not been identified or addressed in prior work.
• For SAG, while it is motivated by self-consistency decoding, our work is the first to adapt and evaluate it in the context of action generation for interactive language agents. Unlike traditional uses of self-consistency that rely on access to ground-truth answers or additional LLM as a judge [1], we leverage majority voting over observations from the environment without external supervision. Importantly, this method remains effective even for the smallest 0.5B model evidenced by +1.91 performance improvement in average for 8 tasks (Table 2).

We hope this clarifies that our contributions extend beyond applying existing ideas.

[1] Chen et al., Universal Self-Consistency for Large Language Model Generation

---

W2. The paper proposes the first-thought prefix (FTP) to improve teacher quality but only validates it on a single model (Qwen2.5-32B-Instruct). It remains a question whether this specific prompting strategy is effective for other, more powerful reasoning models like DeepSeek-R1 or the OpenAI-o1, which may have different internal reasoning structures.

Although we have not yet evaluated FTP on DeepSeek-R1 or OpenAI models because of computational-budget constraints, we agree that testing a reasoning model as teacher would strengthen our claim. We therefore experimented with a feasible open-source reasoning model, QwQ-32B. However, we observed that QwQ-32B solved most problems entirely inside its thinking steps without tools, so meaningful agentic traces are hard to be collected. This result suggests that strong reasoning models may bypass tool use, limiting straight-forward agent distillation. We have included this discussion to the revision and leave using reasoning models as the teacher for future work.

---

W3. The experimental analysis, while broad, lacks depth in key areas. The individual contributions of FTP and SAG are not clearly disentangled, making it difficult to quantify their separate impacts on the final performance. Key hyperparameters, such as the sampling temperature for SAG, are presented without ablation studies to justify their chosen values.

We already provide an ablation of FTP alone and FTP+SAG in Table 2, showing that applying each method progressively improves performance. To further isolate the effect of SAG, we conducted additional experiments using SAG alone on the 0.5B, 1.5B, and 3B models. As shown in the table below, SAG alone improves average performance by up to +2.2 points (1.5B) and consistently yields benefits across all scales.

• D denotes the distilled agent.

We provide the ablation on sampling temperature in our response to Q1.

---

Q1. In describing Self-Consistent Action Generation (SAG), you state that you sample with a "high temperature" to encourage diversity, yet the value used is 0.4. In many contexts, this is considered a relatively low temperature that promotes less randomness. Could you clarify this choice and provide any ablation studies you performed on the effect of temperature on SAG's performance?

Thank you for pointing this out. We acknowledge the term "high temperature" in Line 159 can be misleading. Our intent was to indicate a shift from deterministic decoding (temperature = 0) to sample diverse actions. We chose 0.4, as performance did not vary significantly across temperatures in SAG.

As suggested by the reviewer, we conducted a temperature ablation study on SAG for math tasks using the 1.5B model with temperatures 0.2, 0.4, 0.6, 0.8, and 1.0. The results are shown below and have been included in the revision.

---

Q2. Could you provide more detail on the evaluation protocol? Specifically, were the reported scores, especially for benchmarks with small size like AIME, averaged over multiple runs with different random seeds to ensure statistical significance?

We report all experimental results with a single run due to computational budget constraints, as noted in Checklist item 7. We have clarified this detail in the experimental setup section of the revision.

To assess variance on smaller benchmarks, we additionally ran AIME with 5 different seeds for each model size with the FTP and SAG methods. The results below report the average and standard deviation:

These results show that the distilled agents yield consistent performance on AIME with relatively low standard deviation across model sizes. We have also included these statistics in the revised manuscript to enhance statistical rigor on the small benchmark.

---

Q3. For the student models, you chose to use LoRA for fine-tuning. Given the relatively small size of both the models and the datasets, what was the rationale for using LoRA instead of full fine-tuning? Did you experiment with full fine-tuning, and if so, how did the results compare?

The primary rationale for using LoRA was its efficiency in terms of both reduced training cost and ease of deployment. Although we did not experiment with full fine-tuning, we believe this is a valuable direction and plan to include it in the future revision.

---

Q4. An interesting trend appears on the MATH500 benchmark, where standard CoT prompting outperforms your distillation method for the 3B and 7B models, reversing the trend seen with the 0.5B and 1.5B models. Do you have a hypothesis for this scale-dependent performance difference?

Thank you for this insightful observation. Our hypothesis is as follows:

• Smaller models (0.5B, 1.5B) benefit greatly from tool use, as it compensates for their limited arithmetic ability.
• Larger models (3B, 7B) have stronger internal computation skills, making tool use less beneficial on simpler benchmarks like MATH500.

However, agent distillation remains effective for 3B and 7B models. As shown in Figure 3, the 3B distilled agent performs better on harder MATH500 problems (level 5) with the agentic method. Moreover, on benchmarks involving more complex calculations (e.g., GSM-Hard, OlymMATH), the performance gains from agent distillation are consistent across 3B and 7B models. We have clarified this task-dependence in the experimental results section.

---

Q5. The methodology states that you "filter out any sequences that result in parsing or execution errors" (lines 160-161). Logically, this filtering should result in a zero error rate for the final selected action. However, the results in Figure 6 show that SAG still produces a non-zero error rate. Could you clarify why the error rate is non-zero, or how is it calculated?

We thank the reviewer for highlighting this. The filtering rule, discarding samples with parsing or execution errors, applies only when at least one valid action exists.

When all generated actions fail, we retain randomly-selected one failed action and feed its error message back as an observation, allowing the model to self-correct. The non-zero error rate in Figure 6 reflects these cases. We have clarified this in the method section to avoid confusion.

Dear Reviewer qN16,

We would like to follow up regarding Q3, as we have now conducted the full fine-tuning (FT) experiment that we were unable to include in the rebuttal due to resource constraints. Below, we share the preliminary results on Qwen2.5-1.5B-Instruct:

As shown, full fine-tuning does not yield better performance on average compared to LoRA-based fine-tuning. While further hyperparameter tuning may help, these results suggest that full fine-tuning could hurt generalization possibly due to overfitting. We have incorporated this result and analysis into the revised manuscript.

Thank you again for your time and thoughtful review. We hope that our rebuttal has addressed most of the concerns you raised. If there are any remaining questions, concerns, or clarifications you would like us to address, we would be sincerely grateful for the opportunity to discuss them during the discussion phase.

Best regards,

The Authors