PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.29538v1cs.LGcs.AI

TL;DR

DPO training spends computation on response scoring, sequence execution, and backpropagation, while prior work often changes token supervision or computation kernels. This paper samples response positions, prefixes, and gradient branches to prune those stages and analyzes the resulting objective and gradient changes. Controlled comparisons report memory savings, faster branch-pruned training, and a small accuracy trade-off with open-ended score changes varying by benchmark.

  • Problem

    Prior work changes token supervision or output-score computation, leaving the question of which DPO-step computations can be omitted and how savings relate to objective changes.

    Source: §1 · p. 1

  • Method

    The framework independently samples response prefixes, scoring positions, and gradient branches before their corresponding operations, with inverse-probability weighting for the branch estimator.

    Source: §1 · p. 1

  • Results

    Controlled comparisons show that full supervision retains most output-memory savings, while branch pruning reduces training time across model families with slightly lower mean accuracy and benchmark-dependent open-ended score changes.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The resource benefits of stochastic DPO pruning relate to the computation omitted and the resulting changes in the training objective.

    Source: §5 · p. 8

    Evaluation covers parameter-efficient text preference training with an online reference on one eight-accelerator setup, and savings may vary with hardware, batch geometry, and reference execution.

    Source: Limitations · p. 8

Abstract

from arXiv · show

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.

1 Introduction

The paper asks which computations in a DPO step can be omitted and how their resource savings relate to changes in the optimization objective. It unifies sampling over response execution, token scoring, and backpropagation, then evaluates their resource and accuracy trade-offs.

1 Introduction

  • The paper asks which parts of a DPO step can be omitted and how their resource savings relate to changes in the optimization objective.Prior work changes token supervision or uses memory-efficient kernels, motivating a distinction between supervision changes and computation omission.Source: §1 · p. 1
  • The framework separately samples executed response prefixes, scored positions, and branches carrying gradients, omitting computation before sequence execution, vocabulary projection, and backpropagation.Inverse-probability weighting makes the branch gradient conditionally unbiased for the sampled objective.Source: §1 · p. 1
  • The analysis characterizes position and prefix objective changes, relates optimality gaps to retention and truncation, and gives a conditionally unbiased branch estimator with a convergence rate.The DPO gradient includes contributions from chosen and rejected responses that share policy parameters.Source: §1 · pp. 1–2
  • Controlled comparisons find that full response supervision retains most output-memory savings, while branch pruning reduces training time across model families with slightly lower mean accuracy and benchmark-dependent open-ended scores.The framework samples response prefixes, scoring positions, and gradient branches before their corresponding operations.Source: §1 · pp. 1–2

2 Related Work

Related work spans methods that change token-level credit, methods that reduce preference-training computation, and analyses of stochastic computation and preference-optimization dynamics. This paper positions its contribution as sampling throughout a DPO step while distinguishing resource savings from objective changes.

Token-level credit in preference optimization

  • Token-credit methods select or weight tokens using confidence, reward, contrastive predictions, transport, gradient attribution, or other signals; this framework instead places supervision selection within a complete DPO step.The paper notes that matched-selection controls motivate uniform membership and full-retention controls distinguish supervision deletion from implementation savings.Source: §2 · p. 2

Efficient preference training

  • Efficient-training work reduces memory or throughput costs through kernels, attention, checkpointing, or offline token selection; this framework studies sampling across DPO computation and its objective effects.Controlled profiles separate implementation savings from reduced supervision and execution.Source: §2 · p. 2

Stochastic computation and optimization dynamics

  • Prior analyses cover stochastic-gradient sampling and preference-optimization dynamics; this paper separates objective changes from gradient-estimation noise under composed sampling.The paper also analyzes shared-parameter response changes in an appendix.Source: §2 · pp. 2–3

3 Stochastic Graph Pruning

The method independently prunes response positions, response prefixes, and backward branches in a DPO step. Its analysis distinguishes the objective changes from position and sequence pruning from the conditionally unbiased gradient estimation of branch pruning.

