P$^2$ Law: Scaling Law for Post-Training After Model Pruning

本文提出P² Law作为剪枝后大型语言模型后训练的首个缩放定律，通过结合模型规模、后训练数据量、剪枝率和初始损失预测后训练损失，并在多种剪枝方法和模型上验证其有效性和部分泛化能力。

Large Language Model, Pre-training, Fine-tuning, Efficiency, Scaling Laws

Xiaodong Chen, Yuxuan Hu, Xiaokang Zhang, Yanling Wang, Cuiping Li, Hong Chen, Jing Zhang

Renmin University of China, Zhipu AI, Engineering Research Center of Database and Business Intelligence, MOE, China, Key Laboratory of Data Engineering and Knowledge Engineering, MOE, China

Generated by grok-3

Background Problem

大型语言模型（LLM）基于Transformer架构，在多个领域和任务中表现出色，但其参数规模的快速增长导致硬件需求激增，限制了实际部署。为解决这一问题，模型剪枝技术被广泛用于压缩模型规模以降低硬件需求。然而，剪枝通常会导致性能下降，因此后训练（Post-Training）成为恢复性能的重要步骤。现有研究表明，使用大量数据进行后训练（如持续预训练）能有效恢复性能，但硬件成本高昂，而数据量达到一定规模后性能提升有限。因此，如何平衡后训练成本与模型性能，找到最优的后训练数据量，成为一个亟待解决的问题。本文通过研究剪枝后模型的后训练损失变化，提出了P² Law这一缩放定律，旨在预测后训练损失并指导资源分配。

Method

本文提出了一种新的缩放定律——P² Law，用于预测剪枝后大型语言模型的后训练损失。其核心思想和实现步骤如下：

核心思想：基于Chinchilla缩放定律，结合剪枝前模型规模（ $N_0$ ）、后训练token数量（ $D$ ）、剪枝率（ $\rho$ ）以及剪枝前模型损失（ $\mathcal{L}_0$ ），构建一个数学模型来预测剪枝后模型的后训练损失（ $\mathcal{L}$ ）。
公式定义：P² Law最终参数化为： $\mathcal{L}(N_0, D, \rho, \mathcal{L}_0) = \mathcal{L}_0 + \left(\frac{1}{\rho}\right)^{\gamma} \left(\frac{1}{N_0}\right)^{\delta} \left(\frac{N_C}{N_0^{\alpha}} + \frac{D_C}{D^{\beta}} + E\right)$ 其中， $N_C, D_C, E, \alpha, \beta, \gamma, \delta$ 为常数，通过实验拟合得到。
必要条件：P² Law需满足三个条件：(1) 后训练损失随token数量增加而减少；(2) 在相似剪枝率下，小模型的损失下降更快；(3) 相对后训练损失与剪枝率呈幂律关系。
评估指标：除传统R²和Huber损失外，提出新指标ASD（平均斜率差异），用于评估预测损失曲线斜率与实际曲线的匹配度，强调预测后训练数据量拐点的准确性。
参数化选择：通过对比三种候选参数化形式，基于R²、Huber损失和ASD指标，选择最优形式（ $\mathcal{L}_1$ ）。 批判性思考：P² Law的数学形式虽然基于Chinchilla定律，但引入多个参数（如 $\gamma, \delta$ ）增加了拟合复杂度，可能导致过拟合风险，尤其是在数据量有限的情况下。此外，ASD指标仅关注后半段训练过程的斜率，可能忽略早期训练动态对整体趋势的影响，指标设计合理性有待进一步验证。

Experiment

实验在Llama-3和Qwen-2.5系列模型上进行，涵盖多种模型规模（0.5B至8B参数），并采用三种剪枝方法：深度剪枝、宽度剪枝和2:4半结构化剪枝。

实验设置：剪枝率在15%-35%之间变化，后训练数据从SlimPajama数据集中抽取（0.5B或1B token），使用4台Nvidia A800-80G和4台A6000-48G GPU进行训练，总耗时500小时。
结果分析：(1) P² Law在大多数情况下能准确拟合实际后训练损失曲线，尤其是在深度剪枝和2:4半结构化剪枝下；(2) 泛化实验显示，P² Law在更大数据集规模和更高剪枝率下表现良好（R²接近0.97，ASD较低）；(3) 但在模型规模泛化上，R²和Huber损失表现较差（R²甚至为负），仅ASD指标显示预测曲线斜率与实际趋势一致；(4) Llama-3.1-8B在宽度剪枝下的异常表现未被P² Law有效预测。
实验设计评价：实验覆盖了多种模型和剪枝方法，设置较为全面，但受限于GPU资源，未能验证更大规模模型（>8B）的适用性。此外，异常结果（如Llama-3.1-8B宽度剪枝）未深入分析原因，可能影响定律的普适性。泛化实验中，模型规模的预测结果不佳，表明P² Law在跨规模泛化上的局限性，实验证据支持力度不足。
结果与预期匹配度：P² Law在部分场景下（如数据集规模和剪枝率泛化）与预期一致，但在模型规模泛化和异常场景下表现未达预期，提示其理论和应用范围需进一步完善。

Further Thoughts

P² Law的提出为剪枝后模型的后训练资源分配提供了一个有价值的理论框架，但其局限性也启发了一些进一步的研究方向。首先，论文未能在更大规模模型（>8B参数）上验证定律的适用性，而当前LLM的发展趋势（如GPT-4、Llama-70B等）表明，大规模模型的剪枝和后训练可能是更关键的应用场景，未来研究应优先解决计算资源限制，验证P² Law在这些场景下的表现。其次，异常结果（如Llama-3.1-8B宽度剪枝）的出现提示不同模型架构和剪枝方法可能对后训练损失有复杂影响，是否可以通过引入架构特异性参数或剪枝方法特异性调整来增强P² Law的适应性？此外，P² Law与现有缩放定律（如Chinchilla）的关系值得深入探讨，是否可以构建一个统一的缩放框架，同时涵盖预训练和剪枝后训练的损失预测？最后，结合其他领域的研究，如神经网络剪枝在计算机视觉中的应用，是否能为LLM剪枝后训练提供新的视角，例如借鉴视觉模型中基于梯度敏感性的剪枝策略来优化后训练数据选择？这些问题值得进一步探索，以提升P² Law的理论深度和实际应用价值。