DIAG: Diagnostic Iterative Alignment and Generation for Data-Efficient Mathematical Preference Distillation
2026-08-24 • Computation and Language
Computation and Language
AI summaryⓘ
The authors address a problem in training language models for math reasoning: as models get better, existing practice problems become too easy or too hard, making it hard to get useful feedback. They propose a method called DIAG that adjusts the practice problems based on the model's current skill level, focusing on areas where the model struggles the most. DIAG has two parts: figuring out which topics give the best learning signals, and creating new problems targeting the model’s mistakes. Their experiments show that this method helps models learn more effectively without needing extra training time.
Large Language ModelsIterative Preference OptimizationMathematical ReasoningExploration-Exploitation Trade-offEmpirical BayesPreference LearningCompetence BoundaryKL-Regularized ReweightingDiagnostic Iterative AlignmentActive Learning
Authors
Guhan Chen, Songtao Tian, Bohan Li, Hejin Wang, YeXin Xie, Zixiong Yu
Abstract
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.