3.1 Formulation

  • DPO compares chosen and rejected mean policy–reference log-ratios with a logistic loss; three independently configurable interventions prune response execution, scoring positions, or backward branches.The interventions are sampled before their respective computation stages.Source: §3.1 · p. 3

3.2 Stochastic Position Pruning

  • Uniform position sampling preserves mean reward but changes the nonlinear loss; its variance bound decays with response-length lower bound T0 at fixed retention and vanishes at α = 1.Applying the language-model head only at selected positions skips other vocabulary projections; the loss bound requires no curvature assumption, while the sharper stationary-point displacement bound does.Source: §3.2 · pp. 3–4

3.3 Stochastic Sequence Pruning

  • Sequence pruning truncates responses before model execution, implicitly weights earlier positions, and can shift the mean margin; its optimality bound has no response-length decay and is vacuous at tested ρmin = 1/2.The expected prefix gradient is unbiased for the prefix objective, which can differ from the full-response gradient.Source: §3.3 · p. 4

3.4 Stochastic Branch Pruning

  • Branch pruning backpropagates through one selected response and applies inverse-probability weighting, yielding a gradient conditionally unbiased for the sampled positions and prefixes.It adds gradient noise without objective bias; under the stated strong-convexity and smoothness assumptions, projected SGD has O(1/N) mean-squared error.Source: §3.4 · p. 5

3.5 Composition

  • The composed step truncates responses, selects scoring positions, then selects a backward branch and corrects for its probability; each mode is optional.The composition bound allocates total tolerance across position retention, sequence truncation, and optimization steps; branch pruning adds no objective bias.Source: §3.5 · pp. 5–6

4 Experiments

Across the experiments, branch pruning reduces training time across model families, while position and sequence pruning primarily reduce peak memory. The reported savings coexist with slightly lower mean accuracy and benchmark-dependent open-ended score changes.

4.2 Resource Effects and Attribution

  • With all response supervision retained, selected-head execution saves 8,905 MB and sampling adds 642 MB, accounting for 93.3% and 6.7% of savings, respectively.The separate matched profile compares Qwen output-head attribution.Source: Table 1 · p. 6, §4.2 · p. 6
  • Position+Branch reaches 1.443× speedup and 48.6% lower HBM, but its 11,471 MB saving is below the 15,855 MB sum of separate savings; the cause remains unresolved.Branch alone reduces step time by 29.9% and peak HBM by 26.7% against Full-concat.Source: §4.2 · pp. 6–7, Figure 2 · p. 7
  • Sequence pruning executes 0.85–0.94 of Full’s input tokens and saves 8–11% peak HBM, while steps take 18.430–18.482 s, at most 2.1% below Full.The token counts include prompt and padding; timing differences between cutoffs are smaller than repeated-measurement spread.Source: Table 2 · p. 7, §4.2 · p. 7

4.3 Parameter Sensitivity

  • Increasing Position retention α from 0.3 to 0.9 raises mean accuracy from 57.90 to 58.32, while peak HBM changes little at 11,211–11,580 MB and AH is non-monotone.The memory pattern agrees with the implementation-dominated attribution above.Source: Table 3 · p. 7, §4.3 · p. 7
  • Branch probabilities yield similar training times of 127.7–128.4 min, alongside varying open-ended point estimates; the analysis retains qw = 0.5.HBM was not measured across qw.Source: §4.3 · p. 7

4.4 The Eight-Configuration Factorial

  • Branch raises MT and WB point estimates with AH close to Full, while configurations containing Position or Sequence fall below Full on AH; generation means are descriptive single-seed summaries.Figure 3 plots training time versus seven-task mean accuracy for the eight factorial runs, with a staircase marking the observed single-seed frontier.Source: §4.4 · p. 7, Figure 3 · p. 8
  • All four Branch configurations finish in 126–128 min versus 185–188 min without Branch, a 31–32% saving; all 12 single-axis contrasts lower mean accuracy.The factorial uses one training seed and one selection seed.Source: §4.4 · p. 7, Table 4 · p. 8

