OpenAI Block-sparse GPU Kernels

OpenAI has released highly-optimized GPU kernels designed for neural networks with block-sparse weights. These kernels allow for the training of models that are significantly wider and deeper than traditional dense networks within the same parameter and computational budget, often running orders of magnitude faster than cuBLAS or cuSPARSE.

Technical Implementation of Block-Sparse Kernels

Block-sparse kernels optimize linear operations by treating sparsity at the level of blocks rather than individual weights. This approach reduces computational and storage costs, as the kernels skip computations for any block that is zero.

Key Specifications

Supported Layers: The kernels are designed for use in fully connected and convolutional layers. In convolutional layers, sparsity is applied to input and output feature dimensions, while spatial connectivity remains unaffected.
Optimized Block Sizes: The implementation is optimized for block sizes of 8x8, 16x16, and 32x32.
Configurability: The sparsity pattern at the block level is completely configurable.
Efficiency: Both the computational cost and the parameter storage cost are proportional only to the number of non-zero weights, not the total number of input/output features.

Performance and Benchmarks

Testing conducted on an NVIDIA Titan X Pascal GPU with CUDA 8 demonstrated significant speed-ups over standard libraries. For a wide state of 12,288 hidden units, a block size of 32x32, and a minibatch size of 32, the kernels outperformed cuBLAS, with even larger speed-ups observed when compared to cuSPARSE.

Applications and Research Results

OpenAI utilized these kernels to explore several architectural patterns, most notably "small-world" neural networks, which mimic the connectivity of biological brains to allow information to propagate quickly despite high sparsity.

Small-World LSTMs

By implementing small-world sparse connectivity, OpenAI trained LSTMs with nearly 20,000 hidden units—five times wider than a dense network with a similar parameter count. This architecture improved results in generative modeling of text and semi-supervised sentiment classification.

Sentiment Representation Learning

In experiments comparing dense weight matrices against block-sparse variants with equivalent parameter counts, the sparse models outperformed dense models across all sentiment datasets. Specifically, on the document-level IMDB dataset, the sparse model reduced the error rate from 5.91% (Miyato et al, 2016) to 5.01%.

Image and Text Compression

Text Modeling: Bits-per-character results dropped from 1.059 to 1.048 for models with approximately 100 million parameters.
Image Modeling: By replacing regular 2D convolutional kernels with sparse kernels in a modified PixelCNN++ model for CIFAR-10, OpenAI deepened the network while keeping other hyperparameters fixed, reducing bits-per-dimension from 2.92 to 2.90.

Future Research Directions

OpenAI identifies three primary areas for further exploration using block-sparse kernels:

Inference Speed-up: Investigating the wall-clock time speed-up possible at inference time when combining these kernels with post-training pruning.
Dynamic Sparsity Structure: Exploring whether gradients can be used to learn the optimal sparsity structure during development, similar to synaptic pruning in biological brains.
Scaling Applications: Identifying specific application domains where the ability to train models with massive weight matrices (maintaining the same parameter and computational cost as smaller dense models) provides the most significant performance gains.