稀疏性概述¶
稀疏性是一种从神经网络中移除参数以减少内存开销或延迟的技术。通过仔细选择如何剪枝元素，可以在内存开销和延迟方面实现显著的降低，同时在模型质量（准确性/f1）方面付出相对较低或不付出代价。
目标¶
我们认为当前稀疏性研究人员/用户面临的主要问题是碎片化。研究人员理所当然地希望展示端到端的结果，但这意味他们花费大量时间来研究如何与 PyTorch 集成以及解决诸如以下实现问题：
我应该何时进行掩码处理？
我应该何时/如何存储压缩表示？
我想要原地或非原地掩码更新吗？
我如何调用稀疏矩阵乘法而不是密集矩阵乘法？
我们认为上述问题可以通过 torchao 一次性解决，让研究人员专注于真正重要的事情——推动稀疏内核性能或更精确的剪枝算法。
更具体地说，我们希望为稀疏内核（张量子类）和剪枝算法（torch.ao.pruning.Sparsifier）提供教程和 API，供用户扩展。我们的目标是提供模块化的构建块，不仅可以加速推理，还可以加速训练，并且可以很好地与 torchao 量化工作流程结合使用。
从头开始训练稀疏模型，并进行硬件加速，同时将精度损失降至最低。
通过自定义剪枝算法恢复剪枝模型的精度损失。
在支持稀疏性的硬件上加速掩码/剪枝模型以实现性能提升。
设计¶
稀疏性，像量化一样，是准确性/性能的权衡，我们不仅关心加速，还关心架构优化技术的准确性下降。
在量化中，理论性能增益通常由我们量化到的数据类型决定——从 float32 量化到 float16 会产生理论上的 2 倍加速。对于剪枝/稀疏性，类比变量将是稀疏度级别/稀疏度模式。对于半结构化稀疏性，稀疏度级别固定为 50%，因此我们预期理论上会提高 2 倍。对于块稀疏矩阵和非结构化稀疏性，加速是可变的，取决于张量的稀疏度级别。
稀疏性和量化之间的一个主要区别在于如何确定准确性下降：通常，量化的准确性下降由选择的比例和零点决定。然而，在剪枝中，准确性下降由掩码决定。稀疏性和量化密切相关，并共享诸如量化/稀疏感知训练等精度缓解技术。
通过仔细选择指定的元素并重新训练网络，剪枝可以实现可忽略不计的精度下降，在某些情况下甚至可以提供轻微的精度提升。这是一个活跃的研究领域，目前尚未达成共识。我们希望用户能够有一个目标稀疏模式，并将其剪枝到该模式。
给定一个目标稀疏模式，剪枝/稀疏化模型可以被视为两个独立的子问题
准确性 - 我如何找到一组满足目标稀疏模式的稀疏权重，以最小化模型的准确性下降？
性能 - 我如何加速我的稀疏权重进行推理并减少内存开销？
我们的工作流程设计为包含两个部分，分别回答每个问题：
一个前端 Python 用户友好的 API，用于查找任何任意稀疏模式的稀疏权重。
一个后端稀疏内核/操作的集合，用于减少内存/延迟。
这两个部分之间的交接点是以密集格式存储的稀疏权重，其中缺失元素的位置为 0。这是一个自然的交接点，因为稀疏矩阵乘法和使用此张量的密集矩阵乘法在数值上是等效的。这让我们为后端提供了一个清晰的用户契约，针对给定的稀疏模式：
如果您的密集矩阵可以转换为 **2:4 稀疏格式**，我们可以将矩阵乘法加速至 **1.7 倍**，且无数值损失。
这也使得现有密集格式稀疏权重的用户能够利用我们快速的稀疏核。我们预计许多用户会提出自己的自定义前端掩码解决方案或使用其他第三方解决方案，因为这是一个活跃的研究领域。
下面，我们提供了一个使用 PyTorch API 加速具有 2:4 稀疏性 + bf16 的模型的示例。
import torch
from torch.sparse import to_sparse_semi_structured, SparseSemiStructuredTensor
from torch.ao.pruning import WeightNormSparsifier
# bfloat16 CUDA model
model = model.half().cuda()
# Accuracy: Finding a sparse subnetwork
sparse_config = []
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear):
sparse_config.append({"tensor_fqn": f"{name}.weight"})
sparsifier = WeightNormSparsifier(sparsity_level=1.0,
sparse_block_shape=(1,4),
zeros_per_block=2)
# attach FakeSparsity
sparsifier.prepare(model, sparse_config)
sparsifier.step()
sparsifier.squash_mask()
# now we have dense model with sparse weights
# Performance: Accelerated sparse inference
for name, mod in model.named_modules():
if isinstance(mod, torch.nn.Linear):
mod.weight = torch.nn.Parameter(to_sparse_semi_structured(mod.weight))
从根本上讲，该流程通过操作 torch.Tensors 来实现。在前端，我们通过其完全限定名在 sparse_config 字典中指定张量。前端旨在遵循量化 API，包含一个 prepare 函数，该函数将 FakeSparsity 参数化附加到配置中指定的张量。
FakeSparsity 是一种参数化，它模拟非结构化稀疏性，其中每个元素都有一个掩码。因此，我们可以用它来模拟我们想要的任何稀疏模式。
用户将使用自己的自定义代码训练准备好的模型，并在必要时调用 .step() 来更新掩码。一旦找到合适的掩码，他们就会调用 squash_mask() 将掩码融合到权重中，创建一个在正确位置带有 0 的密集张量。
然后，用户将通过使用量化流程进行量化块稀疏 CPU 推理或在指定的权重张量上调用 to_sparse_semi_structured 来转换模型以进行加速稀疏推理。
背景¶
本节提供了一些关于神经网络剪枝/稀疏性的背景，以及一些常见剪枝/稀疏性术语的定义。在学术界/工业界，**剪枝**和**稀疏性**通常互换使用，指代同一事物。这可能会令人困惑，尤其是因为稀疏性是一个含义过载的术语，可以指代许多其他事物，例如稀疏张量表示。
请注意，本节侧重于**剪枝**，而不是**稀疏训练**。区别在于，在**剪枝**中，我们从预训练的密集模型开始，而在**稀疏训练**中，我们从头开始训练稀疏模型。
为了避免混淆，我们通常尝试使用稀疏性来指代张量。请注意，稀疏张量可以指具有许多零值的密集张量，也可以指使用稀疏表示存储的张量。我们将该流程描述为**剪枝**，并将所得模型描述为**剪枝**模型。
大致而言，实现性能更好的剪枝模型的工作流程如下：
剪枝的一般思想是，我们可以遮蔽训练好的神经网络中的一些权重，并恢复任何精度损失。由此产生的剪枝模型可以在优化内核上运行，这些内核利用这种稀疏性来加速推理。
将剪枝参数归零并不会直接影响模型的延迟/内存开销。这是因为密集张量本身仍然包含剪枝元素（0 元素），并且在矩阵乘法期间仍会使用这些元素进行计算。为了实现性能提升，我们需要将密集核替换为稀疏核。
从广义上讲，这些稀疏表示允许我们跳过涉及剪枝元素的计算，以加速矩阵乘法。为此，这些优化的稀疏内核在以更高效格式存储的稀疏矩阵上工作。一些稀疏张量布局与特定后端紧密耦合，例如 NVIDIA 2:4，而另一些则更通用，并受多个后端支持（CSC 受 FBGEMM 和 QNNPACK 支持）。
| 名称 | 描述 | 稀疏矩阵的存储方式 |
| COO (sparse_coo) | 用于存储稀疏矩阵的坐标格式。矩阵存储为非稀疏数据向量和这些元素在密集矩阵中的索引位置的组合。 | 稀疏矩阵 = {索引：坐标位置张量，数据：对应索引位置的值张量} |
| BSR (sparse_bsr) | 用于存储稀疏矩阵的块稀疏行格式。矩阵存储为数据块以及这些块在密集矩阵中的索引位置。与 COO 非常相似，只是单个数据由块组成，而不是标量。 | 稀疏矩阵 = {索引：坐标位置的张量，对于矩阵是二维的；数据：对应于索引位置的块张量}，其中块是对应于稀疏模式的矩阵。 |
| CSR (sparse_csr) / CSC (sparse_csc) | 压缩稀疏行/列格式，用于存储稀疏矩阵。稀疏矩阵存储为列/行上的数据块，以及这些行/列在密集矩阵中的索引。这是存储块稀疏矩阵最紧凑的格式。 | 稀疏矩阵 = {索引：1D 列索引张量，索引指针：1D 张量，指定列在行中的开始和结束索引（从第 0 行开始），数据：对应于索引位置的块张量}。 |
| NVIDIA 2:4 压缩表示 | 用于 2:4 半结构化稀疏性的定制 NVIDIA 压缩存储格式。我们将稀疏矩阵存储为一个压缩密集矩阵（原始大小的一半），其中包含未剪枝的元素和一个位掩码索引。当将我们的稀疏矩阵乘以另一个密集矩阵时，我们使用掩码索引到密集矩阵并与我们的压缩密集矩阵相乘。 | 稀疏矩阵 = {位掩码：剪枝元素的 2 位索引，压缩密集矩阵：包含所有未剪枝的元素，大小为原始密集矩阵的一半} |
表 4.1：常见稀疏张量布局概述。
虽然剪枝的一般想法很简单，但用户在成功剪枝模型之前必须弄清楚许多细节。
这些可以大致分解如下：
**剪枝配置** - 我应该剪枝哪些层？我应该剪枝到什么稀疏度级别？
**剪枝标准** - 我应该如何决定移除哪些参数？
**剪枝策略** - 移除参数后，我如何恢复任何精度下降？
**稀疏模式** - 当我剪枝模型时，我是否应该尝试使用特定的稀疏模式？不同的硬件后端支持不同稀疏模式的加速推理。
剪枝配置¶
神经网络中的所有层并非都同等重要。有些层可能比其他层对剪枝更敏感。用户必须决定剪枝哪些层，以及每个层的**稀疏度级别**，即该权重张量中 0 的百分比。剪枝配置对剪枝模型的精度和加速都有影响。
确定给定模型的最佳剪枝配置和稀疏度级别是一个开放性问题，目前尚无通用解决方案。这部分原因在于最佳剪枝配置取决于后续的剪枝标准和策略，而且决定如何剪枝模型以及如何恢复损失的精度有无限种方法。
一种确定剪枝哪些层以及剪枝程度的常用方法是进行敏感性分析，通过以不同稀疏度级别剪枝模型中的每一层，并查看随后的精度下降（不重新训练）。这为用户提供了每层的稀疏度-精度曲线，用户可以将其作为确定最佳剪枝配置的代理。
剪枝标准¶
用户必须决定从神经网络中移除参数的标准。就像确定最佳剪枝配置一样，确定最佳剪枝标准是一个开放的研究问题，并且取决于前面提到的其他因素。
最常见的剪枝标准是使用权重大小。其思想是，低幅值权重对模型输出的贡献小于高幅值权重。如果我们想要移除参数，我们可以移除绝对值最小的权重。
然而，即使采用简单的剪枝标准，如权重大小，用户仍需考虑其他因素：
局部与全局范围
**局部范围**意味着稀疏性掩码仅根据层统计信息计算。
优点：掩码计算简单。
缺点：精度与稀疏性的权衡可能次优。
**全局范围**意味着稀疏性统计不受单一层限制，必要时可跨多层。
优点：无需逐层阈值。张量统计信息在层间共享，并通过层间归一化实现。
缺点：计算掩码时复杂性增加。
用于掩码计算的张量
**权重**：仅使用权重张量来计算掩码。这种方法对于推理最简单，因为权重张量是常数。
**梯度**：根据权重和梯度范数计算重要性。常用于基于预训练的方法。目前 CTR_mobile_feed 使用基于梯度的剪枝算法。
**激活**：在一些研究论文中，与感兴趣的权重一起应用的激活的范数用于计算重要性得分。
原地或非原地掩码更新
**原地**更新稀疏张量，通过执行 W = W (Mask)。一旦权重张量更新，稀疏值就会被归零，无法恢复。
**优点**：只需要存储稀疏张量的一份副本（+掩码）
**缺点**：一旦掩码应用于权重，它就会被归零，所有历史记录都丢失了。这些权重无法重新增长。
**非原地**更新不直接修改张量，而是执行以下操作：W' = W (Mask) 和 dW' = dW (Mask)
**优点**：原始张量得以保留（被掩码的元素不会通过反向传播更新）。如果掩码改变，权重可以重新增长。这对于 PAT 是必需的。
**缺点**：除了未掩码的权重 (W) 之外，掩码的权重 (W') 也会被计算并驻留在内存中以进行前向/后向计算。
| 名称 | 描述 | 注意事项 |
| 幅度/显著性 | 移除范数最低的参数（常用 L1 范数） | 已证明对 2:4 半结构化稀疏性效果良好。通过一次性幅度剪枝后重复训练循环，能够实现与原始模型相同的准确性。 |
| 运动剪枝 | 这些方法旨在利用梯度信息来决定要删除哪些参数。其思想是删除在微调过程中变化不大的参数。 | 常见于预训练模型。 |
| 低秩分解 | 这些方法旨在用 SQx 替换 Wx，其中 S 和 Q 是低秩矩阵。 | 通常这些方法使用某种形式的逐层重建，其中不是通过损失函数训练模型来恢复丢失的精度，而是寻求匹配逐层统计信息（找到 SQx 以最小化 L2(SQx, Wx)）。 |
| 随机 | 随机删除参数 |
表 4.2：一些常见剪枝标准的描述。
剪枝策略¶
这是一个通用术语，描述了用户试图从其剪枝模型中恢复任何精度下降的方法。在剪枝模型后，通常会看到模型的精度下降，因此用户通常会重新训练剪枝模型以纠正此问题。剪枝策略还决定了在模型训练期间何时以及多久剪枝模型。
剪枝策略和剪枝标准之间的界限并不明确，尤其是在剪枝感知训练方法（在训练期间更新掩码）的情况下。我们有时使用**剪枝** **算法**一词来指代这两个项目的组合。这两个因素，连同剪枝配置，最终控制着剪枝模型的最终精度。
| 剪枝策略 | 描述 | 注意事项 |
| 零次剪枝 | 剪枝一次，不重新训练模型 | 这些方法依赖于更复杂的剪枝标准。 这在文献中有时被称为“一次性”，但我们将“一次性”用于指剪枝一次并重新训练一次。 |
| 一次性 | 剪枝一次，重新训练模型一次 | 英伟达已证明一次性 2:4 半结构化稀疏剪枝在各种常见视觉/自然语言处理模型中具有良好的泛化能力。\\ 重新训练策略是简单地重复训练过程。 |
| 迭代 | 剪枝模型，重新训练，重复 | 我们可以迭代地增加稀疏度级别，或迭代地剪枝模型中的不同层。 |
| 剪枝感知训练 | 掩码在训练期间学习 | 由 CTR_feed 用于其当前的剪枝算法。 |
| NAS / 多掩码 | 训练期间使用多个掩码。这可以被认为是神经架构搜索的一种形式。 | 由 PySpeech (FastNAS) 使用 |
| 逐层重建 | 我们不是通过损失函数进行再训练，而是尝试通过使用类似于知识蒸馏的双模型方法，从每一层中恢复尽可能多的信息。 | 请参阅 https://arxiv.org/pdf/2204.09656.pdf |
表 4.3：一些常见剪枝策略的描述。
稀疏模式¶
稀疏模式描述了剪枝参数在模型/张量中的排列方式。
回想一下，通常需要使用优化的稀疏核才能获得性能提升。根据权重张量的格式和稀疏度级别，稀疏矩阵乘法可能比其密集对应物更快。如果张量不够稀疏，它也可能更慢。
在最一般的层面上，剪枝是非结构化的——每个参数都有自己的掩码。这提供了最大的灵活性，但需要非常高的稀疏度（>98%）才能提供性能优势。为了在较低稀疏度级别提供加速推理，硬件后端已增加了对特殊稀疏模式的支持。
我们力求剪枝模型，使权重张量呈现与我们的推理后端相同的稀疏模式。如果我们能够恢复丢失的精度，同时保持稀疏模式，我们就可以在稀疏硬件上运行此模型以加速推理，而无需精度损失。我们也可以在目标后端上运行剪枝到不同稀疏模式的模型，但会牺牲一些额外的精度损失。
特定的后端硬件及其相应的稀疏模式，以及剪枝配置最终决定了我们观察到的性能提升。如果使用不同的剪枝标准剪枝模型，只要它遵循相同的稀疏模式和稀疏级别，它将具有相同的性能特征。例如，如果我们决定移除最高幅度的权重而不是最低幅度的权重，我们不会预期这会改变剪枝模型的性能特征。
| 稀疏模式 | 掩码可视化
（50% 稀疏度） |
||||||||||||||||||||||||||||||||
| 非结构化稀疏性 |
|
||||||||||||||||||||||||||||||||
| 2:4 半结构化 |
|
||||||||||||||||||||||||||||||||
| 块稀疏性 |
|
||||||||||||||||||||||||||||||||
| 结构化稀疏性 |
|
表 4.4：一些常见稀疏模式的描述。
有关我们支持的 API 和基准测试的更多信息，请参阅 稀疏性 README。
