PROOF-Gen: From Optimized Data to Better Distillation

Standard tool-calling distillation pipelines rely on a generate-and-filter mechanism, where teacher models generate trajectories and only those that pass verification are kept for training. This approach discards a significant portion of data—often near-misses where the teacher fails at a single decision point—leaving the student model without signal on the hardest scenarios. This paper asks: can we recover these failed trajectories to improve the distillation of tool-calling capabilities?

The authors introduce PROOF-Gen (Per-scenario Reflective Optimization to Overcome Failed Generation). Instead of searching for a single prompt that works across all tasks, PROOF-Gen optimizes a unique, disposable prompt for each failed scenario. A reflector model analyzes the failed execution trace and evaluation feedback, then writes a corrective cheatsheet that is appended to the teacher's system prompt. The teacher re-executes the task until it passes. Crucially, this cheatsheet is stripped before the student model is fine-tuned, ensuring the student learns from clean, scaffold-free demonstrations.

PROOF-Gen recovers 93% of failed scenarios on the τ2-bench telecom domain. When fine-tuned on this combined data (filtered + recovered trajectories), small models like Qwen3-4B-Instruct-2507 and Gemma 4 E4B-it show significant gains in end-to-end task completion. On τ2-bench, Qwen3-4B improves from a Pass^1 of 0.132 to 0.529. In production settings, the method lifts goal completion by 6.3 percentage points and shows positive transfer across multiple locales, confirming that the corrective knowledge is largely language-agnostic.

This work demonstrates that the ceiling on distillation quality is often set by data generation rather than training algorithms. By treating failures as recoverable data rather than noise, PROOF-Gen allows teams to expand task coverage and improve agent performance without needing to retrain the teacher or change the student's interaction style. It effectively decouples problem-solving capability from the teacher's voice, enabling the distillation of stronger reasoning into deployable models.