[论文解读] PLATON: Pruning Large Transformer Models with Upper Confidence Bound of Weight Importance

Qingru Zhang, Simiao Zuo+5 more|arXiv (Cornell University)|Jun 25, 2022

Topic Modeling被引用 11

一句话总结

PLATON 提出了一种针对大型 Transformer 模型的结构化剪枝方法，通过在重要性估计中引入不确定性，利用敏感度的上置信界（UCB）提升训练稳定性和性能。通过将指数平滑的敏感度与基于时间变化的不确定性量化相结合，PLATON 降低了重要性分数的变异性，实现了更可靠的剪枝——尤其在高稀疏度场景下表现优异，在 90% 稀疏度下于 SQuAD 上实现了超过 5% 的 F1 分数提升。

ABSTRACT

Large Transformer-based models have exhibited superior performance in various natural language processing and computer vision tasks. However, these models contain enormous amounts of parameters, which restrict their deployment to real-world applications. To reduce the model size, researchers prune these models based on the weights' importance scores. However, such scores are usually estimated on mini-batches during training, which incurs large variability/uncertainty due to mini-batch sampling and complicated training dynamics. As a result, some crucial weights could be pruned by commonly used pruning methods because of such uncertainty, which makes training unstable and hurts generalization. To resolve this issue, we propose PLATON, which captures the uncertainty of importance scores by upper confidence bound (UCB) of importance estimation. In particular, for the weights with low importance scores but high uncertainty, PLATON tends to retain them and explores their capacity. We conduct extensive experiments with several Transformer-based models on natural language understanding, question answering and image classification to validate the effectiveness of PLATON. Results demonstrate that PLATON manifests notable improvement under different sparsity levels. Our code is publicly available at https://github.com/QingruZhang/PLATON.

研究动机与目标

• 解决因在小批量上估计的重要性分数方差过高而导致的迭代剪枝训练不稳定性问题。
• 降低在训练过程中因敏感度估计噪声而剪除关键权重的风险。
• 提升大型 Transformer 模型在高稀疏度剪枝场景下的泛化能力和稳定性。
• 开发一种方法，平衡对不确定性高的低重要性权重的探索与对高重要性权重的利用。
• 通过量化组级别敏感度分数的不确定性，实现更可靠的结构化剪枝。

提出的方法

• PLATON 将敏感度计算为权重梯度与权重值的点积的绝对值：$ I = |\theta^\top \nabla_\theta \mathcal{L}(\theta) | $。
• 应用指数移动平均（EMA）以在时间上平滑敏感度：$ \overline{I}^{(t+1)} = \beta_1 \overline{I}^{(t)} + (1 - \beta_1) I^{(t)} $。
• 不确定性通过与平滑敏感度的绝对偏差来量化：$ U^{(t)} = |I^{(t)} - \overline{I}^{(t)}| $。
• 对不确定性应用第二次 EMA：$ \overline{U}^{(t+1)} = \beta_2 \overline{U}^{(t)} + (1 - \beta_2) U^{(t)} $。
• 最终的重要性分数为平滑敏感度与不确定性的乘积：$ S^{(t)} = \overline{I}^{(t)} \cdot \overline{U}^{(t)} $。
• 通过选择重要性分数最低的权重或组（采用 top-r% 策略）执行剪枝。

实验结果

研究问题

• RQ1 能否有效建模重要性分数估计中的不确定性，以提升大型 Transformer 模型剪枝的稳定性？
• RQ2 通过上置信界（UCB）引入不确定性，是否能降低迭代剪枝过程中敏感度分数的变异性？
• RQ3 在高稀疏度剪枝设置中，基于 UCB 的重要性分数是否优于标准敏感度或基于幅值的度量？
• RQ4 所提出的方法在结构化剪枝中表现如何，即整个组（如列）被剪枝时？
• RQ5 与现有迭代剪枝基线相比，该方法在多大程度上降低了训练不稳定性并提升了泛化能力？

主要发现

• 在将 DeBERTaV3-base 剪枝至 90% 稀疏度时，PLATON 在 SQuAD v1.1 基准上实现了超过 5.0% 的 F1 分数提升，显著优于先前方法。
• 在 SQuAD v1.1 上微调后的 BERT-base 模型中，PLATON 在保留 50% 参数时仍保持 78.5% 的 EM 和 87.2% 的 F1，优于基线方法。
• 与基于标准敏感度的剪枝相比，PLATON 中重要性分数的变异性显著降低，如图 1（底部）所示，从而实现了更稳定的训练。
• PLATON 的训练过程比现有迭代剪枝方法（如 movement pruning）更稳定，且对超参数选择不那么敏感。
• 在结构化剪枝中，PLATON 在 SQuAD v1.1 上于保留 40% 参数时实现了 78.0% 的 EM 和 86.9% 的 F1，优于非结构化基线方法。
• 该方法通过保留高不确定性但低敏感度的权重，有效防止了关键权重的过早剪枝，从而实现了对模型容量的更好探索。