利用组合优化进行神经网络修剪

现代神经网络在语言、数学推理和视觉 等各种应用领域都取得了令人瞩目的成绩。然而，这些网络通常使用需要大量计算资源的大型架构。这可能使得向用户提供此类模型变得不切实际，尤其是在可穿戴设备和智能手机等资源受限的环境中。一种广泛用于降低预训练网络推理成本的方法是通过删除部分权重来对其进行剪枝，但这种方法不会显著影响效用。在标准神经网络中，每个权重定义两个神经元之间的连接。因此，在剪枝后，输入将通过一组较小的连接传播，从而需要更少的计算资源。

原始网络与修剪后的网络。

剪枝方法可应用于网络训练过程的不同阶段：训练后、训练中或训练前（即，在权重初始化之后立即进行）。在本文中，我们将重点介绍训练后设置：给定一个预先训练好的网络，我们如何确定应该剪枝哪些权重？一种流行的方法是幅度剪枝，它可以删除幅度最小的权重。虽然有效，但这种方法并不直接考虑删除权重对网络性能的影响。另一种流行的范例是基于优化的剪枝，它根据删除权重对损失函数的影响程度来删除权重。虽然在概念上很有吸引力，但大多数现有的基于优化的方法似乎面临着性能和计算要求之间的严重权衡。进行粗略近似的方法（例如，假设对角线Hessian 矩阵）可以很好地扩展，但性能相对较低。另一方面，虽然进行较少近似的方法往往表现更好，但它们的可扩展性似乎要差得多。

在ICML 2023上发表的 “像 CHITA 一样快：使用组合优化的神经网络修剪”中，我们描述了如何开发一种基于优化的方法来大规模修剪预训练的神经网络。CHITA（代表“组合无 Hessian 迭代阈值算法”）在可扩展性和性能权衡方面优于现有的修剪方法，它通过利用多个领域的进步来实现这一点，包括高维统计、组合优化和神经网络修剪。例如，CHITA 的速度可以比最先进的ResNet修剪方法快 20 倍到 1000 倍，并且在许多设置中将准确率提高了 10% 以上。

贡献概述

与流行方法相比，CHITA 有两项显著的技术改进：

高效利用二阶信息：使用二阶信息（即与二阶导数相关）的修剪方法在许多情况下都达到了最先进的水平。在文献中，此信息通常用于计算 Hessian 矩阵或其逆，这是一种很难扩展的运算，因为 Hessian 大小与权重数量成二次函数关系。通过仔细的重新表述，CHITA 使用二阶信息，而无需明确计算或存储 Hessian 矩阵，从而实现了更高的可扩展性。

组合优化：流行的基于优化的方法使用一种简单的优化技术，即单独修剪权重，即在决定修剪某个权重时，它们不会考虑其他权重是否已被修剪。这可能会导致修剪重要的权重，因为单独看来不重要的权重可能会在其他权重被修剪时变得重要。CHITA 通过使用更先进的组合优化算法来避免这个问题，该算法考虑了修剪一个权重对其他权重的影响。

在下面的章节中，我们讨论 CHITA 的修剪公式和算法。

计算友好的修剪公式

有许多可能的剪枝候选，它们通过仅保留原始网络中权重的子集而获得。令k为用户指定的参数，表示要保留的权重数量。剪枝可以自然地表述为最佳子集选择(BSS) 问题：在所有可能的剪枝候选（即权重子集）中，仅保留k个权重，选择损失最小的候选。

修剪作为 BSS 问题：在所有具有同样总权重数量的可能修剪候选者中，最佳候选者定义为损失最小的候选者。此图显示了四个候选者，但这个数字通常要大得多。

解决原始损失函数上的剪枝 BSS 问题通常计算起来很困难。因此，与OBD和OBS等以前的工作类似，我们使用二阶泰勒级数用二次函数近似损失，其中 Hessian 用经验Fisher 信息矩阵估计。虽然通常可以高效地计算梯度，但由于 Hessian 矩阵的规模庞大，计算和存储 Hessian 矩阵的成本过高。在文献中，通常通过对 Hessian（例如对角矩阵）和算法（例如单独剪枝权重）做出限制性假设来应对这一挑战。

