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

📄 arXiv: 2610.29538v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-10-21

💡 一句话要点

提出随机剪枝框架以加速DPO训练并降低内存峰值

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 直接偏好优化 随机剪枝 计算图优化 训练加速 内存管理

📋 核心要点

  1. 现有的直接偏好优化方法在训练过程中面临高计算成本，尤其是在输出评分和反向传播阶段。
  2. 本文提出了一种随机剪枝的统一框架，通过独立采样不同的响应位置和偏好分支来优化DPO计算图。
  3. 实验结果显示，分支剪枝实现了1.426倍的速度提升，并在两个模型系列中训练时间减少约30%。

📝 摘要（中文）

直接偏好优化（DPO）在输出评分、序列执行和反向传播中消耗大量训练成本，仅减少监督并不能有效节省计算资源。本文提出了一种统一的随机剪枝框架，独立地为评分、执行和反向传播采样响应位置、响应前缀和偏好分支。通过在相应操作之前进行选择，该框架在每个训练步骤中省略了计算。分析表明，位置和序列剪枝主要降低了峰值内存，输出内存节省主要归因于保留完整响应监督的实现。分支剪枝相较于连接的DPO基线实现了1.426倍的步骤时间加速，并在两个模型系列中相对标准DPO减少了约30%的训练时间，尽管平均准确率略低且开放式评分在基准上有所变化。这些发现将随机图剪枝的资源效益与省略的计算和训练目标的变化联系起来。

🔬 方法详解

问题定义：本文旨在解决直接偏好优化（DPO）在训练过程中高昂的计算成本，尤其是在输出评分、序列执行和反向传播阶段的资源消耗问题。现有方法在减少监督的情况下并未有效降低计算负担。

核心思路：提出的随机剪枝框架通过在执行前独立采样响应位置、响应前缀和偏好分支，省略了每个训练步骤中的计算，从而降低了内存使用和训练时间。

技术框架：该框架包括三个主要模块：1) 位置采样模块，用于选择评分的响应位置；2) 前缀采样模块，负责选择执行的响应前缀；3) 分支采样模块，处理反向传播中的偏好分支。

关键创新：最重要的技术创新在于通过随机剪枝实现了计算的省略，进而优化了训练效率和内存使用。这一方法与传统的DPO方法相比，显著提高了训练速度和资源利用率。

关键设计：在实现中，采用了逆概率加权的方法来获得条件无偏的单分支梯度估计器，并保留了完整的响应监督以最大化输出内存的节省。

📊 实验亮点

实验结果表明，分支剪枝相较于连接的DPO基线实现了1.426倍的步骤时间加速，并在两个模型系列中训练时间减少约30%。尽管平均准确率略有下降，但在开放式评分方面的变化依赖于具体基准，显示出该方法的有效性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、推荐系统和其他需要高效训练的机器学习任务。通过降低训练时间和内存消耗，研究成果能够推动更大规模模型的训练和应用，提升实际系统的性能和响应速度。

📄 摘要（原文）

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.