DIAG: Diagnostic Iterative Alignment and Generation for Data-Efficient Mathematical Preference Distillation

Guhan Chen; Songtao Tian; Bohan Li; Hejin Wang; YeXin Xie; Zixiong Yu

cs.CLAug 24, 2026

DIAG: Diagnostic Iterative Alignment and Generation for Data-Efficient Mathematical Preference Distillation

Authors: Guhan Chen, Songtao Tian, Bohan Li, Hejin Wang, YeXin Xie, Zixiong Yu

Organizations: Tsinghua University · Kyoto University · Nanjing University

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.

Figures & tables

Figure 1: An illustrative example of DIAG. Left: The student directly compares likelihoods while ignoring prior probabilities, producing an incorrect posterior. Right: DIAG diagnoses the misconception from the student’s erroneous reasoning trace and generates a targeted follow-up with extreme priors that makes the error more salient, yielding more informative preference pairs for alignment.

Figure 2: The DIAG Framework for Iterative Preference Distillation. DIAG identifies the student’s competence boundary and uncertainty, and generates targeted practice to sustain informative preference pairs, improving both supervision quality and data efficiency for math reasoning alignment.

Method	Standard			Exam		Competition (Hard)			Avg.
Method	GSM8K	MATH 500	Minerva Math	Gaokao 2023en	College Math	Olympiad Bench	AIME24 (Avg@8)	AMC23 (Avg@8)	Avg.
Qwen2.5-Math-7B
Base (Start)	84.2	70.2	26.5	58.7	35.0	33.8	24.2	54.7	48.4
Numina $\rightarrow$ IDPO	85.4	70.4	29.8	59.0	37.8	35.9	22.5	56.6	49.7
Static Gen $\rightarrow$ IDPO	86.5	69.6	33.8	60.8	38.7	35.7	22.1	59.7	50.9
DIAG (Ours)	87.1	70.8	34.9	61.0	41.9	38.4	25.0	63.1	52.8

Table 1: Main Results on Mathematical Reasoning (Fixed Total Budget). We compare four data sources under the same total student-sample budget : ( Base ) the starting checkpoint, ( Numina ) IDPO on a static open-source pool, ( Static Gen ) IDPO on a fixed teacher-generated pool conditioned on (topic, difficulty), and DIAG (Ours) IDPO on online teacher generation conditioned on (failed question, wrong answer, topic, difficulty).

Task	Sampling Strategies
Task	w/o Explore	w/o EB	Ours
GSM8K	87.6	86.4	87.1
MATH500	70.4	70.6	70.8
MinervaMath	35.7	35.3	34.9
Gaokao2023	59.5	59.7	61.0
Olympiad	37.5	38.7	38.4

Table 2: Ablation on Sampling / Scheduling. We ablate the uniform exploration stream ( w/o Explore ) and EB smoothing ( w/o EB ) under a fixed rollout budget.

Method / Variant	Avg.	$\Delta$
Full DIAG	52.8	–
Impact of conditioning on failed question
w/o Wrong-QA (no failed context)	51.3	-1.5
Impact of conditioning on wrong answer/trace
w/o Wrong-A (no wrong answer)	52.1	-0.7

Table 3: Ablation on Teacher Conditioning Signals.

\Delta

denotes performance drop relative to the full method.

Method	BBH	Human Eval	MMLU	TruthfulQA
Method	BBH	Human Eval	MMLU	MC1	MC2
Qwen3-8B-Base	75.5	64.6	76.9	33.5	51.1
+ Static Gen	75.6 $\,{}^{\uparrow}$	64.6	76.9	33.6 $\,{}^{\uparrow}$	51.2 $\,{}^{\uparrow}$
+ DIAG	75.6 $\,{}^{\uparrow}$	64.7 $\,{}^{\uparrow}$	76.9	33.9 $\,{}^{\uparrow}$	51.0

Table 4: Cross-Task Generalization and Alignment Tax Evaluation. Bold denotes the best result; ↑ indicates improvement over the base model.

Figure 3: Effective sample yield across iterations. DIAG maintains a higher signal yield

\rho_{t}

than baselines and avoids degeneration into uninformative regimes.

Figure 4: Pass-rate distribution. Our method concentrates more mass in the intermediate region, reducing both all-correct and all-wrong degeneracy.

