Stochastic Pruning of the DPO Computation Graph: 1.4x Faster Training and Lower Peak Memory

Yubing Wu; Wei Chen; Junmei Yang; Delu Zeng; Tiewen Pan

cs.LGOct 21, 2026

Stochastic Pruning of the DPO Computation Graph: 1.4x Faster Training and Lower Peak Memory

Authors: Yubing Wu, Wei Chen, Junmei Yang, Delu Zeng, Tiewen Pan

Organizations: South China University of Technology · Shenzhen Technology University

Abstract

Direct preference optimization (DPO) incurs training costs in output scoring, sequence execution, and backpropagation, so reducing supervision alone does not necessarily save computation. We introduce a unified framework for stochastic pruning of the DPO computation graph that independently samples response positions for scoring, response prefixes for execution, and preference branches for backpropagation. By making these selections before the corresponding operations, the framework omits computation at each stage of a training step. Our analysis characterizes the objective changes induced by position and sequence pruning and uses inverse-probability weighting to obtain a conditionally unbiased single-branch gradient estimator for the sampled objective. Controlled measurements show that position and sequence pruning primarily reduce peak memory, with most output-memory savings attributable to an implementation that retains full response supervision. Branch pruning yields a $1.426\times$ step-time speedup over a concatenated DPO baseline and reduces training time by about 30% relative to stock DPO across two model families, with slightly lower mean accuracy and benchmark-dependent changes in open-ended scores. These findings relate the resource benefits of stochastic graph pruning to the computation omitted and the resulting changes in the training objective.

Figures & tables

Figure 1: Stochastic interventions on the DPO computation graph. The two response branches share policy parameters. Rewards average policy–reference log-ratios over scored response positions. Fixed-reference computations use matching prefixes and positions and are omitted. Sampling outcomes and counts are schematic.

Figure 2: Resource ratios at the default settings. Dashed lines mark each baseline. Panel (c) counts executed input tokens, including prompt and padding.

Configuration	Retention	Step (s)	HBM (MB)
Full-concat	1.0	17.081	23,627
Selected head	1.0	16.827	14,722
Selected head	0.6	16.791	14,080

Table 1: Qwen output-head attribution in a separate matched profile. Retention

1.0

preserves all response supervision. Shading highlights full-supervision selected-head execution. Details: Appendix D.7 .

Sequence setting	Tokens	HBM (MB)	$\Delta$ HBM	Step (s)
Full	1.000 $\times$	18,571	—	18.826
$\mathcal{U}(0.25,1)$	0.854 $\times$	16,718	$-10.0\%$	18.480
$\mathcal{U}(0.5,1)$	0.893 $\times$	16,827	$-9.4\%$	18.430
$\mathcal{U}(0.75,1)$	0.943 $\times$	17,072	$-8.1\%$	18.482
fixed $0.75$	0.894 $\times$	16,561	$-10.8\%$	18.471

Table 2: Sequence profile against Full. Tokens counts executed input, including prompt and padding, relative to Full. Shading marks the default

\mathcal{U}(0.5,1)

Axis	Setting	Resource	Mean accuracy (%)	Mean generation score
Position	$\alpha{=}0.30$	11,211 MB	57.90	47.73
Position	$\alpha{=}0.60$	11,320 MB	58.06	48.38
Position	$\alpha{=}0.90$	11,580 MB	58.32	48.63
Sequence	$\mathcal{U}(0.25,1)$	16,718 MB	58.09	47.91
Sequence	$\mathcal{U}(0.50,1)$	16,827 MB	58.04	47.33
Sequence	$\mathcal{U}(0.75,1)$	17,072 MB	58.04	49.53

Table 3: Ablation on Qwen3-4B/HelpSteer3-Preference. Shading marks defaults. Means follow Sec. 4.1 . Full task scores: Tab. 10 .

Method	Wall (min)	Mean accuracy (%)	Mean generation score
Full	187.8	58.33	49.33
Position	185.9	58.06	48.38
Sequence	186.4	58.04	47.33
Branch	128.4	58.08	50.25
Position+Sequence	185.0	57.97	48.85
Position+Branch	127.8	57.94	48.17

Table 4: Qwen3-4B factorial, one training and one selection seed. Bold/underline mark best/second-best values, including ties. Full scores: Tab. 9 .

Figure 3: Training time versus seven-task mean accuracy for the eight factorial runs. The gray staircase is the observed frontier of single-seed point estimates.

Config.	Wall (min)	Speedup	Mean acc. (%)	Mean gen. score
Qwen3-4B, seed 42
Full	187.8	$1.000\times$	58.33	49.33
Branch	128.4	$\bm{1.463\times}$	58.08	50.25
Qwen3-4B, seed 43
Full	187.0	$1.000\times$	58.26	47.64
Branch	131.0	$\bm{1.428\times}$	57.81	48.49