4.5 Branch Pruning Across Seeds and Models

  • 29.9–32.7% lower wall-clock time across two Qwen seeds and Mistral, with mean accuracy lower by 0.25–0.45 pp for Full–Branch trainings.The Mistral profile reports 1.507× speedup and 1,347 MB (6.4%) saved against Full.Source: Table 5 · p. 8, §4.5 · p. 8

5 Conclusion

The conclusion separates output scoring, sequence execution, and backward computation to examine resource savings and optimization changes. It links those savings to trade-offs in accuracy and open-ended results.

5 Conclusion

  • The framework separates output scoring, sequence execution, and backward computation to locate resource savings and optimization changes.Source: §5 · p. 8
  • Branch-pruning time benefits reproduce across Qwen seeds and Mistral, alongside lower mean accuracy and mixed open-ended results.Most output-memory savings come from implementation rather than supervision deletion; configuration choice should reflect resource constraints and acceptable objective changes.Source: §5 · p. 8
  • Configuration choice should follow the resource constraint and the acceptable objective change.Source: §5 · p. 8

Limitations

The evaluation is limited to parameter-efficient text preference training with an online reference on one eight-accelerator setup. The appendix defines the notation and establishes results for position, prefix, and branch sampling.

Limitations

  • Evaluation covers parameter-efficient text preference training with an online reference on one eight-accelerator setup.Savings may vary with hardware, batch geometry, and reference execution.Source: Limitations · p. 8

A Notation

  • DPO uses chosen and rejected responses, token log-ratios to form rewards and a margin, and a logistic per-pair loss.Source: A Notation · p. 11
  • Sequence, position, and branch pruning respectively retain response prefixes, select scoring positions, and sample a backward branch.The notation also names the sampling laws, active branches, probability correction, and scoring operator.Source: A Notation · p. 11

B.1 Proof of Proposition 1

  • Uniform position sampling preserves the expected side reward, but mask-dependent loss gradients need not be unbiased for the full-loss gradient.The expected loss difference is quantified separately.Source: §B.1 · p. 12

B.2 Proof of Lemma 1

  • A nonnegative normalized non-increasing cutoff-weight vector corresponds to a cutoff distribution; uniform weights recover the fixed full-response cutoff.Source: §B.2 · p. 12

B.3 Proof of Lemma 2

  • Prefix sampling yields an unbiased gradient of its sampled objective, which need not match the full-response gradient.The temporal shift and Jensen gap can cancel at a fixed policy, so equal loss values there do not establish equal gradients.Source: §B.3 · p. 13

B.4 Proof of Proposition 2

  • Inverse-probability weighting makes the sampled single-branch gradient conditionally unbiased for the objective defined by the sampled positions and prefixes.With pruning disabled, this recovers the full-response objective.Source: §B.4 · p. 13

B.5 Proof of Corollary 1 and the retention rule

The appendix derives bounds for position-pruning objectives and gradients, then relates gradient error to stationary-point displacement under curvature conditions. It also gives a sufficient curvature condition for stability.

B.5 Proof of Corollary 1 and the retention rule

  • Uniform position sampling gives F_A ≥ F, and the retention condition can be chosen to satisfy a specified loss tolerance.The bound uses the position-retention rate and variance bound.Source: §B.5 · p. 13

B.6 Proof of Eq. (31) and the stationary displacement

  • Under the stated stationary-point and curvature assumptions, a gradient difference of O(r_A) yields displacement of O(r_A).The gradient bound uses the position-pruning variance and gradient bounds.Source: §B.6 · p. 14

B.7 Curvature constants from the margin geometry

  • If margins are bounded by C, gradient covariance is bounded below by νI, and the margin Hessian by H, then β^2m_Cν > βH suffices for stability.For a linear-margin model, H = 0, so positive ν suffices.Source: §B.7 · p. 14

B.8 Proof of Eq. (13)

The appendix bounds the effects of random prefix truncation and derives variance properties for the single-branch estimator. It also distinguishes independent estimator averaging from redrawing only the branch.

