Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MENTOR: A Reinforcement Learning Framework for Enabling Tool Use in Small Models via Teacher-Optimized Rewards

ChangSu Choi, Hoyun Song, Dongyeon Kim, WooHyeon Jung, Minkyung Cho, Sunjin Park, NohHyeob Bae, Seona Yu, KyungTae Lim

TL;DR

Mentor introduces a reinforcement learning framework that distills a large teacher's tool-use policy into smaller language models by combining GRPO-based policy optimization with a dense, teacher-guided reward design. It overcomes supervised fine-tuning’s poor generalization and sparse-reward RL's exploration inefficiency by training the student to internalize a generalizable tool-use methodology rather than mimic fixed demonstrations. The approach leverages reference trajectories from the teacher to construct a three-component reward (Correctness, Teacher-Alignment, Tool Validation) within a math-focused curriculum to promote robust exploration and policy alignment. Across in-domain and out-of-domain benchmarks, Mentor yields substantial gains in tool-use efficiency and cross-domain generalization, demonstrating practical promise for on-device, tool-augmented reasoning in small models.

Abstract

Distilling the tool-using capabilities of large language models (LLMs) into smaller, more efficient small language models (SLMs) is a key challenge for their practical application. The predominant approach, supervised fine-tuning (SFT), suffers from poor generalization as it trains models to imitate a static set of teacher trajectories rather than learn a robust methodology. While reinforcement learning (RL) offers an alternative, the standard RL using sparse rewards fails to effectively guide SLMs, causing them to struggle with inefficient exploration and adopt suboptimal strategies. To address these distinct challenges, we propose MENTOR, a framework that synergistically combines RL with teacher-guided distillation. Instead of simple imitation, MENTOR employs an RL-based process to learn a more generalizable policy through exploration. In addition, to solve the problem of reward sparsity, it uses a teacher's reference trajectory to construct a dense, composite teacher-guided reward that provides fine-grained guidance. Extensive experiments demonstrate that MENTOR significantly improves the cross-domain generalization and strategic competence of SLMs compared to both SFT and standard sparse-reward RL baselines.

MENTOR: A Reinforcement Learning Framework for Enabling Tool Use in Small Models via Teacher-Optimized Rewards

TL;DR

Mentor introduces a reinforcement learning framework that distills a large teacher's tool-use policy into smaller language models by combining GRPO-based policy optimization with a dense, teacher-guided reward design. It overcomes supervised fine-tuning’s poor generalization and sparse-reward RL's exploration inefficiency by training the student to internalize a generalizable tool-use methodology rather than mimic fixed demonstrations. The approach leverages reference trajectories from the teacher to construct a three-component reward (Correctness, Teacher-Alignment, Tool Validation) within a math-focused curriculum to promote robust exploration and policy alignment. Across in-domain and out-of-domain benchmarks, Mentor yields substantial gains in tool-use efficiency and cross-domain generalization, demonstrating practical promise for on-device, tool-augmented reasoning in small models.

Abstract

Distilling the tool-using capabilities of large language models (LLMs) into smaller, more efficient small language models (SLMs) is a key challenge for their practical application. The predominant approach, supervised fine-tuning (SFT), suffers from poor generalization as it trains models to imitate a static set of teacher trajectories rather than learn a robust methodology. While reinforcement learning (RL) offers an alternative, the standard RL using sparse rewards fails to effectively guide SLMs, causing them to struggle with inefficient exploration and adopt suboptimal strategies. To address these distinct challenges, we propose MENTOR, a framework that synergistically combines RL with teacher-guided distillation. Instead of simple imitation, MENTOR employs an RL-based process to learn a more generalizable policy through exploration. In addition, to solve the problem of reward sparsity, it uses a teacher's reference trajectory to construct a dense, composite teacher-guided reward that provides fine-grained guidance. Extensive experiments demonstrate that MENTOR significantly improves the cross-domain generalization and strategic competence of SLMs compared to both SFT and standard sparse-reward RL baselines.
Paper Structure (46 sections, 8 equations, 7 figures, 12 tables, 1 algorithm)

This paper contains 46 sections, 8 equations, 7 figures, 12 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Tool-Augmented Language Models
  4. Reasoning Distillation through SFT
  5. Reinforcement Learning and Reward Design
  6. Mentor: Model Enhancement via Teacher-Optimized Reward
  7. Reference Trajectory Generation
  8. Adapting GRPO for Reasoning Distillation
  9. Generating Rollouts
  10. Student Policy Optimization
  11. Teacher-Guided Reward Design
  12. Correctness Reward
  13. Teacher-Alignment Reward
  14. Tool Validation Reward
  15. Experiments and Results
  16. ...and 31 more sections

Figures (7)

  • Figure 1: Comparative analysis of the tool effect on model performance and tool invocation patterns between LLMs and SLMs.
  • Figure 2: Overview of the Mentor training framework. A problem-solving trajectory ($\tau$) consists of a sequence of Reasoning (R), Tool-Use (T), and Observation (O), and a final Answer (A). Each student rollout is evaluated by a set of reward models, which generate a reward signal by aligning the student's actions against the teacher's reference trajectory.
  • Figure 3: Correlation between task performance (Math) and alignment score (AS).
  • Figure 4: Comparison of tool invocation patterns between student models and the teacher. Each bar represents the percentage point difference in invocation frequency for a specific tool between the student model and the teacher.
  • Figure 5: Tool-use efficiency on in-domain and out-of-domain tasks, measured by the distribution of tool calls per sample.
  • ...and 2 more figures