Table 5: Full–Branch validation. Speedup uses each pair’s Full training time. Bold marks the better value within each pair. Score means follow Sec. 4.1 ; task scores are in Tab. 13 .

Qwen3-4B
Configuration	Step (s)	HBM (MB)
Full	18.204	18,571
Branch	12.161	17,319
Mistral-7B
Full	22.438	21,185
Branch	14.888	19,838

Table 6: Paired Full–Branch step-time and peak-memory profiles. Bold marks lower values within each model.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Qwen3-4B
Benchmark	All/All	Random/Random	High/High	Low/Low
ARC	$62.80{\scriptstyle\pm 1.41}$	$62.97{\scriptstyle\pm 1.41}$	$62.46{\scriptstyle\pm 1.42}$	$62.97{\scriptstyle\pm 1.41}$
BBH	$56.16{\scriptstyle\pm 3.01}$	$56.16{\scriptstyle\pm 3.01}$	$56.08{\scriptstyle\pm 3.01}$	$56.33{\scriptstyle\pm 3.01}$
MMLU-Pro	$42.46{\scriptstyle\pm 0.45}$	$42.25{\scriptstyle\pm 0.45}$	$42.24{\scriptstyle\pm 0.45}$	$42.50{\scriptstyle\pm 0.45}$
MuSR	$42.46{\scriptstyle\pm 1.73}$	$43.39{\scriptstyle\pm 1.74}$	$42.72{\scriptstyle\pm 1.74}$	$43.25{\scriptstyle\pm 1.74}$
TruthfulQA	$37.94{\scriptstyle\pm 1.70}$	$38.07{\scriptstyle\pm 1.70}$	$36.11{\scriptstyle\pm 1.68}$	$38.56{\scriptstyle\pm 1.70}$

Appendix

Table 7: UltraFeedback Qwen3 task scores (%) with evaluator-reported standard errors. Mean accuracy is the eight-benchmark point mean. Non-All rules use retention 0.8. These errors do not measure training variability.

Method	ARC	BBH	MMLU-Pro	MuSR	IFEval	TruthfulQA	Mean accuracy (%)
All/All	$25.51{\scriptstyle\pm 1.27}$	$28.46{\scriptstyle\pm 2.64}$	$10.59{\scriptstyle\pm 0.28}$	$35.85{\scriptstyle\pm 1.71}$	$11.46{\scriptstyle\pm 1.37}$	$22.40{\scriptstyle\pm 1.46}$	22.38
Random/Random	$28.16{\scriptstyle\pm 1.31}$	$27.47{\scriptstyle\pm 2.55}$	$10.86{\scriptstyle\pm 0.28}$	$36.64{\scriptstyle\pm 1.72}$	$15.53{\scriptstyle\pm 1.56}$	$25.09{\scriptstyle\pm 1.52}$	23.96
Low/Low	$27.73{\scriptstyle\pm 1.31}$	$28.47{\scriptstyle\pm 2.69}$	$10.84{\scriptstyle\pm 0.28}$	$39.95{\scriptstyle\pm 1.73}$	$11.46{\scriptstyle\pm 1.37}$	$27.54{\scriptstyle\pm 1.56}$	24.33
Low/High	$24.91{\scriptstyle\pm 1.26}$	$28.92{\scriptstyle\pm 2.69}$	$11.66{\scriptstyle\pm 0.29}$	$35.58{\scriptstyle\pm 1.70}$	$18.30{\scriptstyle\pm 1.66}$	$22.64{\scriptstyle\pm 1.47}$	23.67
High/Low	$26.11{\scriptstyle\pm 1.28}$	$29.47{\scriptstyle\pm 2.72}$	$11.66{\scriptstyle\pm 0.29}$	$34.79{\scriptstyle\pm 1.69}$	$12.57{\scriptstyle\pm 1.43}$	$23.01{\scriptstyle\pm 1.47}$	22.93
High/High	$25.51{\scriptstyle\pm 1.27}$	$29.99{\scriptstyle\pm 2.76}$	$11.16{\scriptstyle\pm 0.29}$	$34.39{\scriptstyle\pm 1.68}$	$12.20{\scriptstyle\pm 1.41}$	$23.13{\scriptstyle\pm 1.48}$	22.73

Appendix

Table 8: UltraFeedback Pythia-2.8B task scores (%) at retention 0.8. Per-task errors are evaluator-reported standard errors. Mean accuracy is the six-task point mean; GSM8K and MATH-Hard are excluded because all variants score zero on both.

	Efficiency		Capability (accuracy, %)								Open-ended