B.8 Proof of Eq. (13)

  • Prefix truncation bounds the expected loss and gradient changes in terms of the expected removed response fractions.The expected gradient bias is at most G(β^2U + 2β)(1 − ρ_min).Source: §B.8 · p. 14
  • At ρ_min = 1/2, the optimality-gap bound reaches the full loss range and is vacuous; the loss bound is half that range.The bounds become restrictive only for cutoffs closer to one.Source: §B.8 · p. 14

B.9 Proof of Proposition 3 and Eq. (41)

  • At the symmetric branch probability, the estimator variance is ∥c_w − c_l∥^2; averaging independent estimators divides variance by m, while branch-only redraw reduces only the branch component.Source: §B.9 · p. 15

B.10 Proof of the composed bias and error budget

The proof combines a cutoff-averaged position-pruning bias bound with a prefix-bias bound to obtain a composed bound. It then relates the composed objective losses using Jensen’s inequality and the prefix-retention range.

B.10 Proof of the composed bias and error budget

  • The averaged position-pruning bound uses the fact that a prefix’s within-prefix sum of squares cannot exceed the whole response’s sum of squares.For sampled lengths above one, the cutoff law bounds the expected length factor by 1/(M_0 − 1); length-one prefixes have zero variance.Source: §B.10 · p. 15
  • The composed bias adds a cutoff-averaged position-pruning term, C_A r_A|P, and a prefix-bias term, C_P(1 − ρ_min).The position term is averaged over the cutoff law because a fixed prefix can have greater token-score variance than the full response.Source: §B.10 · p. 15
  • The resulting parameter-distance bound uses μ_F, which lower-bounds F’s curvature on the relevant segment, and constants evaluated for the composed objective.The supplied passage refers to the resulting bound as Eq. (45).Source: §B.10 · pp. 15–16
  • For the composed optimal loss, Jensen’s inequality gives F_AP ≥ F_P, while F ≤ F_P + 2βU(1 − ρ_min).The proof chains these inequalities to bound the composed loss.Source: §B.10 · p. 16

C Shared-Parameter Analysis

The shared-parameter analysis examines how a pair-loss update can affect individual token probabilities, while noting that its single-pair calculation does not model the full optimizer. The accompanying experimental material specifies UltraFeedback and HelpSteer3 settings and the geometry of retention-sweep comparisons.

C Shared-Parameter Analysis

  • The appendix expands the shared-parameter coupling introduced earlier; within fixed mask membership, token positions receive coefficients based on selected chosen or rejected status.The coefficients are 1/k_w for selected chosen tokens, −1/k_l for selected rejected tokens, and zero otherwise.Source: C Shared-Parameter Analysis · p. 16, §C.1 · p. 16

C.1 Shared-Parameter Updates and Signed Token Movement

  • An individual chosen-token probability can decrease after a step even when the pair margin increases, if its token gradient has a negative inner product with the pair gradient.Token confidence alone does not determine these inner products, and the calculation is for a single pair rather than the actual minibatch and AdamW update.Source: §C.1 · p. 16

D Experimental Settings and Results

  • The appendix documents experimental settings and additional results across UltraFeedback motivation experiments, HelpSteer3, sweeps, task tables, and supplemental validation.These materials are organized across Sections D.1–D.7.Source: D Experimental Settings and Results · p. 16

D.1 The UltraFeedback Experiments

  • UltraFeedback experiments use retention 0.8 and seed 42 per rule; Low−Random mean-accuracy differences are +0.37, +0.51, and +0.26 points for Pythia-2.8B, Qwen3-4B, and Qwen3-8B, respectively.Pythia’s mean excludes GSM8K and MATH-Hard because all variants score zero on both; the collections’ scores are not comparable across UltraFeedback and HelpSteer3.Source: §D.1 · pp. 16–17

D.2 The HelpSteer3 Corpus and Recipe

  • The HelpSteer3 preference conversion retains 36,264 of 38,459 training pairs and 1,917 of 2,017 validation pairs, removing ties and empty responses.The main recipe uses Qwen3-4B with sequence and prompt limits of 2,048 and 1,024, respectively.Source: §D.2 · p. 17