CHITA 使用了一种有效的剪枝问题 (BSS 使用二次损失) 重构，避免了明确计算 Hessian 矩阵，同时仍然使用该矩阵的所有信息。这是通过利用经验 Fisher 信息矩阵的低秩结构实现的。这种重构可以看作是一个稀疏线性回归问题，其中每个回归系数对应于神经网络中的某个权重。在获得该回归问题的解后，设置为零的系数将对应于应剪枝的权重。我们的回归数据矩阵为 ( n x p )，其中n是批次 (子样本) 大小，p是原始网络中的权重数量。通常n << p ，因此存储和操作该数据矩阵比使用 ( p x p ) Hessian的常见剪枝方法更具可扩展性。

CHITA 将二次损失近似（需要昂贵的 Hessian 矩阵）重新表述为线性回归 (LR) 问题。LR 的数据矩阵在p中是线性的，这使得重新表述比原始二次近似更具可扩展性。

可扩展的优化算法

CHITA 将剪枝问题简化为以下稀疏性约束下的线性回归问题：最多k 个回归系数可以为非零值。为了解决这个问题，我们考虑对众所周知的迭代硬阈值(IHT) 算法进行修改。IHT 执行梯度下降，每次更新后都会执行以下后处理步骤：将 Top- k之外的所有回归系数（即幅度最大的k 个系数）设置为零。IHT 通常可以很好地解决问题，它会迭代探索不同的剪枝候选并联合优化权重。

由于问题规模较大，具有恒定学习率的标准 IHT可能会出现收敛速度非常慢的问题。为了加快收敛速度，我们开发了一种新的线搜索方法，利用问题结构来寻找合适的学习率，即导致损失大幅下降的学习率。我们还采用了几种计算方案来提高 CHITA 的效率和二阶近似的质量，从而开发出一种改进版本，我们称之为 CHITA++。

实验

我们将 CHITA 的运行时间和准确性与使用不同架构的几种最先进的修剪方法（包括ResNet和MobileNet）进行了比较。

运行时间：CHITA 比执行联合优化（而不是单独修剪权重）的同类方法更具可扩展性。例如，在修剪 ResNet 时，CHITA 的速度可以达到 1000 倍以上。

修剪后准确率：下面，我们将 CHITA 和 CHITA++ 与幅度修剪 (MP)、 Woodfisher (WF) 和组合脑外科医生 (CBS)的性能进行比较，以修剪 70% 的模型权重。总体而言，我们看到 CHITA 和 CHITA++ 取得了良好的改进。

各种方法在 ResNet20 上的剪枝后准确率。报告了剪枝 70% 模型权重的结果。

MobileNet 上各种方法的后剪枝准确率。报告了剪枝 70% 模型权重的结果。

接下来，我们报告对较大网络 ResNet50 进行剪枝的结果（在此网络上，ResNet20 图中列出的一些方法无法扩展）。在这里，我们将其与幅度剪枝和M-FAC进行比较。下图显示，CHITA 在各种稀疏度水平下都实现了更好的测试准确率。

使用不同的方法获得的修剪网络的测试精度。

结论、局限性和未来工作

我们介绍了一种基于优化的预训练神经网络修剪方法 CHITA。CHITA 通过有效利用二阶信息并借鉴组合优化和高维统计的思想，提供了可扩展性和具有竞争力的性能。

CHITA 专为非结构化剪枝 而设计，可以移除任何权重。理论上，非结构化剪枝可以显著降低计算要求。然而，在实践中实现这些减少需要支持稀疏计算的特殊软件（可能还有硬件）。相比之下，结构化剪枝会移除整个结构（如神经元），可能会提供更容易在通用软件和硬件上实现的改进。将 CHITA 扩展到结构化剪枝将会很有趣。

致谢

这项工作是 Google 和麻省理工学院研究合作的一部分。感谢 Rahul Mazumder、Natalia Ponomareva、Wenyu Chen、Xiang Meng、Zhe Zhao 和 Sergei Vassilvitskii 在准备这篇文章和论文时提供的帮助。还要感谢 John Guilyard 为这篇文章制作图形。