Method	Wall	Speed	ARC	GSM8K	TQA	BBH	MATH	MMLU	MuSR	Mean accuracy (%)	MT	AH	WB	Mean generation score
Full	187.8	1.000 $\times$	62.54	87.41	53.88	55.64	63.59	42.28	42.99	58.33	6.649	50.12	3.139	49.33
Position	185.9	1.010 $\times$	62.46	86.43	53.51	55.58	63.26	42.18	42.99	58.06	6.829	44.85	3.200	48.38
Sequence	186.4	1.008 $\times$	62.37	86.28	53.40	55.92	63.27	42.20	42.86	58.04	6.432	45.80	3.186	47.33
Branch	128.4	1.463 $\times$	62.03	86.66	53.06	55.73	63.65	42.16	43.25	58.08	6.800	50.11	3.263	50.25
Position+Sequence	185.0	1.015 $\times$	62.20	85.97	53.06	55.74	63.63	42.07	43.12	57.97	6.893	45.54	3.208	48.85

Appendix

Table 9: Full factorial task scores underlying Tab. 4 . Metrics and score units follow Sec. 4.1 . Bold/underline mark the best/second-best displayed column values.

Setting	ARC	GSM8K	TQA	BBH	MATH	MMLU	MuSR	Mean accuracy (%)	MT	AH	WB	Mean generation score
Position $\alpha{=}0.3$	62.29	86.20	53.15	55.61	62.78	42.15	43.12	57.90	6.568	45.56	3.195	47.73
Position $\alpha{=}0.6$	62.46	86.43	53.51	55.58	63.26	42.18	42.99	58.06	6.829	44.85	3.200	48.38
Position $\alpha{=}0.9$	62.54	87.34	53.83	55.64	63.54	42.26	43.12	58.32	6.707	48.00	3.081	48.63
Sequence $\mathcal{U}(0.25,1)$	62.37	86.35	53.30	55.80	63.53	42.15	43.12	58.09	6.552	47.05	3.117	47.91
Sequence $\mathcal{U}(0.5,1)$	62.37	86.28	53.40	55.92	63.27	42.20	42.86	58.04	6.432	45.80	3.186	47.33
Sequence $\mathcal{U}(0.75,1)$	62.46	86.28	53.44	56.03	63.16	42.16	42.72	58.04	6.874	47.92	3.193	49.53

Appendix

Table 10: Per-task single-axis sweeps, one training seed. Shading marks defaults; rows specify settings. Metric units follow Sec. 4.1 . Evaluator standard errors (

0.45

–

3.02

pp) do not test aggregate differences or training variability.

Selection	Retention	ARC	GSM8K	TruthfulQA	BBH	MATH-Hard	MMLU-Pro	MuSR	Mean accuracy (%)
Low/Low	0.1	63.48	85.82	54.57	55.54	62.39	42.26	42.86	58.13
Random/Random	0.1	62.88	86.73	53.11	55.41	63.36	42.00	42.59	58.01
Low/Low	0.2	63.82	86.58	55.19	55.32	60.29	42.33	42.59	58.02
Random/Random	0.2	63.05	86.35	53.62	55.06	63.14	42.02	42.46	57.96
Low/Low	0.3	63.82	87.26	55.49	55.47	58.47	42.35	42.59	57.92
Random/Random	0.3	63.23	86.50	54.46	55.51	63.20	42.04	42.72	58.24

Appendix

Table 11: HelpSteer3-Preference task scores (%). Low and Random are compared within a retention level. Retentions 0.1–0.4 and 0.5–0.9 use different training schedules. Mean accuracy is the seven-task mean. ∗ At

0.9

both means use six tasks, excluding MATH-Hard, which is missing for Random. The last two rows are baselines: the untrained Qwen3-4B starting checkpoint and unmasked DPO, which provides context but does not use a matched training schedule. A dash denotes a missing result.

$\alpha$	Low/Low	Random/Random	Difference
Batch 32
0.1	58.13	58.01	+0.12
0.2	58.02	57.96	+0.06
0.3	57.92	58.24	$-0.32$
0.4	58.00	58.20	$-0.20$
Batch 64

Appendix

Table 12: Qwen3-4B on HelpSteer3-Preference: mean capability scores (%) for confidence-ranked (Low/Low) versus matched uniform (Random/Random) masks at retention

\alpha

. Comparisons are within rows. ∗ At

0.9

both means use six tasks, excluding MATH-Hard, which is missing for Random. Differences use displayed means.