D.3 Retention-Sweep Geometry

  • Retention-sweep comparisons are interpreted within each retention level because the 0.1–0.4 and 0.5–0.9 blocks use different training schedules.Within each block, Low–Random pairs share geometry and data order; the blocks do not form one controlled dose-response curve.Source: §D.3 · p. 17

D.4 The Factorial and Parameter Sweeps

The factorial and supplemental experiments report descriptive quality contrasts, branch-pruning wall-clock differences, and validation across two model families. The appendix also documents the configurations, task tables, and measurement controls underlying those results.

D.4 The Factorial and Parameter Sweeps

  • The factorial uses position retention α = 0.6, prefix cutoffs sampled from U(0.5, 1.0), and one backward branch per pair with q_w = q_l = 0.5.All eight configurations were trained end-to-end; composition applies truncation, then position pruning within the prefix, then branch pruning.Source: §D.4 · p. 17
  • The single-axis sweeps match nominal position and branch sparsity, while sequence-sweep executed-token fractions exceed response-token retention and are a distinct quantity.Peak HBM was not measured across branch-retention settings, so those comparisons use training wall-clock only.Source: §D.4 · p. 17
  • Branch pruning changes mean accuracy by −0.26, −0.12, −0.19, and −0.28 pp over Full, Position, Sequence, and Position+Sequence, while its wall-clock effect is consistently negative across four pairings.The corresponding training wall-clock differences are −59.4, −58.1, −60.1, and −58.0 min; factorial contrasts are descriptive single-seed comparisons.Source: §D.4 · pp. 17–18
  • Adding position, sequence, or branch pruning changes mean generation scores with mixed signs across the reported pairings.For branch pruning, the four reported changes are +0.92, −0.21, +0.35, and −0.50; the metric uses the scaling in Section 4.1.Source: §D.4 · pp. 17–18

D.5 Full Task Tables

  • The task tables report per-task scores, but evaluator standard errors do not measure training-seed variability or test aggregate differences.Table 10 covers single-axis sweeps; Table 11 gives HelpSteer3 scores with comparisons within retention levels, and Table 9 gives factorial task scores.Source: §D.5 · p. 18, Table 10 · p. 19, Table 9 · p. 19

D.7 Supplemental Attribution and Full–Branch Validation

  • A controlled output-profile comparison reports 8,905 MB of memory difference for full-supervision selected-head execution and a further 642 MB for random deletion.The profile compares Full-concat with selected-head execution at retention 1.0 and 0.6; full-retention selected-head rewards match Full-concat exactly.Source: §D.7 · p. 19
  • In matched supplemental runs, Branch-minus-Full mean-accuracy differences are −0.25 and −0.45 pp for two Qwen seeds and −0.31 pp for Mistral.Both model families fit the same 2 × 4 × 8 geometry without a batch-size change; the open-ended score summaries do not characterize training-seed variability.Source: §D.7 · pp. 19–20

E Additional Token-Selection Analysis

The token-selection analysis compares confidence-ranked position selection with uniform random selection at matched retained ratios. Its results show modest, benchmark-dependent differences, motivating uniform position pruning before vocabulary projection.

E Additional Token-Selection Analysis

E.1 Supervision Membership Is a Limited Lever

  • At equal retained ratios on UltraFeedback, Low/Low exceeds Random/Random mean accuracy by 0.37, 0.51, and 0.26 points for Pythia-2.8B, Qwen3-4B, and Qwen3-8B, respectively.On HelpSteer3-Preference, the difference changes sign across retention levels and stays within 0.32 points.Source: §E.1 · pp. 20–21
  • Task-level differences can diverge from aggregate comparisons: at low retention, Low trails Random on MATH-Hard while leading TruthfulQA.The comparisons motivate uniform position pruning before vocabulary projection, while the factorial separately assesses quality costs relative to Full.Source: §E.1 · p. 21

Related papers

Loading 2610.29538v1…