Technical Summary: MENTOR

Problem Statement

The integration of Large Language Models (LLMs) with external tools (e.g., code interpreters, retrieval APIs) has enabled advanced reasoning capabilities. However, the high inference costs of large-scale models limit their practical deployment, motivating the distillation of these tool-use capabilities into Small Language Models (SLMs).

Current distillation methods primarily rely on Supervised Fine-Tuning (SFT), where student models are trained to strictly replicate teacher-generated trajectories. The authors identify that this rigid alignment with static token distributions stifles the exploratory autonomy of capacity-limited SLMs, leading to poor out-of-domain (OOD) generalization.

While Reinforcement Learning (RL) offers an alternative by enabling models to learn through interaction, designing reward structures for SLMs presents a severe dilemma:

1. Sparse Outcome Rewards: Providing rewards only for final correctness offers insufficient guidance, leading to inefficient exploration and tool avoidance.
2. Strict Trajectory Matching: Enforcing exact sequence replication imposes an overly restrictive optimization burden on SLMs, hindering their ability to learn underlying strategic logic and adapt to unseen tools.

The paper posits that SLMs exhibit a "fragile tool-use behavior" characterized by a disconnect between tool invocation frequency and success rates, often resulting in invalid tool calls and strategic misalignment with the teacher.

Methodology: MENTOR

To address these limitations, the authors propose MENTOR (Model Enhancement by Non-rigid Teacher-Optimized Rewards), a reinforcement learning framework designed to bridge the strategic gap between teacher and student through a flexible, process-aware reward structure.

Core Framework

MENTOR utilizes the Group Relative Policy Optimization (GRPO) algorithm. Instead of treating the teacher's trajectory as a rigid target for sequential matching, it serves as a strategic reference to guide the student's exploratory process.

Flexible Reward Design

The core innovation lies in a composite reward signal $R(O^{(s)}, O^{(t)})$ that balances three components:
$$R = w_c R_c + w_a R_a + w_v R_v$$

1. Correctness Reward ($R_c$): A binary signal (1 or 0) indicating whether the student's final answer matches the teacher's outcome (which is verified against ground truth). This ensures the model remains task-oriented.
2. Teacher-Alignment Reward ($R_a$): Unlike strict sequence matching, this component employs a set-based matching approach. It rewards the student if the set of tools used matches the teacher's set of tools, regardless of the execution order or call counts. This encourages the student to adopt the teacher's high-level tool-selection strategy while allowing flexibility in the reasoning path.
3. Tool Validation Reward ($R_v$): A binary signal that penalizes invalid or malformed tool calls. This component is crucial for SLMs, which are prone to execution errors, ensuring that the model learns to generate executable code.

The framework allows the student to explore alternative valid reasoning paths that utilize the correct tools, rather than forcing a literal token-for-token imitation of the teacher's trajectory.

Key Contributions

1. Identification of the Reward Dilemma: The authors empirically demonstrate that capacity-limited SLMs face a bottleneck where sparse outcome signals fail to induce tool use, while strict trajectory constraints lead to excessive invalid calls and poor generalization.
2. The MENTOR Framework: They introduce a flexible, process-aware reward structure that mitigates this dilemma. By decoupling strategic alignment (tool selection) from rigid sequential constraints, MENTOR provides moderated guidance that accommodates the SLM's limited capacity.
3. Empirical Validation: Extensive experiments across multiple benchmarks show that MENTOR significantly outperforms both SFT and strict RL baselines, particularly in OOD scenarios.

Experimental Results

The authors evaluated MENTOR using Qwen3-235B as the teacher and various Qwen2.5 and Qwen3 models (1.5B to 8B) as students.

• In-Domain Performance: On mathematical reasoning benchmarks (e.g., AIME, AMC, MATH), MENTOR consistently achieved the highest accuracy among all baselines, including SFT and other RL variants. For instance, on the 7B model, MENTOR improved overall Math accuracy to 28.63% compared to 26.64% for Strict RL and 22.48% for SFT.
• Out-of-Domain Generalization: MENTOR demonstrated superior zero-shot generalization on unseen tools and domains (BFCL-v4 and RAG tasks). The flexible alignment allowed SLMs to adapt to new tool-calling patterns more effectively than models trained with rigid trajectory matching.
• Behavioral Analysis:
  • Tool Usage: MENTOR-trained models exhibited tool usage distributions closely aligned with the teacher, whereas SFT models retained pre-existing, suboptimal invocation habits.
  • Stability: MENTOR reduced the frequency of invalid tool calls and eliminated the "redundant verification loops" observed in SFT and strict RL baselines, where models would continue calling tools even after deriving the correct answer.
  • Ablation Studies: Removing the teacher-alignment reward ($R_a$) or the tool validation reward ($R_v$) led to significant performance drops, confirming the necessity of the composite reward structure.

Significance and Claims

The paper claims that within verifiable tool-use environments, flexible tool-use alignment is more effective than strict trajectory replication for developing adaptable small models.

• Strategic vs. Imitative Learning: The findings suggest that SLMs do not need to memorize specific token sequences to learn tool use; rather, they benefit from learning the strategic logic of tool selection (i.e., which tools to use) while retaining the autonomy to determine how to sequence them.
• Bridging the Capacity Gap: MENTOR effectively narrows the performance gap between 1.5B–8B models and their 235B teacher, offering a viable path for deploying capable, on-device AI agents without relying on expensive cloud infrastructure.
• Limitations: The authors modestly note that their approach assumes the student model has some initial ability to generate executable tool calls (often requiring a brief SFT warm-up). They also acknowledge that the current evaluation is limited to controlled, executable-tool settings (math and retrieval) and does not yet address the complexities of open-ended real-world interactive environments (e.g., web browsing, multi-step state tracking).

In conclusion, MENTOR proposes that the future of tool-use distillation lies not in forcing smaller models to mimic larger ones perfectly, but in providing them with flexible, teacher-guided signals that respect their capacity constraints while fostering robust, generalizable reasoning strategies.