Model / seed	Config.	ARC	GSM8K	TQA	BBH	MATH	MMLU	MuSR	Mean accuracy (%)	MT	AH	WB	Mean generation score
Qwen3-4B / 42	Full	62.54	87.41	53.88	55.64	63.59	42.28	42.99	58.33	6.649	50.12	3.139	49.33
	Branch	62.03	86.66	53.06	55.73	63.65	42.16	43.25	58.08	6.800	50.11	3.263	50.25
Qwen3-4B / 43	Full	62.54	87.04	53.88	55.55	63.45	42.34	42.99	58.26	6.623	44.36	3.233	47.64
	Branch	62.03	86.20	53.05	55.55	62.57	42.14	43.12	57.81	6.750	45.03	3.295	48.49
Mistral-7B / 42	Full	67.06	52.46	64.82	46.37	12.50	30.47	40.34	44.86	5.130	24.60	2.361	33.17
	Branch	65.27	51.18	62.87	46.51	12.96	30.32	42.72	44.55	5.373	24.00	2.405	33.93

Appendix

Table 13: Complete Full–Branch scores in the units of Sec. 4.1 . Qwen seed-42 values repeat the factorial evaluations in Tab. 4 .

Direct Preference Optimization

Direct Preference Optimization trains generative models from preferred and rejected outputs without fitting an explicit reward model.

See all 160 papers

Jun 11, 2026

Direct Preference Optimization for Chatbot Fine-Tuning: An Empirical Study

We present an approach to fine-tuning large language models using Direct Preference Optimization (DPO), a reinforcement learning technique. Our experimental results demonstrate that DPO simplifies the training pipeline, improves computational efficiency, and achieves competitive performance. The evaluation using BLEU, ROUGE, and cosine similarity metrics indicates effective learning and convergence, though further investigation is needed to address observed training instability.

Dezhi Yu, Yvonne Qiu, ShuoJia Fu

LLM Alignment Direct Preference Optimization LLM Fine-Tuning LLM Post-Training Fine-Tuning Language Models

Ask PDF

Mar 18, 2026

Sharpness-Aware Minimization in Logit Space Efficiently Enhances Direct Preference Optimization

Direct Preference Optimization (DPO) has emerged as a popular algorithm for aligning pretrained large language models with human preferences, owing to its simplicity and training stability. However, DPO suffers from the recently identified squeezing effect (also known as likelihood displacement), where the probability of preferred responses decreases unintentionally during training. To understand and mitigate this phenomenon, we develop a theoretical framework that models the coordinate-wise dynamics in logit space. Our analysis reveals that negative-gradient updates cause residuals to expand rapidly along high-curvature directions, which underlies the squeezing effect, whereas Sharpness-Aware Minimization (SAM) can suppress this behavior through its curvature-regularization effect. Building on this insight, we investigate logits-SAM, a computationally efficient variant that perturbs only the output layer with negligible overhead. Extensive experiments on Pythia-2.8B, Mistral-7B, and Gemma-2B-IT across multiple datasets and benchmarks demonstrate that logits-SAM consistently improves the effectiveness of DPO and integrates seamlessly with other DPO variants. Code is available at https://github.com/RitianLuo/logits-sam-dpo.

Haocheng Luo, Zehang Deng, Thanh-Toan Do +3

Monash University, Australia · Swinburne University of Technology, Australia

LLM Alignment Direct Preference Optimization Sharpness-Aware Minimization

Ask PDF

Jun 10, 2026

Boosting Direct Preference Optimization with Penalization

Offline preference optimization has become a practical substitute for reinforcement learning from human feedback, but pairwise objectives such as Direct Preference Optimization (DPO) and its variants use only the chosen and rejected responses stored in a static dataset. This leaves a useful signal unused: the response that the reference model itself would generate for the same prompt. We propose Direct Preference Optimization with Penalization (DPOP), a simple extension of DPO that augments the base preference loss with a gated penalty on reference-greedy responses. DPOP activates this penalty only when the current policy still assigns a lower likelihood to the preferred response than to the rejected response. On AlpacaEval 2.0, DPOP improves length-controlled win rate over DPO, SimPO, and AlphaDPO on both Llama-3-8b-it and Gemma-2-9b-it, achieving relative gains of 5.3% and 4.4% over baselines on the two models, respectively. Ablations further show that a SimNPO-style length-normalized penalty is stronger than NPO and token-level unlikelihood in this setting.

Pengwei Sun

Stanford University.

Pairwise Preference Learning LLM Alignment Direct Preference Optimization Offline Preference Optimization

Ask PDF

Abstract

Direct Preference Optimization

Explore similar work

Direct Preference Optimization for Chatbot Fine-Tuning: An Empirical Study

Sharpness-Aware Minimization in Logit Space Efficiently Enhances Direct Preference Optimization

Boosting Direct Preference Optimization with Penalization