From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

技术摘要：从层到子模块 (SUBFIT)

问题陈述
大语言模型（LLM）的训练后压缩旨在通过移除架构组件来降低部署成本。虽然删除整个层或连续块的方法很有效，但它们会破坏残差流（residual stream），从而导致质量剧烈下降。近期的“基于替换”的方法通过使用轻量级的拟合模块来替代被移除的组件来解决这一问题。然而，现有方法存在两个限制性的设计约束：

1. 全层粒度： 它们将整个 Transformer 层视为移除的原子单元。
2. 连续选择： 它们在连续的深度范围内移除组件。

作者认为这些约束过于严格。预训练 Transformer 中的冗余并不局限于连续区域，也不在 Attention 和 FeedForward (FFN) 输出之间均匀分布。因此，需要不同的策略来近似不同的子模块类型，且可移除的组件不需要聚集在连续的深度范围内。

方法论：SUBFIT
本文介绍了 SUBFIT（子模块级拟合残差替换），这是一种在单个 Attention 和 FFN 子模块粒度上运行的训练后压缩方法。该过程包含两个主要阶段：

1. 非连续选择：

   • SUBFIT 根据稀疏预算 $s$，在模型深度范围内独立地选择目标数量的 Attention 和 FFN 子模块。
   • 选择是非连续的；模块根据评分独立选择，而不形成连续块。
   • 评分策略：
     • Attention： 模块根据在校准激活值上计算的“影响得分”进行排名，该得分衡量了子模块旋转残差流的程度及其相对幅度。
     • FFN： 对于基座模型，使用结合了移除效应和代理误差的替换感知得分；对于指令微调模型，使用仅基于余弦相似度的得分，以避免来自校准不良的残差幅度的噪声。
   • 排序： 先选择并替换 Attention 子模块。随后在部分压缩后的模型上对 FFN 子模块进行评分和选择，以补偿由 Attention 替换引起的残差流偏移。

2. 拟合残差替换：

   • 每个选定的子模块都被替换为一个轻量级的参数化代理 $\hat{F}_\ell$，该代理旨在近似原始子模块的残差贡献（$y = F_\ell(x)$），而非其内部计算。
   • 代理结构： 代理使用仿射校正（逐元素门控、偏置和中心化）以及由低秩算子 $U_\ell = V_\ell^\top W_\ell$ 定义的结构化交互项。
   • 非对称参数化：
     • Attention： 每个被移除的 Attention 子模块都接收其独立的低秩基底 ($V_\ell$) 和参数，以反映不同深度下的异构角色。
     • FFN： 在所有选定的 FFN 层中共享一个输入基底 $V$。只有输出因子 ($W_\ell$) 和仿射项保持层特异性。这种跨层参数共享显著降低了部署成本。
   • 拟合： 通过在校准统计数据上进行解析拟合（岭正则化最小二乘法和特征分解）来恢复参数，无需梯度优化或全模型微调。

核心贡献

1. 识别约束： 作者指出了现有基于替换的压缩方法中全层粒度和连续选择的局限性。
2. SUBFIT 框架： 一种在子模块级（Attention 和 FFN）进行联合选择与替换的新颖方法，并采用非连续选择。
3. 非对称设计： 一种特定的架构，其中 Attention 旁路使用独立的基底，而 FFN 旁路共享通用的输入基底以最小化参数开销。
4. 全面评估： 在 10 个 LLM（5 个基座模型，5 个指令微调模型）和 5 个稀疏水平（12.5% 到 37.5%）上进行了广泛测试，证明了其相对于四个基准方法的改进。

实验结果

• 困惑度（基座模型）： 在 10 个模型-稀疏度设置中的 9 个中，SUBFIT 实现了最低的困惑度退化。在激进压缩下（例如在 25% 稀疏度时，SUBFIT 的困惑度退化为 2.42×，而最强基准为 4.34×），其相对于最强基准的优势显著扩大。
• 下游任务准确率（指令微调模型）： 在 25% 稀疏度下，SUBFIT 平均保留了 84.6% 的密集模型下游准确率，优于最强基准（81.6%）。它是唯一能在 25% 稀疏度下保持 80% 以上准确率，以及在 37.5% 稀疏度下保持 73% 以上准确率的方法。
• 稳定性： 与表现出不稳定性（例如在高稀疏度下标准差增加数个数量级）的基准相比，SUBFIT 的困惑度退化方差显著较低。
• 效率： 该方法带来了可衡量的推理加速（首字延迟/time-to-first-token 提升 1.18×–1.40×）和 KV 缓存节省（与稀疏度成比例），且无需特殊的稀疏内核或定制硬件。
• 消融研究：
  • 残差补偿： 替换 FFN 是稳定性的主导因素；同时替换 Attention 和 FFN 会产生最佳结果。
  • 共享基底： 共享 FFN 基底实现的性能与逐层基底相当，但参数量显著减少。
  • 秩敏感性： FFN 的质量在秩达到 4096 时提升，而 Attention 在秩为 256 时趋于饱和，这证明了非对称秩设置的合理性。

意义与主张
本文主张粒度是训练后压缩中一个尚未得到充分探索的维度。通过从层级替换转向子模块级替换，SUBFIT 在相同的参数预算下比连续的全层替换恢复了更多的密集模型行为。结果表明，非连续、特定于子模块的恢复提供了更强的综合质量-效率权衡，尤其是在激进压缩的情况下。该方法旨在易于部署（闭式拟合），并兼容现有的推理栈。作者谦虚地建议，这种方法可能会推广到其他移除单元，如 Attention 头或单个神经元，但在本研究中并未声称能立即推广到这些特定单元。