Method	Eff. Samples	Avg.	$\Delta$
Qwen2.5-Math-7B	–	48.4	–
+ SPIN Chen et al. 2024d	8k	48.9	+0.5
+ NuminaMath Li et al. 2024a	8k	49.0	+0.6
+ Static IDPO Zhang et al. 2025	8k	49.7	+1.3
+ LLM2LLM Lee et al. 2024	8k	50.1	+1.7
+ Ours	8k	50.7	+2.3

Table 5: Iso-effective comparison. All methods use

\sim

8k effective instances and the same gradient steps. LLM2LLM-style adopts error-driven generation but trains with DPO for fair comparison.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Prompt for Student Model and Teacher Model Answer Generation.

Panel of Figure 6 — Figure 6: Prompt for Teacher Model Question Generation.

Panel of Figure 7 — Figure 7: Case Studies. Top pair (Improper Integral): The student applies FTC across a singularity at $x=0$ , arriving at $-2$ ; the follow-up shifts the singularity to $x=2$ , where the trap value $-3/2$ is negative for a strictly positive integrand. Bottom pair (Bayesian Inference): The student normalizes likelihoods without prior weighting ( $5/8$ instead of $5/7$ ); the follow-up pushes priors to $1\%$ vs. $99\%$ , where the same shortcut gives $62.5\%$ but the correct posterior is $\approx 1.66\%$ .

LLM Alignment

Large language model alignment studies methods for making models follow human intentions, values, and safety constraints across contexts.

See all 613 papers

The Quality-Utility Paradox: Why High-Reward Data Impairs Small Model Mathematical Reasoning
Jun 15, 2026Haolong Qian, Xianliang Yang, Yinuo ma +6CoT Distillation Language Model Distillation
Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China · Microsoft Research Asia, Microsoft, Beijing, China · Faculty of Computer Science and Artificial Intelligence, Shenzhen University of Advanced Technology, Shenzhen, China +1
INSPIRE: An Internalize-Then-Improve Approach for Example-Driven Mathematical Reasoning
Aug 27, 2026Shuai Wang, Jiayi Kuang, Yinghui Li +4LLM Mathematical Reasoning Mathematical Reasoning
Sun Yat-sen University · Tsinghua University · Beihang University +1
HiPO: Hierarchical Preference Optimization for Adaptive Reasoning in LLMs
Apr 22, 2026Darsh Kachroo, Arjun Prasaath Anbazhagan, Adriana Caraeni +2LLM Alignment Direct Preference Optimization
Vellore Institute of Technology, Chennai · Northwestern University · University of Massachusetts Amherst +2

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

The Quality-Utility Paradox: Why High-Reward Data Impairs Small Model Mathematical Reasoning

Haolong Qian, Xianliang Yang, Yinuo ma, Lirong Che+5

Tsinghua Shenzhen International Graduate School, Tsinghua University, Shenzhen, China · Microsoft Research Asia, Microsoft, Beijing, China · Faculty of Computer Science and Artificial Intelligence, Shenzhen University of Advanced Technology, Shenzhen, China+1

CoT Distillation Language Model Distillation Rejection Sampling LLM Mathematical Reasoning

Open paper AskPDF arXiv

Knowledge distillation from powerful reasoning models is widely used to improve Small Language Models (SLMs) on mathematical reasoning, often assuming that traces with higher reward model scores provide more useful supervision. We identify a counterintuitive Quality-Utility Paradox in mathematical reasoning distillation. Data refined or synthesized by a stronger Oracle obtains higher perceived quality according to reward models, yet consistently underperforms traces generated by the SLM itself and selected through rejection sampling across Qwen2.5, LLaMA-3, and DeepSeek families. Our analysis shows that Oracle refinement couples logical repair with distributional drift away from the SLM's native reasoning distribution. This drift increases the learner's adaptation cost and can outweigh the benefit of improved reasoning logic. To test this mechanism, we introduce Style-Aligned Refinement, which preserves the native trajectory of the SLM while retaining logical repair from the Oracle. This intervention lowers adaptation cost and restores downstream utility. These findings suggest that effective mathematical reasoning distillation should jointly optimize perceived solution quality and learner-data compatibility, rather than relying solely on reward-model scores. The datasets and code are available at https://github.com/Dracoqhl/Quality-Utility-Paradox.

Abstract

LLM Alignment

Explore similar work

The Quality-Utility Paradox: Why High-Reward Data Impairs Small Model Mathematical Reasoning

INSPIRE: An Internalize-Then-Improve Approach for Example-Driven Mathematical Reasoning

HiPO: Hierarchical Preference Optimization for Adaptive Reasoning in LLMs