Stochastic Pruning of the DPO Computation Graph: Faster Training and Lower Peak Memory
ORIGINAL / Stochastic Pruning of the DPO Computation Graph: 1.4x Faster Training and Lower Peak Memory
The study randomly samples and prunes response positions, response prefixes and preference branches in the DPO computation graph before the corresponding operations, skipping computation at each training stage. Its novelty lies in linking resource savings to changes in the training objective and using inverse-probability weighting for a conditionally unbiased single-branch gradient estimator, which is useful for practitioners concerned with DPO training efficiency.
01 ABSTRACT
The authors propose a unified framework that independently samples response positions, response prefixes and preference branches of the DPO computation graph and selects them before the corresponding operations to omit computation. The analysis characterizes objective changes from position and sequence pruning and uses inverse-probability weighting to obtain a conditionally unbiased single-branch gradient estimator. Controlled measurements show position and sequence pruning mainly reduce peak memory, while branch pruning achieves a 1.426x step-time speedup over a concatenated DPO baseline.
02 KEY FINDINGS
- Position and sequence pruning mainly reduce peak memory, with most output-memory savings coming from an implementation that retains full response supervision.
- Branch pruning yields a 1.426x step-time speedup over a concatenated DPO baseline and about 30% less training time than stock DPO.
- The framework selects before sampling to skip computation in scoring, execution and backpropagation, using inverse-probability weighting to keep a conditionally unbiased gradient estimator.
AI GENERATED SUMMARY / DISCOVERED BY ARXIV CS.LG