OpenAI 块稀疏 GPU 内核

OpenAI 已发布专为具有块稀疏权重的神经网络设计的高度优化 GPU 内核。这些内核使得在相同参数和计算预算下，可以训练出比传统密集网络宽得多、深得多的模型，并且通常比 cuBLAS 或 cuSPARSE 快几个数量级。

块稀疏内核的技术实现

块稀疏内核通过将稀疏性视为块级别而非单个权重来优化线性操作。这种方法降低了计算和存储成本，因为内核会跳过任何为零的块的计算。

关键规范

支持的层：内核专为全连接层和卷积层设计。在卷积层中，稀疏性应用于输入和输出特征维度，而空间连接保持不变。
优化的块大小：实现针对 8x8、16x16 和 32x32 块大小进行了优化。
可配置性：块级别的稀疏模式完全可配置。
效率：计算成本和参数存储成本仅与非零权重的数量成正比，而与输入/输出特征的总数无关。

性能和基准测试

在 NVIDIA Titan X Pascal GPU 上使用 CUDA 8 进行的测试表明，相较于标准库，这些内核展示了显著的加速。对于 12,288 个隐藏单元的宽状态、块大小为 32x32、小批量大小为 32 的情况，这些内核的性能优于 cuBLAS，并且与 cuSPARSE 相比观察到更大的加速。

应用和研究结果

OpenAI 使用这些内核探索了几种架构模式，最值得注意的是“小世界”神经网络，它们模仿生物大脑的连接方式，以便在高稀疏性的情况下仍能让信息快速传播。

小世界 LSTM

通过实现小世界稀疏连接，OpenAI 训练了具有近 20,000 个隐藏单元的 LSTM——这比参数数量相似的密集网络宽五倍。该架构在文本生成建模和半监督情感分类方面取得了改进。

情感表示学习

在将密集权重矩阵与参数数量相当的块稀疏变体进行比较的实验中，稀疏模型在所有情感数据集上均优于密集模型。具体来说，在文档级 IMDB 数据集上，稀疏模型将错误率从 5.91%（Miyato 等人，2016）降低到 5.01%。

图像和文本压缩

文本建模：对于大约 1 亿参数的模型，每字符比特数从 1.059 下降到 1.048。
图像建模：通过在 CIFAR-10 的修改版 PixelCNN++ 模型中用稀疏内核替换常规的 2D 卷积内核，OpenAI 在保持其他超参数不变的情况下加深了网络，使每维比特数从 2.92 降低到 2.90。

未来研究方向

OpenAI 确定了使用块稀疏内核进行进一步探索的三个主要领域：

推理加速：研究在将这些内核与训练后剪枝结合时，推理阶段可能实现的墙钟时间加速。
动态稀疏结构：探索梯度是否可以在开发过程中用于学习最优的稀疏结构，类似于生物大脑中的突触修剪。
规模应用：确定特定的应用领域，在这些领域中，能够训练具有巨大权重矩阵的模型（保持与较小密集模型相同的参数和计算成本）能够带来最显著的性能提升。