Training Small Language Models for Transferable Tool Selection
Zhaoyi Zhou, Akshayaa Magesh, Jeongyeol Kwon, Artem Kuriksha, Daniel Jiang, Andrea Zanette, Ayush Jain
cs.LG, cs.AI
2026-10-23
SFT plus GRPO trains Qwen3-8B to split instructions and select tools, reaching 41.5% exact match on ToolBench G3, above prompted Llama-4-Maverick (402B) at 30.6%.
Once a tool catalog passes ten thousand entries, the descriptions no longer fit in context, and choosing a tool blocks the call itself. BM25 ranks by word overlap. Dense retrievers rank by embedding similarity. Both score resemblance. Neither knows whether a step actually needs the tool. Instructions often bundle several requests, and one query lets the words of one request drown out another. A high rank is not the same thing as a required tool.
Recent systems put Claude or ChatGPT in charge of the retriever: split the instruction, fetch candidates per substep, then name a tool. Tool selection is a subroutine an agent repeats, and calling a closed model every time is a poor deployment fit. An off-the-shelf small model does badly at this Plan-Retrieve-Select loop. A trained one is only useful if it can switch tool catalogs and retrievers without a new training run.
The models are Qwen3 at 1.7B, 4B, and 8B. One checkpoint, two prompts: plan, then select. The retriever sits outside the model.
Planning splits the instruction into short substeps, each naming one capability. Each substep is searched on its own. The default retriever is BM25, which needs no corpus-specific training. Tools scoring above τ=3 are kept and cut to at most K=15. Selection reads those descriptions, commits to one tool per substep, and returns the union. Two assumptions fix the task shape: each substep maps to exactly one tool, and substeps are independent, so a later step never reads an earlier return. Each call is stateless, and the context does not grow with the catalog.
The base model rarely finishes the loop, so reinforcement learning from scratch starts slowly. Llama-4-Maverick writes supervised trajectories as a warm start. ToolBench gold tools come from reference solution traces and include tools the user never mentioned. Pairs that share no content word are dropped. A teacher keeps the rest only when the query paraphrases the tool, and gold tools the decomposition never uses are removed. Plan reasoning is rewritten so gold labels the teacher saw do not leak into the chain of thought.
Supervised fine-tuning teaches the format and a basic split. GRPO then trains planning and selection apart. GRPO samples a group of outputs for one input and updates from which samples look better inside that group. During selection training the substeps are the teacher's, not the model's own plans, so a bad plan does not corrupt the reward.
The planning reward needs no extra human labels. Each substep goes to the retriever and is matched, in order, to one gold tool that is still free. Precision is the share of substeps that find one, which penalizes extra substeps. Recall is the share of gold tools covered, discounted by rank, so a gold tool near the top scores higher. The reward is their F1. The selection reward is 1 only when the model picks exactly that substep's gold tool. At test time, success means the predicted set equals the filtered gold set for the whole instruction.
Training uses ToolBench G3: 21,361 instructions and 1,419 APIs. The test split has 2,124 instructions. Exact match is stricter than Recall@K. One extra tool or one missing tool fails the example. The unit being selected is a single API.
At 8B, the base scores 20.5%, supervised fine-tuning 36.2%, and supervised fine-tuning plus RL 41.5%. RL from the base, with no warm start, reaches 35.7%. The full recipe scores 37.6% at 1.7B and 40.3% at 4B. Supervised fine-tuning alone clears 30% at every size.
Prompted Llama-4-Maverick (402B) peaks at 30.6% with a Search-R1-style single stream, the best of four protocols. The same Plan-Retrieve-Select workflow scores 10.1% with Maverick, below untrained Qwen3-8B. ReAct scores 28.4%, with up to five model calls per substep. Selection here is one call per substep. Trained Qwen3-1.7B still beats Maverick's best prompted result.
| Method | Setup | G3 exact match |
| Llama-4-Maverick | hierarchical search | 8.6% |
| Llama-4-Maverick | this workflow, prompt only | 10.1% |
| Llama-4-Maverick | ReAct | 28.4% |
| Llama-4-Maverick | Search-R1 style | 30.6% |
| Qwen3-8B | base | 20.5% |
| Qwen3-8B | RL only | 35.7% |
| Qwen3-8B | SFT only | 36.2% |
| Qwen3-1.7B | SFT+RL | 37.6% |
| Qwen3-4B | SFT+RL | 40.3% |
| Qwen3-8B | SFT+RL | 41.5% |
The same checkpoints transfer with no further training. G2 holds 11,474 tools, about 8 times the G3 catalog, with about 9% overlap. The 8B model falls from 41.5% to 25.8% (4B to 23.5%) and still beats Maverick plus Search-R1 at 20.8%. Gorilla and CRAFT have one gold tool each, so they test selection rather than planning. The 8B model scores 29.9% and 69.0%. Maverick's best scores there are 16.0% and 47.9%. On CRAFT the untrained 8B already scores 53.4%; SFT moves it to 66.4%, and RL to 69.0%. Most of the cross-benchmark gain is the supervised warm start.
TaskBench is multi-tool and out of domain. Maverick with Search-R1 scores 80.6%, against 72.8% for the 8B model. Planning there needs some domain knowledge, and skipping the decomposition does better. The 8B model still beats Maverick on ReAct (60.2%) and on this same workflow (36.9%).
Training uses BM25. Swapping in COLT, a dense retriever trained on ToolBench, raises success at all three sizes. The paper does not print the COLT percentages. BM25 is the training retriever so the transfer numbers do not depend on an in-domain expert.
In the 8B ablations on the G3 test set, with selection and retrieval held fixed, an untrained Qwen3-8B planner scores 27.6%, a Maverick planner 17.2%, and the trained planner 41.5%. Always taking the BM25 top-1 tool scores 28.6%. Success rises until the candidate cap K reaches 20 and then flattens. The default K=15 is not the peak.
On this narrow task, a trained 8B model beats a prompted 402B model. The checkpoint can change catalogs and retrievers without retraining, and selection is one call per substep. In-domain exact match is still 41.5%, so most instructions miss the gold set. On TaskBench the larger model without decomposition scores higher. The 8B selector fits when the catalog is large and the substeps are mostly independent. It does not cover pipelines that pass results between steps, or planning that itself needs domain knowledge.
The workflow allows one tool per substep and no data flow between steps. Any instruction where a later call must read an earlier return is out of scope.
Training and evaluation share the filtered gold sets, not the original ToolBench labels. The comparison with Maverick is fair on those labels and is not comparable to older Recall@K numbers. Tools the query does not surface are removed, so exact match is both narrower and easier.
Selection is trained on the teacher's substeps and run on the model's own plans. That gap is unmeasured. The small model beats a prompted 402B teacher, not a 402B model trained with the same reward.
G2 drops from about 40% to about 25%. Transfer works, and the drop is large. The COLT gain is tied to a ToolBench expert retriever. The training catalog has only 1,419 APIs; the ten-thousand-tool regime appears at transfer time. The gold filter is not checked against an outside annotator, and the main text does not report enough GRPO hyperparameters to reproduce the run.