Source-linked AI summary
DIAG: Diagnostic Iterative Alignment and Generation for Data-Efficient Mathematical Preference Distillation
Guhan Chen, Songtao Tian, Bohan Li, Hejin Wang, YeXin Xie, Zixiong Yu
TL;DR
Static practice pools can become mismatched to improving students, creating too few informative preference pairs. DIAG diagnoses yield to control exploration and topic priorities, then generates teacher-written variants from failure traces; experiments report improved data efficiency and reasoning performance, including under iso-effective training.
Problem
Static practice pools become increasingly mismatched as students improve, producing too few informative preference pairs for iterative mathematical reasoning alignment.
Method
DIAG diagnoses signal yield, allocates exploration and topic quotas with Empirical Bayes shrinkage, and generates targeted practice conditioned on student failure traces.
Results
DIAG consistently improves data efficiency and final reasoning performance across diverse students and challenging math evaluations, with gains persisting under iso-effective training.
Takeaways & Limitations
DIAG can distill more informative preference supervision by reshaping practice near the student’s moving competence boundary.
Takeaways & Limitations
DIAG relies on a capable teacher for question synthesis, and its evaluation is limited to mathematical reasoning with reliable verifiers.
Abstract
from arXiv · showhide
Iterative preference optimization is essential for aligning Large Language Models on mathematical reasoning tasks, yet its efficiency is often throttled by signal scarcity: as the model improves, static problem sets become increasingly mismatched to the model's evolving competence, producing rollouts that are either too easy or too hard and therefore non-informative, which leads to a scarcity of valid preference pairs. We propose DIAG, a Diagnostic Iterative Alignment and Generation framework that adaptively reshapes the practice distribution to increase informative supervision and focus training near the student's current competence boundary. DIAG consists of two phases: (1) diagnosing valid preference-pair yield to calibrate the exploration-exploitation trade-off and allocate topic quotas via an Empirical Bayes shrinkage estimator, thereby prioritizing high-yield concepts; and (2) generating targeted practice, where a teacher synthesizes variants from the student's failure traces. We further provide a theoretical view interpreting DIAG as a teacher-mediated approximation to KL-regularized reweighting of the practice distribution toward the student's competence boundary, where valid preference-pair yield is maximized. Experiments show that DIAG boosts yield across iterations and delivers stronger reasoning performance under an iso-effective training budget, demonstrating that it can distill more informative preference supervision for mathematical reasoning.
1 Introduction
Static practice distributions can become mismatched to evolving model capabilities, reducing informative preference-pair yield. DIAG responds by diagnosing signal scarcity, adapting topic allocation, and generating mistake-conditioned practice near the student’s decision boundary.
- Static training sets may be too easy, too difficult, or misleading for models at different capability levels, reducing supervision yield.
- Fixed practice distributions can produce too few informative preference pairs even when rollouts are generated on-policy.
- Curriculum feedback in iterative preference optimization is sparse, noisy at low sample counts, and nonstationary across iterations.
- DIAG diagnoses usable signal, balances exploration and exploitation, allocates topic budgets with Empirical Bayes shrinkage, and generates targeted follow-ups from student mistakes.
- DIAG’s mistake-conditioned generation concentrates practice near the student’s decision boundary, yielding higher-quality preference pairs and greater information density per unit budget.
2 Related Work
Prior work improves iterative preference optimization through update rules, pair construction, or supervisory distributions on fixed tasks. DIAG instead reshapes the data distribution by synthesizing questions from erroneous reasoning traces and targeted failure modes.
- Recent iterative preference methods focus on update rules, pair construction, or dynamically selecting data as reasoning abilities evolve.
- On-policy distillation methods transfer teacher knowledge through token-level supervision while operating on fixed tasks or prompts.
- DIAG differs by analyzing erroneous student reasoning traces and synthesizing new questions that force the student to confront the diagnosed misconception.
- Existing data-synthesis approaches generally overlook targeted feedback in their generation paradigms.
- Compared with failure-conditioned approaches that use failed questions or trajectories, DIAG conditions synthesis on the student’s erroneous reasoning trace.
3 Methods
DIAG iteratively diagnoses preference-pair yield, allocates exploration and exploitation across topics, and generates targeted variants from boundary mistakes. Its theoretical interpretation links this process to reweighting practice toward prompts with maximal valid-pair yield.
- Overview: DIAG replaces a fixed prompt pool with a Diagnose-and-Generate loop that reshapes the next iteration’s practice distribution from student failures.
- Phase I: Diagnose & Allocate: Phase I estimates prior yield, sets the exploration–exploitation split, and assigns topic quotas using Empirical Bayes-smoothed utility estimates.
- Natural Preference Signal: Natural preference pairs require rollouts containing at least one correct and one incorrect response, while degenerate pools can cause gradient starvation.
- Phase I: Diagnose & Allocate: Exploration increases when informative pairs are scarce and decreases when the current curriculum already produces abundant contrastive signal.
- Phase II: Generate Practice: Phase II samples recent boundary mistakes and asks a teacher to generate variants preserving the concept and difficulty while changing surface form.
- Theoretical interpretation: For rollout width K, valid-pair probability is 1 − p^K − (1 − p)^K, peaking at p = 1/2 and vanishing for all-correct or all-wrong rollouts.
4 Experiments
DIAG is evaluated across multiple student models and mathematical benchmarks under matched student-sample budgets. It achieves the strongest average performance, with gains especially visible on harder testbeds, while ablations show benefits from exploration, Empirical Bayes smoothing, and richer error conditioning.
- Evaluation Setup: Experiments cover eight mathematical benchmarks and three open-weight student models spanning multiple capability profiles.The evaluation includes GSM8K, MATH500, MinervaMath, Gaokao2023En, College Math, OlympiadBench, AIME24, and AMC23.
- Evaluation Setup: All methods share the same total student-sample budget across seven IDPO iterations with 3 × 10^3 prompts per iteration.The decoding protocol is held identical across methods, with Qwen3-235B serving as teacher.
- Main Results: DIAG achieves the strongest average performance across all three student models under identical student-sample budgets.The gains are most pronounced on harder testbeds, where generic practice pools more often enter degenerate regimes and reduce natural preference supervision.
- Ablation Study: AMC23 performance falls from 63.1 to 58.1 without exploration, while removing Empirical Bayes smoothing produces broad degradation with Avg 51.8.These ablations indicate that exploration supports robustness on competition-style evaluation and EB smoothing stabilizes low-sample topic scoring.
- Ablation Study: Removing failed-context conditioning drops performance by 1.5, whereas removing the wrong answer drops it by 0.7 relative to full conditioning.The richer signal combines topic, difficulty, failed question, and wrong answer.
- Cross-Task Generalization: DIAG maintains comparable knowledge and truthfulness performance while showing mild positive transfer on reasoning and code-generation tasks.This cross-task evaluation is presented as evidence that the method does not incur an alignment tax.
5 Analysis
DIAG improves effective preference-pair yield by moving practice toward the student’s competence boundary. Its advantage persists when effective training instances and gradient steps are held constant, supporting higher supervision quality rather than merely greater sample quantity.
- Effective Sample Yield over Training: Effective yield ρ_t is the fraction of questions whose K rollouts contain both correct and incorrect outcomes.This measures the availability of valid natural preference pairs across iterations.
- Effective Sample Yield over Training: DIAG increases ρ_t from ∼0.59 to ∼0.83, a +24 pp improvement over 10 iterations, while Static Gen and static Numina decline to around 0.55.The result indicates that targeted generation keeps practice near the student’s moving competence boundary.
- Pass-Rate Distribution: DIAG concentrates pass-rate mass in the intermediate region, reducing all-correct and all-wrong regimes associated with wasted budget and gradient starvation.The analysis interprets mistake-conditioned generation as improving difficulty matching rather than simply increasing difficulty.
- Iso-Effective Comparison: Under ∼8k effective instances and identical gradient steps, DIAG remains consistently better than comparison methods.The iso-effective comparison separates supervision quality from improvements attributable purely to data quantity.
- Iso-Effective Comparison: SPIN improves by +0.5 and LLM2LLM-style by +1.7, while DIAG adds wrong-reasoning conditioning and EB-guided scheduling.The paper links these additions to more targeted misconception follow-ups and improved signal yield.
6 Conclusion
The paper identifies signal scarcity as the bottleneck in iterative mathematical preference distillation and proposes DIAG to reshape practice as the student improves. Across evaluations, DIAG improves data efficiency and reasoning performance, including under iso-effective training.
- Conclusion: DIAG reshapes the practice distribution through signal-aware pacing, Empirical-Bayes topic prioritization, and mistake-conditioned generation.Together, these components form its Diagnose-and-Generate framework.
- Conclusion: DIAG consistently improves data efficiency and final reasoning performance across diverse students and challenging mathematical evaluations.The gains persist under iso-effective training, indicating higher-quality supervision beyond increased signal yield.
Limitations
DIAG’s Diagnose-and-Generate loop has two stated limitations: it depends on a capable teacher for question synthesis and has only been evaluated in mathematically verifiable settings.
- DIAG relies on a capable teacher model to synthesize targeted follow-up questions, which may be unavailable or unaffordable in some settings.
- The evaluation is restricted to mathematical reasoning with reliable verifiers, leaving transfer to domains without automatic correctness checks unclear.
Ethics Statement
The work reports compliance with the ACL Ethics Policy and uses publicly available resources without human subjects or private data. It also cautions that model outputs may remain incorrect or misleading and should be independently verified in high-stakes settings.
- Experiments use publicly available datasets and widely accessible models, with citations provided throughout.
- The method involves no human subjects, private data, or annotation beyond automated answer verification.
- LLMs trained with DIAG may still produce incorrect or misleading reasoning steps.
- Users should independently verify model outputs before relying on them in high-stakes or real-world decisions.
- The theoretical view places useful preference pairs near the student’s competence boundary and interprets DIAG as practice-distribution reweighting toward high-yield regions.
A.1 Signal Yield Peaks at the Competence Boundary
The theoretical analysis shows that valid preference-pair yield is highest near the student’s competence boundary, while degenerate prompts are inefficient. DIAG approximates yield-focused distribution reweighting using topic statistics and mistake-conditioned generation, with empirical conditions governing its improvement.
- Valid-pair yield is uniquely maximized at pass probability p = 1/2, while all-correct and all-wrong prompts provide low usable supervision.
- With M sampled prompts, the expected number of effective DPO training instances is M Jθ,K(µ).
- DIAG approximates ideal KL-regularized exponential reweighting toward high-yield prompts using topic-level pair statistics and local mistake contexts.
- Mistake contexts come from recent valid-pair prompts, whereas degenerate all-correct or all-wrong prompts affect pacing but are excluded as exploitation anchors.
- DIAG improves expected yield when mistake-conditioned generation produces fewer degenerate follow-up prompts than metadata-only generation on average.
- Empirical-Bayes topic scores shrink low-count yield estimates toward the global mean, then increasingly reflect topic-specific rates as sample counts grow.
- DPO updates are locally most sensitive when preferred and dispreferred responses have equal model preference, supporting practice near the moving competence boundary.
B Complexity Analysis
DIAG frames computational efficiency around the yield of valid preference pairs rather than raw rollout or teacher-generation volume. By maintaining higher yield through targeted practice and mistake-conditioned generation, it produces more effective supervision under matched budgets.
- Cost and effective supervision: Valid preference-pair yield determines how many student rollouts provide contrastive supervision, because a usable pair requires both correct and incorrect responses.When static or practice-agnostic distributions produce all-correct or all-wrong rollouts, yield falls and computation is wasted.
- Cost and effective supervision: DIAG and competing methods use the same total teacher-generated-question budget; their difference is how that budget reshapes practice.Static generation uses topic and difficulty, whereas DIAG conditions online generation on recent student failures.
- Cost and effective supervision: Higher ρt gives DIAG a lower compute cost per effective preference pair and more gradient signal under the same rollout budget.The student rollout budget is fixed, so the dominant efficiency factor is valid-pair yield.
- Empirical Bayes allocation: The Empirical Bayes topic score models each topic’s valid-pair probability and shrinks low-count estimates toward the global rate.As topic observations increase, the empirical topic rate dominates; with few observations, shrinkage limits reactions to noise.
- Prompted generation: Teacher prompts incorporate the failed question and erroneous reasoning trace to synthesize variants targeting the same misconception while changing surface features.Student and teacher prompt templates support the Diagnose-and-Generate loop.
E.1 Training Configuration
The experiments use DPO training with fixed optimization settings, diverse student rollouts, and targeted teacher generation. Case studies illustrate how generated variants preserve misconceptions while exposing them under changed surface conditions.
- Training and decoding: All student models use DPO with AdamW, peak learning rate 5 × 10−7, global batch size 128, and 2 epochs per iteration.The setup also uses β = 0.1, label smoothing 0.1, frozen previous checkpoints as reference policies, and 4096-token sequence truncation.
- Data and verification: The data taxonomy spans Algebra, Geometry, Calculus, Discrete Mathematics, and Probability, with fine-grained sub-topics and four difficulty levels.Failure traces are sampled from the mistake buffer using uncertainty-weighted selection.
- Training and decoding: Student rollouts sample K = 8 responses per prompt at temperature 1.0, top-p 0.99, and top-k 50, while teacher question generation uses temperature 1.0 and top-p 0.99.Teacher answer generation uses greedy decoding.
- Teacher generation: The teacher is instructed to generate rigorous, original questions aligned with topic and difficulty that directly target student weaknesses and yield numerical or formula answers.The prompt explicitly frames the teacher as a curriculum designer for StudentLM.
- Case studies: Shifting an improper-integral singularity from x = 0 to x = 2 preserves the continuity misconception while defeating a location-specific heuristic.The naive evaluation gives −3/2, whereas recognizing divergence yields 0, creating a binary rollout split.
- Case studies: With extreme priors of 1% versus 99%, ignoring priors still gives 5/8 ≈62.5%, while the correct posterior is approximately 1.66%, a nearly 40× discrepancy.The contrast makes the Bayesian oversight more visibly wrong than in the original setting.