🤖 AI

PROOF-Gen: From Optimized Data to Better Distillation

The paper introduces PROOF-Gen, a method that uses per-scenario prompt optimization to recover successful trajectories from teacher model failures, thereby significantly improving the performance of distilled tool-calling agents on benchmarks and in deployed pipelines.

Original authors: Anh Ta, Junjie Zhu, Shahin Shayandeh

Published 2026-08-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Anh Ta, Junjie Zhu, Shahin Shayandeh

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the world of artificial intelligence, a common way to teach a smaller, faster computer program to perform complex tasks is to have a much larger, smarter program demonstrate how to do it. This process, known as distillation, works like a master craftsman showing an apprentice a perfect technique. The apprentice watches the master's successful attempts and learns to mimic them. However, this method has a blind spot. If the master makes a mistake, the apprentice simply ignores that attempt and waits for the next successful demonstration. In the specific field of tool-calling—where an AI must decide which digital tools to use to solve a problem, such as checking a database or sending a message—this creates a significant gap. The hardest problems are often the ones where the master almost succeeds but fails at the very last step. Because the standard teaching method discards these near-misses, the apprentice never learns how to navigate the specific decision point that caused the failure, leaving the AI unable to solve the most difficult scenarios.

Researchers at Apple have developed a new method called PROOF-Gen to fix this gap. Instead of throwing away the failed attempts, they treat them as valuable learning opportunities. When the master program fails a task, a second, even more advanced program acts as a reflective coach. This coach analyzes exactly where the master went wrong, looking at the sequence of actions and the final error. It then writes a specific set of instructions, or a "guide," designed to guide the master through that exact same problem again. The master program tries the task once more, this time following the coach's new advice. If it succeeds, that successful path is saved as a new lesson. Crucially, before the apprentice sees this new lesson, the guide is removed. The apprentice learns only from the clean, successful sequence of actions, not from the extra instructions that helped create it. This ensures the apprentice learns the skill itself, rather than memorizing a specific set of hints.

The results of this approach are striking. In tests using a benchmark designed to simulate customer service interactions, the standard method of discarding failures meant that the master program succeeded in only about 7 percent of cases. The new method recovered successful solutions for 93 percent of the tasks that the master had originally failed. When smaller AI models were trained on this expanded collection of lessons, their ability to complete tasks improved dramatically. One model jumped from solving 13 percent of tasks to solving 53 percent, while another saw a similar leap in performance. The researchers confirmed that this improvement came from learning the correct path through difficult decision points, not just from getting better at the easy parts of the task. In fact, training only on the easy, pre-selected successes improved the models' ability to make individual tool calls, but it did not help them finish the overall job. Only by including the recovered, difficult examples did the models learn to complete the full task.

This method also offers a practical advantage for companies that want to upgrade their AI systems without changing how the AI sounds or behaves. Because the final lessons are stripped of the coaching instructions, the AI learns the problem-solving capability of the powerful coach while retaining the natural voice and style of the original master program. The researchers tested this in a real-world production system used for customer service, where the method improved the success rate of the AI by 6.3 percentage points. These gains held true across dozens of different languages and regions, showing that the ability to recover from failure is a universal skill that does not depend on a specific language or cultural context. By turning failures into a structured learning process, the researchers have shown that the path to a smarter AI lies not just in finding more perfect examples, but in understanding and correcting the imperfect ones.

Technical Summary: PROOF-Gen

Problem Statement

The standard pipeline for distilling tool-calling capabilities into deployable models relies on Supervised Fine-Tuning (SFT) using trajectories generated by a frontier "teacher" model. This process typically employs a generate-and-filter mechanism: the teacher attempts a task, a verifier scores the trajectory, and only passing trajectories are retained for training.

The paper identifies a critical coverage gap in this approach:

Binary Filtering Discards Near-Misses: In execution-based verification, a trajectory that completes 90% of required tool calls correctly but fails on a single decisive action (e.g., a state-dependent decision point) is discarded entirely.
Lack of Signal from Failures: On the $\tau^2$ -bench, 57% of teacher trials fail, with two-thirds being "near-misses" where tool-call accuracy is high, but the final environment or database state is incorrect.
Ineffectiveness of Partial Credit: Training on filtered data or using turn-level partial rewards improves intermediate action accuracy but fails to improve end-to-end task completion. The student model never observes a complete, successful trajectory navigating the specific decision point where the teacher failed.
Stochastic Sampling Limitations: Simply sampling more trajectories or increasing temperature recovers stochastic errors but not strategic failures, which persist across resamples from the same teacher.

Methodology: PROOF-Gen

The authors introduce PROOF-Gen (Per-scenario Reflective Optimization to Overcome Failed Generation), a method that recovers "golden" trajectories from teacher failures rather than discarding them. The pipeline inverts the standard prompt optimization goal: instead of finding a single prompt that generalizes across tasks, it optimizes a prompt for a single scenario to force a passing trajectory, then discards the prompt.

