Bidirectional KL Constraints: Enabling Small Models to Learn Tool Use

Small language models (SLMs) are widely recognized as valuable for deployment in resource-constrained environments, but a practical problem persists: they struggle to discover effective tool-calling behaviors through reinforcement learning alone, unlike large models. Direct teacher supervision, on the other hand, runs into two walls: behavioral distribution mismatch and the propagation of erroneous preferences. A recent study proposes the SeTPO method, which attempts to find a middle path between these two extremes using bidirectional KL constraints.

The Core Dilemma of Tool Learning for Small Models

External tools can extend the capability boundaries of small language models, enabling them to complete more complex tasks in compute-limited environments. However, the inherent capacity limitations of small models create a thorny contradiction: while reinforcement learning allows models to explore autonomously, the search space is too vast for small models to discover effective tool-use strategies within a reasonable time.

Another path is to have a teacher model directly provide supervision signals. This sounds straightforward and effective, but in practice there are two pitfalls. First, the behavioral distributions of teacher and student naturally differ, and the teacher's decision trajectories may not suit the student's current capability level. Second, the teacher's own preferences may not be entirely correct, and direct imitation risks inheriting the teacher's erroneous preferences as well.

Both paths have their own pitfalls, and SeTPO's approach is not to choose one over the other, but to make them work in concert.

SeTPO's Three-Layer Design Logic

SeTPO stands for Selective Teacher-guided Policy Optimization, and its core mechanism can be broken down into three layers.

The most fundamental layer is student autonomous exploration. The model attempts tool calls in the environment and adjusts its policy based on reward signals—this is the standard reinforcement learning process.

The second layer is selective teacher guidance. The key lies in the word "selective"—SeTPO does not introduce teacher signals on all trajectories, but only on low-reward trajectories, using reverse KL divergence to derive token-level teacher feedback and incorporate it into the policy advantage function. In other words, when the student is doing poorly on its own, the teacher steps in to guide; when the student is doing well, the teacher does not intervene.

The third layer is continuous demonstration supervision. SeTPO maintains an auxiliary forward KL supervision objective throughout training to ensure that behaviors demonstrated by the teacher are not forgotten or degraded during the training process. This design acts as an "anchor" for the training process, preventing the policy from drifting too far during exploration.

Policy Advantage Function and Token-Level Feedback

In policy gradient methods, the advantage function A(s,a) measures how much better or worse taking action a in state s is relative to average. A positive advantage means the action deserves encouragement; a negative one should be suppressed. Traditional reinforcement learning estimates advantages through sampling, but small models suffer from high variance in advantage estimation due to large search spaces and low sampling efficiency. SeTPO's approach is to use the teacher model's reverse KL divergence on low-reward trajectories to construct token-level advantage correction terms. Specifically, when a student's token decision deviates significantly from the teacher's preferred distribution, the reverse KL produces a penalty signal that is added to the original policy advantage, achieving precise correction without introducing complete teacher trajectories. The benefit of this design is that teacher signals only act on positions where the student genuinely errs, avoiding excessive intervention on correct decisions and reducing the risk of amplifying the teacher's erroneous preferences.

Why Bidirectional KL Rather Than Unidirectional

To understand SeTPO theoretically, one must return to the role of KL divergence in policy optimization.

Forward KL tends to make the student's policy cover all modes of the teacher's policy, resulting in comprehensive but potentially imprecise learning; reverse KL tends to concentrate the student's policy on high-probability regions of the teacher's policy, achieving precision but potentially missing certain modes.

SeTPO combines both: reverse KL is used to provide precise teacher corrections at critical moments (low-reward trajectories), while forward KL is used to maintain coverage of teacher behaviors throughout. The researchers provide a theoretical analysis under an idealized bidirectional KL objective, proving that the optimal policy's expected reward is no lower than the teacher's, while retaining a guaranteed proportion of the teacher's probability mass.

The significance of this theoretical conclusion is that it shows students need not choose between "surpassing the teacher" and "being faithful to the teacher"—the bidirectional KL design allows both to hold simultaneously.

The Relationship Between KL Divergence and Policy Optimization

KL divergence (Kullback-Leibler Divergence) measures the difference between two probability distributions and is often used as a constraint term in policy optimization to limit how far the new policy deviates from the reference policy. Its definition is asymmetric: D_KL(P||Q) ≠ D_KL(Q||P), and it is precisely this asymmetry that gives rise to two fundamentally different behavioral modes: forward and reverse. In reinforcement learning, forward KL is typically written as some variant of D_KL(teacher||student), tending to make the student cover all teacher modes; reverse KL does the opposite, concentrating the student on the teacher's high-probability regions, sacrificing some coverage for precision. Moreover, in token-level sequence generation tasks, KL constraints are computed per token, meaning tokens at different positions are subject to different constraint intensities—this is the mathematical foundation that enables SeTPO's "token-level teacher feedback." Understanding the directionality of KL divergence is a prerequisite for understanding why SeTPO uses it "bidirectionally."

Experimental Validation: Performance Across Three Scales

On two tool learning benchmarks, BFCL v4 and API Bank, SeTPO was compared against three categories of baselines: supervised fine-tuning, reinforcement learning, and distillation. Results show that across all three student model scales, SeTPO achieved superior overall performance.

What is notable about this result is its cross-scale stability. A common problem in tool learning for small models is that a method works at a particular scale but fails at another. SeTPO's consistent performance across three scales indicates that the benefits of bidirectional KL constraints do not depend on specific model capacity.

BFCL v4 and API Bank Benchmarks

BFCL (Berkeley Function Calling Leaderboard) is an authoritative benchmark for evaluating large language models' tool-calling capabilities, and its v4 version covers multi-turn dialogues, parallel calls, error recovery, and other test items closer to real-world scenarios. API Bank is another evaluation set focused on API call accuracy and parameter-filling correctness. Together, they provide a relatively comprehensive measure of model performance in tool learning. The significance of choosing these two benchmarks is that they not only test whether a model can call tools, but also the timing of calls, parameter correctness, and coherence in multi-step tasks—dimensions where small models are particularly prone to exposing weaknesses. SeTPO's cross-scale stability on these two benchmarks indicates that its benefits are not incidental results of a particular evaluation set, but have methodological generality.

Implications for the Tool Learning Paradigm

A broader insight from SeTPO is that teacher guidance and student exploration need not be a zero-sum relationship. Through carefully designed constraint mechanisms, teachers can provide precise help when students need it without suppressing students' own exploration capabilities.

This idea of "selective intervention" also has reference value for other training scenarios requiring knowledge transfer—whether it is having large models teach small models or strong models teach weak models, how to calibrate the timing and intensity of intervention is often more critical than whether to intervene at all.