PROOF-Gen: From Optimized Data to Better Distillation

Anh Ta; Junjie Zhu; Shahin Shayandeh

cs.AIAug 24, 2026

PROOF-Gen: From Optimized Data to Better Distillation

Authors: Anh Ta, Junjie Zhu, Shahin Shayandeh

Organizations: Apple

Abstract

Supervised fine-tuning on teacher-generated trajectories is the standard first stage for distilling tool-calling capabilities into deployable models. Post-training pipelines that drive shipped tool-calling agents re-run this stage on a daily or weekly cadence, paying the frontier-teacher cost each cycle, yet the mechanism is generate-and-filter (keep the teacher's passing trajectories, discard the rest) and each cycle leaves behind the same hard scenarios because failures supply no signal. On τ2-bench, 57% of teacher trials fail, two-thirds of them near-misses (most tool calls correct, undone by one decisive error). We introduce PROOF-Gen (Per-scenario Reflective Optimization to Overcome FailedGeneration), which recovers golden trajectories from these failures via per-scenario prompt optimization. For each failed task, a reflector analyzes the execution trace and evaluation feedback, then writes corrective guidance that steers the teacher to a passing trajectory. The guidance is stripped before training, so the student learns from clean demonstrations with no task-specific scaffold. On τ2-bench, per-scenario optimization recovers 93% of failed scenarios. Fine-tuned on the combined data, Qwen3-4B-Instruct-2507 improves from Pass^1=0.132 to 0.529 and Gemma 4 E4B-it gains +7.2pp on BFCL v4 multi-turn. In a deployed pipeline, the method lifts trajectory quality by +6.3pp goal completion and transfers to a deployed on-device model (+1.5pp goal completion; +1.7 to +5.0pp across response-quality metrics), with positive transfer in every locale (non-English average +1.48pp).

Figures & tables

Panel of Figure 1 — Figure 1: Training pipeline. For each task the teacher originally failed, a per-scenario optimization loop (green) iteratively refines a cheatsheet from the trajectory’s feedback until the teacher produces a passing (“golden”) trajectory. The cheatsheet is stripped at scaffold removal, so the student trains only on what the trajectory demonstrates; the fine-tuning procedure is otherwise unchanged.

Panel of Figure 2 — Figure 2: The accuracy asymmetry in failed GPT-4o trajectories on $\tau$ 2-bench (278 tasks, 3 trials each; 473 of 834 trials fail). Most failures execute over half of tool calls correctly (high tool-call accuracy), yet a single missed action cascades into an incorrect end state (environment and database failure). Because tool-call accuracy is already high in failures, turn-level or partial-credit signals saturate without lifting end-to-end completion; what is missing is a complete trajectory through the decisive step. We confirm this in Section 4 .

	$\tau$ 2-bench	BFCL v4
Eval tasks	114	240
User simulation	Stochastic	Deterministic
Primary metric	Pass^1	Task accuracy
Secondary metric	Partial reward	Partial reward

Table 1: Benchmark comparison.

\tau

2-bench uses an LLM user simulator; BFCL uses pre-scripted user turns.

Generate-and-filter (teacher)
	$\tau$ 2-bench	BFCL v4
	telecom	multi-turn
Training pool (tasks)	2,171	560
Filtered (passed)	158	296
Failed	2,013	264
Teacher pass rate	7.3%	52.9%

Table 2: Data-generation pipeline for both benchmarks. The teacher is GPT-4o (temperature 0, one trajectory per task). Telecom recovery runs on a 300-failure sample of the 2,013 failures; all 264 BFCL failures are attempted. Recovery rate = recovered / failures sent to optimizer. The combined set merges filtered and recovered trajectories; “from recovery” is the recovered fraction of that set.

Figure 3: Reflector capability dominates recovery rate (range 24.4%–97.6%). Per-scenario recovery by reflection model on a cross-domain ablation set of 82 failed

\tau

2-bench scenarios (13 airline, 14 retail, 55 telecom), holding the executor (GPT-4o) and optimizer fixed. GPT-5.1 recovers 80/82 (97.6%) versus 20/82 (24.4%) for GPT-4o self-reflection; the largest gap is on telecom (54/55 vs. 4/55). This ablation isolates the effect of reflector capability and is not the source of the production-scale headline rate (93%, 279/300 telecom).

Figure 4: Post-training results on

\tau

2-bench telecom (left) and BFCL v4 multi-turn (right). Solid bars show the primary metric (Pass^1, task accuracy; left axis); hatched bars show partial rewards (action accuracy, turn-level partial reward; right axis). Dashed lines mark the GPT-4o teacher. Filtered-only training improves partial rewards but does not improve end-to-end task completion in any of the four model–benchmark pairs; combined training produces the only consistent gains.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Three failed GPT-4o trajectories from

\tau

2-bench telecom. Green circles indicate tool calls that matched the ground truth; red circles indicate missed calls. Each trajectory completes most required actions but fails at one decision point, causing the entire trajectory to be discarded by the filter.

Figure 6: Cumulative recovery by optimizer iteration on 300 failed telecom tasks (GPT-4o agent, GPT-5.1 reflector). 264 of 279 total recoveries are cheatsheet-driven; the remaining 15 are stochastic replays. Most solvable scenarios resolve within 2–3 iterations.

Agent	Reflection	Recovered	Rate
GPT-5.1 (no reasoning)	none	44/300	14.7%
GPT-5.1 (high)	none	250/300	83.3%
GPT-4o	GPT-5.1 (high), 10 iter	279/300	93.0%
GPT-5.1 (high)	GPT-5.1 (high), 1 iter	295/300	98.3%

Appendix

Table 3: Recovery rate on 300 failed telecom tasks (selected from GPT-4o failures by design). GPT-5.1 defaults to reasoning_effort=none (no reasoning); “high” = reasoning_effort=high. Per-scenario optimization with a weaker agent (GPT-4o) and capable reflector recovers 93.0% using the expensive model only for reflection (avg 2.2 calls per scenario vs.

{\sim}

11.5 execution turns per task for direct generation).

Figure 7: Distribution of per-task pass rates at temperature 1 (10 trials) on

\tau

2-bench flip tasks. Boxes show quartiles; black lines mark means. Baseline rates cluster near zero; combined rates are broadly distributed, confirming that training gains persist under stochastic sampling.

$\tau$ 2-bench (60 samples, 15 per teacher)
Reflection Model	Opus 4.6	Sonnet 4.6
GPT-4o	15/15	14/15
GPT-5-mini	15/15	15/15
GPT-4.1	15/15	15/15
GPT-5.1	15/15	15/15
BFCL v4 (30 samples, GPT-5.4 teacher)

Appendix

Table 4: Cheatsheet integrity audit. Each cell shows the number of trajectories rated Transferable (all discovery steps preserved). Sonnet’s two flags are false positives on manual inspection (Appendix C ).

Experiment	Action	Env	DB
Baseline (no SFT)	41.4%	23.5%	9.3%
+ QLoRA Filtered	66.3%	22.6%	10.9%
+ QLoRA Combined	83.3%	54.5%	15.6%

Appendix

Table 5: Reward breakdown on

\tau

2-bench telecom (base-114 held-out tasks, 3 trials, temp=0). Filtered-only training improves action accuracy but not environment or database accuracy. Combined training improves all three, with the largest gain in environment assertions.

Metric	GPT-4o	GPT-5.1 (high)
Trajectories	259	259
Text responses	3,180	2,863
Avg. length (chars)	308	987
Median length (chars)	253	965
Tool calls / trajectory	5.2	5.2

Appendix

Table 6: Text response characteristics on 259 matched telecom tasks. GPT-4o uses cheatsheet guidance; GPT-5.1 uses direct high reasoning. Tool call counts are comparable; response length differs by 3.8

\times

Figure 8: Median text response length on matched tasks. GPT-4o without scaffold (n=15 tasks) and GPT-4o with cheatsheet (n=259) produce comparable response lengths, while GPT-5.1 responses are 3.8

\times

longer. Whiskers show IQR.

Parameter	Value
LoRA rank	32
LoRA alpha	32
Quantization	4-bit (QLoRA)
Learning rate	2e-4
LR scheduler	cosine
Warmup steps	20

Appendix

Table 7: QLoRA training hyperparameters (shared across both models).

Reflector	Domain	Improved	Replay	Total	Rate
GPT-4o	airline	4	1	5/13	38.5%
	retail	9	2	11/14	78.6%
	telecom	3	1	4/55	7.3%
	overall	16	4	20/82	24.4%
GPT-5-mini	airline	9	1	10/13	76.9%
	retail	9	3	12/14	85.7%

Appendix

Table 8: Recovery breakdown for the 82-task reflector ablation. Improved : cheatsheet enabled the teacher to pass. Replay : teacher passed on stochastic re-execution without a cheatsheet. Replay accounts for 4–7 of the total recoveries per reflector; the majority of recoveries are cheatsheet-driven.

After tool	GPT-4o	GPT-5.1	Ratio
	(+cheatsheet)	(direct)
get_details_by_id	469	997	2.1 $\times$
enable_roaming	292	1,254	4.3 $\times$
get_data_usage	476	1,450	3.0 $\times$
refuel_data	206	835	4.1 $\times$
All text responses	308	987	3.2 $\times$

Appendix

Table 9: Average text response length (characters) after each tool call type on 259 matched tasks. The ratio is consistent across tool types (2.1–4.3

\times

), confirming the style difference is systematic rather than specific to particular actions.

Gemma 4 E4B-it
Parameters	4.5B effective (8B total)
Architecture	Multimodal (text, image, audio)
Layers / Hidden	42 / 2,560
Attention	8Q / 2KV, hybrid sliding+full
Context	128K tokens
Tool format	Gemma special tokens

Appendix

Table 10: Student model comparison. Both share a hidden dimension of 2,560 and operate at comparable effective scale but differ in architecture, modality, and tool-call format.

$\tau$ 2-bench telecom (2,171 non-base tasks)
	Filtered	Recovered	Combined
Trajectories	158	279	437
Unrolled examples	2,049	4,814	6,863
Recovered:filtered	1.77:1
BFCL v4 multi-turn (560 train tasks)
base	96	11	107

Appendix

Table 11: Training data composition. Filtered = tasks the teacher passed. Recovered = tasks recovered via per-scenario optimization. For BFCL, the 264 failures (560 train tasks minus 296 passed) yield 89 recoveries, a 33.7% recovery rate.

	Gemma 4 E4B-it		Qwen3-4B-Inst.
Category	Base	+Comb.	Base	+Comb.
base	17.2	23.3	30.0	28.9
long_context	15.0	18.3	28.3	21.1
miss_func	0.0	15.6	11.7	31.7
miss_param	17.8	21.7	24.4	20.6
Overall	12.5	19.7	23.6	25.5

Appendix

Table 12: BFCL per-category accuracy (%, 60 test tasks each, averaged across 3 runs). Bold indicates improvement over baseline.

Overall accuracy (%, 240 test tasks)
Condition	Run 1	Run 2	Run 3	Spread
Gemma baseline	12.5	12.5	12.5	0.0pp
Gemma + combined	20.4	19.6	19.2	1.2pp
Qwen baseline	23.8	23.8	23.3	0.5pp
Qwen + combined	25.8	26.2	24.6	1.6pp
Per-category max spread (60 test tasks each)

Appendix

Table 13: BFCL v4 repeated evaluation stability. Overall accuracy is stable across runs (spread

\leq

1.6pp). Per-category results at

n{=}60

vary by up to 3.3pp from serving non-determinism alone. The miss_func gain is the most stable signal: Gemma 0.0%

\times

3 (baseline) vs. 15.0–16.7% (combined); Qwen 11.7%

\times

3 (baseline) vs. 31.7%

\times

3 (combined).

Run	Description	Pass^1
Original	Train + eval	0.529
Re-eval	Same checkpoint, new eval	0.518
Retrain	New training run + eval	0.538

Appendix

Table 14: Qwen3-4B-Instruct-2507 combined training stability on

\tau

2-bench telecom (114 tasks, 3 trials, temp=0). Pass^1 ranges from 0.518 to 0.538 across independent training and evaluation runs (spread: 0.020), confirming that the reported improvement is not an artifact of a favorable seed.

LLM Agent Training

Large language model agent training teaches models to plan, use tools, and learn from multi-step interaction with environments.

See all 586 papers

Trajectory-Refined Distillation
Jun 7, 2026Li Jiang, Haoran Xu, Yichuan Ding +1On-Policy Distillation LLM Post-Training
McGill University · Mila Quebec AI Institute · UT Austin
Persistent Teacher Anchoring for Tool-Using Agents
Sep 4, 2026Hyun Bin Park, Kyungho Song, Sangmin Lee +1LLM Agent Training On-Policy Distillation
Sogang University · University of Michigan, Ann Arbor
PivotOPD: Learning to Recover from Pivotal Mistakes in Multi-Turn Agents
Sep 30, 2026Yinghui He, Yapei Chang, Khushi Bhardwaj +4Language Model Distillation LLM Agent Training
Princeton University · NVIDIA · University of Maryland

Jun 7, 2026J/K move · Enter open · S save

Trajectory-Refined Distillation

Li Jiang, Haoran Xu, Yichuan Ding, Amy Zhang

McGill University · Mila Quebec AI Institute · UT Austin

On-Policy Distillation LLM Post-Training Teacher-Student Learning Knowledge Distillation

Open paper AskPDF arXiv

On-policy distillation (OPD) has become a central post-training tool for large language models (LLMs), providing dense per-token teacher supervision along the student's own rollouts. In this work, we identify a common structural cause underlying OPD, which we call prefix failure. Under prefix failure, dense per-token supervision induces a bimodal teacher mixture and fragmented gradients that token-level loss truncation or reweighting fail to address. We thus propose Trajectory-Refined Distillation (TRD), a trajectory-level correction method that revises the student's rollout under the teacher guidance while within on-policy support. By correcting problematic prefixes before distillation, TRD mitigates prefix failure at its source. Moreover, TRD improves exploration by exposing the student to alternative valid derivations under teacher guidance, even when the original rollouts are already correct. TRD can also be applied to on-policy self-distillation (OPSD), a parameter-sharing variant that uses the student model conditioned on privileged information as the teacher. TRD improves average math accuracy while remaining competitive on code. Under OPSD, it solves

9/23

AMOBench questions left unsolved by the Qwen3-8B base model, nearly doubling the strongest dense-KL baseline. Code is available at https://github.com/louieworth/trd.

Abstract

LLM Agent Training

Explore similar work

Trajectory-Refined Distillation

Persistent Teacher Anchoring for Tool-Using Agents

PivotOPD: Learning to Recover from Pivotal Mistakes in Multi-Turn Agents