稀疏性概述¶
稀疏性是一种从神经网络中移除参数的技术，旨在降低其内存开销或延迟。通过仔细选择元素的修剪方式，可以在不显著牺牲模型质量（准确率/f1 分数）的情况下，大幅减少内存开销和延迟。
目标¶
我们认为，当前稀疏性研究人员/用户面临的主要问题是碎片化。研究人员理应致力于展示端到端的结果，但这通常意味着需要花费大量时间来研究如何与 PyTorch 集成以及解决诸如以下实现问题：
应该何时应用掩码？
我应该何时/如何存储压缩表示？
我想要就地更新还是非就地更新掩码？
如何调用稀疏矩阵乘法而不是密集矩阵乘法？
我们认为，上述问题可以由 torchao 一次性解决，让研究人员专注于真正重要的事情——提升稀疏核性能或改进修剪算法的准确性。
更具体地说，我们希望为稀疏核（张量子类化）和修剪算法（torch.ao.pruning.Sparsifier）提供教程和 API，供用户扩展。我们旨在提供模块化的构建块，这些构建块不仅可以用于加速推理，还可以用于加速训练，并且可以与 torchao 的量化工作流程良好地结合。
从头开始训练带有硬件加速的稀疏模型，并实现最小的准确率损失。
使用自定义修剪算法恢复修剪模型的准确率损失。
在支持稀疏性的硬件上加速掩码/修剪模型，以实现性能提升。
设计¶
稀疏性，就像量化一样，是一种准确率/性能的权衡，我们不仅关心加速效果，也关心架构优化技术带来的准确率下降。
在量化中，理论性能增益通常取决于我们量化到的数据类型——从 float32 量化到 float16 可带来理论上的 2 倍加速。对于修剪/稀疏性，类似的变量是稀疏度水平/稀疏模式。对于半结构化稀疏性，稀疏度水平固定为 50%，因此我们期望理论上实现 2 倍的提升。对于块稀疏矩阵和非结构化稀疏性，加速效果是可变的，取决于张量的稀疏度水平。
稀疏性和量化之间的一个关键区别在于准确率下降的决定因素：通常，量化的准确率下降由选择的 scale 和 zero_point 决定。然而，在修剪中，准确率下降由掩码决定。稀疏性和量化密切相关，并共享诸如量化/稀疏感知训练之类的准确率缓解技术。
通过仔细选择指定的元素并重新训练网络，修剪可以实现微不足道的准确率下降，在某些情况下甚至能略微提升准确率。这是一个活跃的研究领域，尚未达成一致共识。我们期望用户心中有一个目标稀疏模式，并按照该模式进行修剪。
给定一个目标稀疏模式，修剪模型可以被视为两个独立的子问题：
准确率 - 我如何找到一组满足目标稀疏模式的稀疏权重，从而最大程度地减少模型的准确率下降？
性能 - 我如何加速稀疏权重的推理并减少内存开销？
我们的工作流程设计包含两个独立回答这些问题的部分：
一个面向用户的 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 来转换模型以进行加速稀疏推理。
背景¶
本节提供了一些关于神经网络修剪/稀疏性的背景信息，并定义了一些常见的修剪/稀疏性术语。在学术界/工业界，修剪（pruning） 和 稀疏性（sparsity） 经常互换使用，指代同一件事。这可能会引起混淆，特别是因为稀疏性是一个具有多重含义的术语，可以指代许多其他事物，例如稀疏张量表示。
请注意，本节侧重于修剪（pruning），而非稀疏训练（sparse training）。区别在于，在修剪中，我们从预训练的密集模型开始；而在稀疏训练中，我们从头开始训练一个稀疏模型。
为了避免混淆，我们通常尽量使用稀疏性来指代张量。请注意，稀疏张量可以指代包含许多零值的密集张量，或使用稀疏表示存储的张量。我们将整个流程描述为修剪，将由此产生的模型描述为修剪模型（pruned model）。
粗略地说，实现性能更好的修剪模型的流程如下：
修剪背后的总体思想是，我们可以对训练好的神经网络的某些权重进行掩码处理，并恢复由此带来的准确率损失。由此产生的修剪模型可以在利用这种稀疏性进行加速推理的优化核上运行。
直接将修剪后的参数置零并不会影响模型的延迟/内存开销。这是因为密集张量本身仍然包含被修剪的元素（那些 0 元素），并在矩阵乘法期间仍会使用这些元素进行计算。为了实现性能提升，我们需要将密集核替换为稀疏核。
粗略地说，这些稀疏表示允许我们跳过涉及修剪元素的计算，以加快矩阵乘法。为此，这些优化的稀疏核处理的是以更高效格式存储的稀疏矩阵。某些稀疏张量布局与特定后端紧密耦合，例如 NVIDIA 2:4，而另一些则更通用，受多个后端支持（CSC 受 FBGEMM 和 QNNPACK 支持）。
| 名称 | 描述 | 稀疏矩阵的存储方式 |
| COO (sparse_coo) | 用于存储稀疏矩阵的 COOrdinate 格式。矩阵存储为非稀疏数据向量和这些元素在密集矩阵中的索引位置的组合。 | 稀疏矩阵 = {Index: 坐标位置张量, Data: 与索引位置对应的值张量 } |
| BSR (sparse_bsr) | 用于存储稀疏矩阵的块稀疏行格式。矩阵存储为数据块和这些块在密集矩阵中的索引位置。与 COO 非常相似，区别在于单个数据由块组成，而非标量。 | 稀疏矩阵 = {Index: 坐标位置张量（对于矩阵是二维的）, Data: 与索引位置对应的块张量 } 其中块是与稀疏模式对应的矩阵。 |
| CSR (sparse_csr) / CSC (sparse_csc) | 用于存储稀疏矩阵的压缩稀疏行/列格式。稀疏矩阵存储为列/行上的数据块以及这些行/列在密集矩阵中的索引。这是存储块稀疏矩阵最紧凑的格式。 | 稀疏矩阵 = {Index: 列索引的 1D 张量, IndexPtr: 指定行（从第 0 行开始）的列开始和结束索引的 1D 张量, Data: 与 Index 位置对应的块张量。} |
| NVIDIA 2:4 压缩表示 | 适用于 2:4 半结构化稀疏性的自定义 NVIDIA 压缩存储格式。我们将稀疏矩阵存储为一个压缩的密集矩阵（½ 原大小），其中包含未修剪的元素和一个位掩码索引。当我们将稀疏矩阵与另一个密集矩阵相乘时，我们使用掩码来索引密集矩阵并与我们的压缩密集矩阵相乘。 | 稀疏矩阵 = {Bitmask: 修剪元素的 2bit 索引 压缩密集矩阵: 包含所有未修剪的元素，大小是原始密集矩阵的一半} |
表 4.1: 常见稀疏张量布局概述。
虽然修剪的总体思想相当简单，但在成功修剪模型之前，用户必须弄清楚许多细节。
这些可以大致分解如下：
修剪配置 - 我应该修剪哪些层？应该修剪到什么稀疏度水平？
修剪标准 - 我应该如何决定移除哪些参数？
修剪策略 - 移除参数后，我如何恢复任何准确率下降？
稀疏模式 - 修剪模型时，我应该尝试使用特定的稀疏模式吗？不同的硬件后端支持不同稀疏模式的加速推理。
修剪配置¶
并非神经网络中的所有层都一样。有些层可能比其他层对修剪更敏感。用户必须决定修剪哪些层以及每层的稀疏度水平，即该权重张量中 0 的百分比。修剪配置会影响修剪模型的准确率和加速效果。
确定给定模型的最佳修剪配置和稀疏度水平是一个开放问题，尚无通用解决方案。这部分是因为最优修剪配置取决于后续的修剪标准和策略，并且决定如何修剪模型以及如何恢复损失的准确率有无数种方法。
一种确定修剪哪些层以及修剪程度的常用方法是进行敏感性分析，即在不同稀疏度水平下修剪模型中的每一层，并观察随后的准确率下降（不进行重新训练）。这为用户提供了每层的稀疏度-准确率曲线，用户可以将其用作确定最佳修剪配置的代理。
修剪标准¶
用户必须决定从神经网络中移除参数的标准。就像确定最佳修剪配置一样，确定最佳修剪标准也是一个开放的研究问题，并且取决于上述其他因素。
最常见的修剪标准是使用权重幅值。其思想是，低幅值的权重对模型输出的贡献小于高幅值的权重。如果我们要移除参数，我们可以移除绝对值最小的权重。
然而，即使使用权重幅值这样简单的修剪标准，用户也必须考虑其他因素：
局部范围 vs 全局范围
局部范围意味着稀疏性掩码仅根据层的统计信息计算。
优点：掩码计算简单
缺点：准确率与稀疏度的权衡可能不是最优。
全局范围意味着稀疏性统计信息不受单个层的限制，如果需要，可以跨越多个层。
优点：无需逐层阈值。张量统计信息在层之间共享，并使用跨层归一化来实现。
缺点：计算掩码时复杂度增加。
用于掩码计算的张量
权重：仅使用权重张量来计算掩码。对于推理而言，这种方法最简单，因为权重张量是恒定的。
梯度：根据权重和梯度范数计算重要性。常见于基于预训练的方法。目前 CTR_mobile_feed 使用基于梯度的修剪算法。
激活值：在一些研究论文中，与相关权重一起应用的激活值的范数被用来计算重要性得分。
就地或非就地掩码更新
就地更新通过执行 W = W (Mask) 来更新稀疏张量。权重张量更新后，稀疏值被置零，无法恢复。
优点：只需存储一份稀疏张量（+ 掩码）
缺点：一旦掩码应用于权重，该权重即被置零，所有历史信息丢失。这些权重无法“重新生长”（regrow）。
非就地更新不直接修改张量，而是执行以下操作：W’ = W (Mask) 和 dW’ = dW (Mask)
优点：原始张量得以保留（被掩码的元素不会通过反向传播更新）。如果掩码改变，权重可以“重新生长”。这对 PAT 是必需的。
缺点：除了未被掩码的权重 (W)，还需要计算被掩码的权重 (W’)，并在前向/后向计算期间驻留在内存中。
| 名称 | 描述 | 注意事项 |
| 幅值 / 显著性 | 移除范数最低的参数（常用 L1 范数） | 已证明与 2:4 半结构化稀疏性配合良好。通过在一次性幅值修剪后重复训练循环，能够达到与原始模型相同的准确率。 |
| Movement Pruning | 这些方法旨在利用梯度信息来决定移除哪些参数。其思想是移除在微调过程中变化不大的参数。 | 常用于预训练模型。 |
| 低秩分解 | 这些方法旨在将 Wx 替换为 SQx，其中 S 和 Q 是低秩矩阵。 | 通常这些方法使用某种层级重建，其中不是通过训练模型来恢复损失的准确率，而是寻求匹配层级的统计信息（找到使得 L2(SQx, Wx) 最小化的 SQx）。 |
| 随机 | 随机移除参数 |
表 4.2: 一些常见修剪标准描述。
修剪策略¶
这是一个通用术语，描述了用户试图从修剪模型中恢复任何准确率下降的方法。修剪模型后，通常会看到模型准确率下降，因此用户通常会重新训练修剪后的模型以弥补这一点。修剪策略还决定了在模型训练期间何时以及多久进行一次修剪。
修剪策略和修剪标准之间的界限并不明确，尤其是在修剪感知训练方法中，它们在训练期间更新掩码。我们有时使用术语修剪算法（pruning algorithm）来指代这两者的组合。这两个因素，连同修剪配置，最终决定了修剪模型的最终准确率。
| 修剪策略 | 描述 | 注意事项 |
| 零样本 (Zero-shot) | 修剪一次，不重新训练模型 | 这些方法依赖于更复杂的修剪标准。 在文献中有时也称为“一次性（one-shot）”，但我们将“一次性（one-shot）”用于指修剪一次并重新训练一次的情况。 |
| 一次性 (One-shot) | 修剪一次，重新训练模型一次 | NVIDIA 已表明，一次性 2:4 半结构化稀疏修剪在各种常见视觉/NLP 模型上具有良好的泛化能力。 \ \ 重新训练策略就是简单地再次重复训练过程。 |
| 迭代式 (Iterative) | 修剪模型，重新训练，重复 | 我们可以迭代地增加稀疏度水平，或者迭代地修剪模型中的不同层。 |
| 修剪感知训练 (Pruning Aware Training) | 掩码在训练过程中学习 | CTR_feed 用于其当前修剪算法。 |
| NAS / 多掩码 (Multimask) | 训练期间使用多个掩码。这可以被认为是神经架构搜索的一种形式。 | PySpeech 使用 (FastNAS) |
| 层级重建 (Layer-wise reconstruction) | 我们不是使用损失函数进行重新训练，而是通过使用类似于知识蒸馏的双模型方法，尝试从每一层恢复尽可能多的信息。 | 参见 https://arxiv.org/pdf/2204.09656.pdf |
表 4.3: 一些常见修剪策略描述。
稀疏模式¶
稀疏模式描述了修剪后的参数在模型/张量中的排列方式。
回想一下，通常需要使用优化的稀疏核才能实现性能提升。根据权重张量的格式和稀疏度水平，稀疏矩阵乘法可能比其对应的密集乘法更快。如果张量不够稀疏，也可能更慢。
在最一般的层面上，修剪是非结构化的——每个参数都有自己的掩码。这提供了最大的灵活性，但需要非常高的稀疏度（>98%）才能提供性能优势。为了在较低稀疏度水平下提供加速推理，硬件后端增加了对特殊稀疏模式的支持。
我们寻求修剪模型，使得权重张量呈现与我们的推理后端相同的稀疏模式。如果在保持稀疏模式的同时能够恢复损失的准确率，我们就可以在稀疏硬件上运行此模型进行加速推理，而不会牺牲准确率。我们也可以在目标后端上运行修剪到不同稀疏模式的模型，但这会带来一些额外的准确率损失。
特定的后端硬件及其对应的稀疏模式，以及修剪配置，最终决定了我们观察到的性能加速效果。如果我们使用不同的修剪标准修剪模型，只要它遵循相同的稀疏模式和稀疏度水平，其性能特征将是相同的。例如，如果我们决定移除最高幅值的权重而不是最低幅值的权重，我们预计这不会改变修剪模型的性能特征。
| 稀疏模式 | 掩码可视化
（50% 稀疏度水平） |
||||||||||||||||||||||||||||||||
| 非结构化稀疏性 |
|
||||||||||||||||||||||||||||||||
| 2:4 半结构化 |
|
||||||||||||||||||||||||||||||||
| 块稀疏性 |
|
||||||||||||||||||||||||||||||||
| 结构化稀疏性 |
|
表 4.4: 一些常见稀疏模式描述。
有关我们支持的 API 和基准测试的更多信息，请参阅稀疏性 README。
