DeepSpeed 稀疏注意力
基于注意力的深度学习模型，如transformers，在捕捉输入序列中标记之间的关系方面非常有效，即使是在长距离的情况下。因此，它们被用于文本、图像和基于声音的输入，其中序列长度可以达到数千个标记。然而，尽管注意力模块在捕捉长期依赖关系方面非常有效，但在实践中，它们对长序列输入的应用受到注意力计算的计算和内存需求的限制，这些需求随着序列长度n呈二次方增长，O(n^2)。
为了解决这一限制，DeepSpeed 提供了一套稀疏注意力内核——一种关键技术，可以通过块稀疏计算将注意力计算的计算和内存需求减少几个数量级。这套工具不仅缓解了注意力计算的内存瓶颈，还能高效地执行稀疏计算。其 API 允许方便地与任何基于 transformer 的模型集成。除了提供广泛的稀疏结构外，它还具有处理任何用户定义的块稀疏结构的灵活性。更具体地说，稀疏注意力（SA）可以设计为计算附近 token 之间的局部注意力，或通过局部注意力计算的摘要 token 进行全局注意力。此外，SA 还可以允许随机注意力，或局部、全局和随机注意力的任意组合，如下图中蓝色、橙色和绿色块所示。因此，SA 将内存占用减少到 O(wn)，其中 1 < w < n 是一个参数，其值取决于注意力结构。
该库基于PyTorch，并通过Triton平台开发所需的内核；内核不是用CUDA编写的，这为将来支持CPU/OpenCL/Vulkan留下了可能性。该库是DeepSpeed的扩展，可以通过DeepSpeed使用，也可以独立使用。
DeepSpeed稀疏注意力内核处理的块稀疏计算分别在以下图中展示了前向和后向传播。在图中，S代表一个块稀疏矩阵，D代表一个密集矩阵。
要了解更多关于稀疏配置的信息，以及如何使用这个库，请查看我们的教程，它提供了详细的信息。
性能结果
- 支持超过10倍长的序列 在一个预训练实验中，我们在三种设置下运行了BERT模型：密集、带激活检查点的密集以及带激活检查点的稀疏（SA）。与密集模型相比，SA在BERT基础版和大版中分别支持10倍和16倍长的序列。下图显示了在BERT基础版和大版模型中可运行的最长序列长度；实验是在单个NVIDIA V100 GPU-32GB内存上以批量大小为1进行的。
- 计算速度提升至6.3倍 我们继续对不同批次大小和序列长度进行了预训练实验，使用了BERT base/large和Megatron GPT2。在这个实验中，我们让训练持续进行100次迭代，并记录最后30次迭代的平均时间。与密集计算相比，SA减少了总计算量并提高了训练速度：随着序列长度的增加，提升效果更明显，BERT base的速度提升了6.3倍，BERT large提升了5.3倍，GPT2提升了6.1倍。以下图表展示了这些结果。
- 更高的准确性 在稀疏注意力（Sparse Transformer, Longformer, BigBird）相关工作中，已经显示出与全注意力相当或更高的准确性。我们的经验也很好地印证了这一点。除了更低的内存开销和更快的计算速度外，我们还观察到在生产环境中，SA（稀疏注意力）达到了更高的准确性和更快的收敛速度。下图展示了基于BERT的长文档理解（2,048序列长度）生产模型的训练准确性。实验在三种设置下进行：从头开始的密集训练、从头开始的SA训练，以及从使用512序列长度的密集训练检查点继续训练的SA。我们观察到，对于从头开始的预训练，SA比密集训练收敛更快且准确性更高。此外，从预训练检查点继续训练的SA在时间和准确性方面表现更好。
- 与最先进的Longformer的比较
我们将SA与Longformer进行了比较，后者是一种最先进的稀疏结构和实现。在我们的实验中，SA使用了
Fixed稀疏性，两种实现的准确性相当。在系统性能方面，SA在训练和推理方面均优于Longformer：- 1.47倍 在Wikitext103上执行预训练MLM的速度更快 我们按照Longformer提供的notebook进行了实验。在这个实验中，我们使用RoBERTa-base检查点预训练了一个MLM模型。这是在8个V100-SXM2 GPU上完成的。下表显示了结果的详细信息，其中使用DeepSpeed稀疏注意力显示了1.47倍的加速。
| 模型 | 局部窗口大小 | BPC | 训练步数 | 每次迭代时间 | 时间改进 | 准确率改进 |
|---|---|---|---|---|---|---|
| RoBERTa 检查点 | 2.5326 | |||||
| Longformer | 512 | 2.6535 | 0 | 1.47 | 1.01 | |
| 稀疏注意力 | 2.6321 | |||||
| Longformer | 1.6708 | 3k | 1.6280 | 1.01 | ||
| 稀疏注意力 | 1.6613 | 1.1059 | ||||
| Longformer | 64 | 5.7840 | 0 | 1.31 | 1.46 | |
| 稀疏注意力 | 3.9737 | |||||
| Longformer | 2.0466 | 3k | 1.4855 | 1.09 | ||
| 稀疏注意力 | 1.8693 | 1.1372 |
- 3.13倍 在BERT-Base上的执行推理速度更快
通过我们上面描述的长文档理解应用程序，我们还检查了在不同窗口大小下测试BERT模型的推理时间，序列长度为
2,048，批量大小为1。在这个实验中，我们注意到用DeepSpeed稀疏注意力替换Bert注意力而不是Longformer注意力时，速度提升了高达3.13倍。下表显示了完整的结果。
| 本地窗口大小 | 时间改进 |
|---|---|
| 512 | 3.13 |
| 256 | 2.29 |
| 128 | 2.16 |
| 64 | 1.5 |
| 32 | 1.24 |
| 16 | 1.23 |
- 灵活处理任何块稀疏结构
DeepSpeed 稀疏注意力套件并不针对任何特定的稀疏结构，而是使模型科学家能够在高效的系统支持下探索任何块稀疏结构。目前，我们已经添加了流行的稀疏结构，例如：
- Fixed（来自 OpenAI 稀疏 Transformer）
- BigBird（来自 Google）
- BSLongformer（来自 AI2 的 Longformer 的块稀疏实现）
我们还定义了一个具有variable结构的模板（上图），该模板可用于简单地自定义任何块稀疏随机/局部/全局注意力模式。除了此列表外，用户还可以添加任何其他稀疏结构，如tutorial部分所述。