Deterministic Differentiable Structured Pruning for Large Language Models

DDP: Deterministic Differentiable Structured Pruning for LLMs

一句话总结

用确定性 soft sigmoid 映射替代随机 hard-concrete 松弛实现 $\ell_0$ 结构化剪枝，消除训练-测试不一致和表达力受限问题，仅优化 mask 参数即可在 30M token 内完成 LLM 结构化剪枝。

核心问题

结构化剪枝可视为学习每个组件的乘性门控 $m_k$ + $\ell_0$ 稀疏约束。传统 hard-concrete 松弛：

1. 训练-测试不一致：训练时随机采样 mask，部署时必须离散化
2. 表达力受限：mask 被限制在近二值 $[0,1]$ 范围

关键公式

统一 Mask 框架

$y = \sum_{k=1}^{K} m_k \cdot f_k(X), \quad m_k \in \mathbb{R}$

Multi-head attention ($K=H$ heads)：

$f_h^{\text{attn}}(X) = \text{Attn}(XW_Q^{(h)}, XW_K^{(h)}, XW_V^{(h)}) \cdot W_O^{(h)}$

MLP channels ($K=C$ intermediate width)：

$f_j^{\text{mlp}}(X) = (\phi(Xu_j) \odot (Xg_j)) \cdot v_j$

$\ell_0$ 约束优化

$\min_{m \in S} \mathcal{L}_{\text{ce}}(\theta, m), \quad \bar{m} = \frac{1}{K}\|m\|_0 = \rho$

增广拉格朗日惩罚：

$\mathcal{L}_{\text{sparsity}}(\|m\|_0) = \lambda_1(\bar{m} - \rho) + \lambda_2(\bar{m} - \rho)^2$

DDP 核心：确定性替代

前向传播用确定性 ReLU 门控：

$m = \text{ReLU}(z)$

mask 空间从近二值 $[0,1]$ 扩展到 $[0, \infty)$。

正则化用退火的确定性 soft sigmoid：

$v = \sigma\left(\frac{(z - \mu_t) \cdot C_0}{\mu_t}\right), \quad s = \text{Clamp}(v(r - l) + l, 0, 1)$

固定参数 $l=-0.1, r=1.1, C_0 \approx 2.4$。

退火调度（$\mu_t \to 0$ 时逼近精确 $\ell_0$）：

$\mu_t = \mu_0 - (\mu_0 - \mu_T)\sqrt{t/T}$

默认 $\mu_0 = 0.5$，$\mu_T \approx 0$。

二值化正则器

$\mathcal{L}_{\text{bin}}(s) = \lambda_3 \cdot \frac{1}{K}\sum_k s_k(1 - s_k)$

在端点 $\{0, 1\}$ 最小化，推动模糊组件做出明确选择。

完整训练目标

$\min_z \mathcal{L}_{\text{ce}}(\theta, m) + \mathcal{L}_{\text{sparsity}}(s) + \mathcal{L}_{\text{bin}}(s)$

其中 $m = \text{ReLU}(z)$，$s = \phi(z; \mu_t)$。

关键图表

Figure 1: DDP 方法概览

Figure 2: 确定性替代映射

退火 $\mu$ 使 soft sigmoid 逐步逼近阶跃函数，实现从连续到离散的渐进转换。

实验结果

Table 2: Dense LLM 剪枝

LLaMA-7B 20% sparsity：

LLaMA-13B 50% sparsity：

Table 3: MoE 剪枝

DeepSeekMoE-16B 60% sparsity：

Table 7: 推理加速（vLLM）

消融实验

三个组件（确定性 + 扩展mask + 二值化）加性贡献。

Token 预算：zero-shot accuracy ~10M token 饱和；PPL 持续改善到 60M token。

DDP vs Hard-Concrete 对比

对我们工作的启示

1. Mask 机制用于结构决策：DDP 的 mask 优化思路可迁移到模型增长——增长后判断哪些新结构该保留/删除
2. 确定性 > 随机性：训练-测试一致性对部署很重要
3. 轻量优化：仅优化 mask（数千万参数） vs 完整权重训练 → 30M token 够用
4. MoE 适配：对 expert-level 剪枝的扩展方式可参考
5. 退火调度：$\mu_t$ 退火实现连续→离散渐进转换的技巧有通用性

局限性

20% sparsity 结果优秀，但 50%+ 时精度下降较多

训练数据质量敏感（C4/LaMini 比 FineWeb-Edu 差 1-2pp）

未报告训练不稳定性问题

高稀疏度下的 speedup 需要更多硬件验证

相关概念

PTQ

LoRA

Taylor pruning