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

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

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.
MethodStandardExamCompetition (Hard)Avg.
GSM8KMATH 500Minerva MathGaokao 2023enCollege MathOlympiad BenchAIME24 (Avg@8)AMC23 (Avg@8)
Qwen2.5-Math-7B
Base (Start)84.270.226.558.735.033.824.254.748.4
Numina →\rightarrow IDPO85.470.429.859.037.835.922.556.649.7
Static Gen →\rightarrow IDPO86.569.633.860.838.735.722.159.750.9
DIAG (Ours)87.170.834.961.041.938.425.063.152.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).
TaskSampling Strategies
w/o Explorew/o EBOurs
GSM8K87.686.487.1
MATH50070.470.670.8
MinervaMath35.735.334.9
Gaokao202359.559.761.0
Olympiad37.538.738.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 / VariantAvg.Δ\Delta
Full DIAG52.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.
MethodBBHHuman EvalMMLUTruthfulQA
MC1MC2
Qwen3-8B-Base75.564.676.933.551.1
+ Static Gen75.6  ↑\,{}^{\uparrow}64.676.933.6  ↑\,{}^{\uparrow}51.2  ↑\,{}^{\uparrow}
+ DIAG75.6  ↑\,{}^{\uparrow}64.7  ↑\,{}^{\uparrow}76.933.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 ρt\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.
MethodEff. SamplesAvg.Δ\Delta
Qwen2.5-Math-7B–48.4–
+ SPIN Chen et al. 2024d8k48.9+0.5
+ NuminaMath Li et al. 2024a8k49.0+0.6
+ Static IDPO Zhang et al. 20258k49.7+1.3
+ LLM2LLM Lee et al. 20248k50.1+1.7
+ Ours8k50.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

Figure 5: Prompt for Student Model and Teacher Model Answer Generation.
Figure 6: Prompt for Teacher Model Question Generation.
Figure 7: Case Studies. Top pair (Improper Integral): The student applies FTC across a singularity at x=0x=0 , arriving at −2-2 ; the follow-up shifts the singularity to x=2x=2 , where the trap value −3/2-3/2 is negative for a strictly positive integrand. Bottom pair (Bayesian Inference): The student normalizes likelihoods without prior weighting ( 5/85/8 instead of 5/75/7 ); the follow-up pushes priors to 1%1\% vs. 99%99\% , where the same shortcut gives 62.5%62.5\% but the correct posterior is ≈1.66%\approx 1.66\% .
More in this topic

LLM Alignment

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

See all 613 papers

Explore similar work

CardsList
Jun 15, 2026

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

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.

Haolong Qian, Xianliang Yang, Yinuo ma +6

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 DistillationLanguage Model DistillationRejection SamplingLLM Mathematical Reasoning
AskPDF
Aug 27, 2026

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

Mathematical reasoning has seen rapid progress in large language models (LLMs), yet existing methods optimize predominantly for final-answer correctness, raising the question whether models truly internalize mathematical concepts or merely memorize solution patterns. In human mathematics education, example-based reasoning such as constructing counterexamples to test theorem boundaries reflects deep conceptual understanding, but remains underdeveloped in current LLMs. Enhancing this capability through preference optimization presents two key challenges: (1) the model's limited example-based reasoning ability makes constructing effective preference pairs inherently difficult; and (2) capability acquisition is progressive, as the model must first learn to adopt this strategy before learning to apply it correctly. Therefore we propose INSPIRE, an Internalize-Then-Improve approach combining Reference-Guided Student Internalization (RGSI), which produces high-quality preference candidates under the policy model's own distribution, with a stage-wise rubric preference training strategy that decomposes learning into method-oriented and correctness-oriented stages. Experiments across multiple model scales and families demonstrate consistent improvements, even surpassing larger open-source models, while evaluations on out-of-distribution benchmarks confirm no degradation in general mathematical reasoning ability.

Shuai Wang, Jiayi Kuang, Yinghui Li +4

Sun Yat-sen University · Tsinghua University · Beihang University +1

LLM Mathematical ReasoningMathematical ReasoningPreference Optimization
AskPDF
Apr 22, 2026

HiPO: Hierarchical Preference Optimization for Adaptive Reasoning in LLMs

Direct Preference Optimization (DPO) is an effective framework for aligning large language models with human preferences, but it struggles with complex reasoning tasks. DPO optimizes for the likelihood of generating preferred over dispreferred responses in their entirety and lacks the granularity to provide feedback on subsections of many-step solutions typical of reasoning tasks. Existing methods excel at either stable preference learning (e.g., DPO variants like KTO and RSO) or structured reasoning (e.g., ReMA's multi-agent RL framework, Tree of Thoughts), but fail to merge these complementary strengths. We propose HiPO (Hierarchical Preference Optimization), an extension of DPO that separates responses into reasoning segments (query clarification and context, reasoning steps, and answer) and computes loss as a weighted sum of the DPO loss for each segment. Our approach enables segment-specific training while maintaining DPO's computational efficiency and training stability. We demonstrate that for multiple 7B LLMs fine-tuned using HiPO and DPO on the Math Stack Exchange preference dataset, the models trained with HiPO outperform the others on a variety of common math benchmarks and achieve greater organization, logical flow, and consistency as measured by GPT-4.1.

Darsh Kachroo, Arjun Prasaath Anbazhagan, Adriana Caraeni +2

Vellore Institute of Technology, Chennai · Northwestern University · University of Massachusetts Amherst +2

LLM AlignmentDirect Preference OptimizationMathematical ReasoningPreference OptimizationLLM Reasoning
AskPDF