The process consists of three main stages:

Per-Scenario Recovery Loop:
- For each task where the teacher (e.g., GPT-4o) fails, a reflector model (a stronger frontier model, e.g., GPT-5.1) analyzes the failed execution trace and verifier feedback.
- The reflector generates a "cheatsheet": a block of natural-language guidance (diagnoses, procedural steps, tool ordering) appended to the teacher's system prompt for that specific task.
- The teacher re-executes the task from scratch using this cheatsheet.
- If the teacher still fails, the reflector revises the cheatsheet based on the new trace and feedback. This loop repeats up to $K=10$ iterations until the teacher passes all evaluators.
- The process uses an iterative prompt optimizer (GEPA) to refine the cheatsheet.
Scaffold Removal:
- Before training, the cheatsheet is stripped from the trajectory.
- The student model is fine-tuned on the recovered trajectory alone. This ensures the student learns the capability demonstrated by the trajectory (the correct sequence of tool calls and state management) without inheriting the scaffold (the specific instructions or guidance used to generate it).
Training Data Composition:
- The final training set is the union of the standard filtered (originally passing) trajectories and the recovered (originally failing, now fixed) trajectories.
- This approach decouples capability (supplied by the reflector's problem-solving) from voice (supplied by the executor's interaction style), allowing teams to upgrade task coverage without disrupting the established product tone or latency profile.

Key Contributions

Recovery of Hard Scenarios: PROOF-Gen transforms the "generate-and-filter" pipeline into a "generate-recover-and-filter" pipeline, recovering trajectories that were previously discarded.
Per-Scenario Prompt Optimization: It introduces the use of automatic prompt optimization as a data-generation engine for single scenarios, rather than a general inference-time optimization tool.
Voice-Preserving Distillation: By keeping the executor (teacher) unchanged and only modifying its prompt via a reflector, the method allows the student to inherit advanced problem-solving capabilities without adopting the reflector's distinct interaction style or verbosity.
Configurable Objectives: The reflector can be directed to optimize for any combination of dimensions (e.g., groundedness, brevity, formatting) in addition to task completion.

Results

Recovery Effectiveness (RQ1)

High Recovery Rates: On the $\tau^2$ -bench telecom domain, per-scenario optimization recovered 93% of sampled teacher failures (279 out of 300).
Reflector Capability: Recovery rates are driven by the reflector's reasoning ability. GPT-5.1 achieved a 97.6% recovery rate on a cross-domain ablation set, compared to 24.4% for GPT-4o self-reflection.
Convergence: Most solvable scenarios (58.7%) were resolved on the first iteration, with 84.1% resolved by the third.

Downstream Impact (RQ2)

End-to-End Gains: Models trained on the combined dataset (filtered + recovered) showed consistent improvements in end-to-end task completion, whereas models trained on filtered-only data (the partial-credit baseline) did not.
- Qwen3-4B-Instruct-2507: Improved from Pass@1 = 0.132 to 0.529 on $\tau^2$ -bench telecom.
- Gemma 4 E4B-it: Gained +7.2 percentage points on BFCL v4 multi-turn.
Partial vs. Full Metrics: Filtered-only training improved action accuracy (+25pp) but failed to improve environment or database state accuracy. Combined training improved all three, closing the gap on the decisive decision points.
Stochastic Robustness: Gains persisted under stochastic sampling (temperature 1.0), with combined models achieving ~54% success on "flip tasks" versus near-zero for baselines.

Production Scaling (RQ3)

Deployment: The method was deployed in a production tool-calling agent post-training system.
Metrics: It improved trajectory quality by +6.3pp in goal completion and +2.5 to +8.0pp across five response-quality metrics.
Transfer: These gains transferred to a deployed on-device model (+1.5pp goal completion) with positive transfer observed in every locale evaluated (non-English average +1.48pp), indicating the corrective knowledge is largely language-agnostic.

Significance and Claims

The paper claims that the ceiling of distillation quality is set upstream at data generation time. If the hardest scenarios are systematically absent from the training distribution (because they are filtered out), no downstream training algorithm (e.g., RL, alternative divergences) can compensate.

PROOF-Gen addresses this by treating teacher failures as recoverable rather than terminal. By using a stronger reflector to guide a weaker executor to a passing state, the method generates clean demonstrations of complex problem-solving that the base teacher could not produce alone. The authors emphasize that this approach allows for capability–voice decoupling, enabling the integration of advanced reasoning from frontier models into smaller, deployed agents without altering the agent's established interaction style or latency profile.

The work is presented as a practical, scalable solution for production systems, validated across controlled benchmarks and a real-world deployment, demonstrating that recovering discarded failures is a necessary step to close the coverage gap in tool-calling agents.