OpenAI Sparse Transformer

OpenAI has developed the Sparse Transformer, a deep neural network designed to predict the next element in sequences of text, images, and sound. By implementing an algorithmic improvement to the attention mechanism, the model can extract patterns from sequences 30x longer than previously possible, enabling the modeling of tens of thousands of elements using hundreds of layers.

Reducing Memory Constraints with Deep Attention

Standard Transformer architectures require an $N \times N$ attention matrix for every layer and attention head. This creates a significant memory bottleneck when processing high-dimensional data such as raw audio or large images. For example, a 64x64x3 pixel ImageNet image requires 154 GB of memory for stored matrices in a deep Transformer (64 layers, 4 heads), far exceeding the 12-32 GB capacity of standard GPUs.

To address this, OpenAI implemented two primary optimizations:

  1. Checkpointing: By recomputing the attention matrix from checkpoints during backpropagation, the largest memory cost becomes independent of the number of layers. This allows for the training of networks with substantially greater depth; OpenAI found that models with up to 128 layers outperformed shallower networks on benchmarks like CIFAR-10.
  2. Operational Adjustments: The team modified the initialization scheme and the ordering of operations within the transformer to support increased depth.

Sparse Attention and Algorithmic Complexity

To handle very large inputs where even a single attention matrix is impractical, the Sparse Transformer utilizes sparse attention patterns. In this approach, each output position computes weightings from only a subset of input positions (e.g., $\sqrt{N}$ elements instead of $N$). This reduces the algorithmic complexity from $O(N^2)$ to $O(N \sqrt{N})$.

Factorized Attention Patterns

To ensure the model retains the ability to learn dynamic, global patterns, OpenAI implemented a two-dimensional factorization of the attention matrix. This allows the network to attend to all positions through two steps of sparse attention:

  • Strided Attention: Each position attends to its own row and column, effectively factorizing the full attention operation. This is particularly useful for data with a two-dimensional structure, such as images.
  • Fixed Attention: The network attends to a fixed column and the elements following the latest column element. This pattern is optimized for data that does not fit a two-dimensional structure, such as text.

Experimental Results and Performance

The Sparse Transformer achieved new state-of-the-art scores for density estimation on the following benchmark datasets:

  • CIFAR-10
  • Enwik8
  • Imagenet 64

OpenAI observed that sparse attention not only performed faster than full attention but also achieved lower loss. The researchers suggest this may be due to an underlying optimization issue with dense attention or a beneficial inductive bias provided by the sparsity patterns.

Generative Capabilities across Modalities

Image Generation

The model demonstrates a grasp of global structure through image completion tasks on 64x64 ImageNet data. Because the models are trained using a maximum likelihood objective, unconditional samples generated with an unadjusted softmax temperature of 1.0 reflect the full distribution of images the model believes exist, which can occasionally result in strange-looking samples.

Raw Audio Generation

By modifying the position embeddings, the Sparse Transformer can generate raw audio waveforms. OpenAI trained the model on raw classical music clips, enabling it to generate sequences of 65,000 elements, which corresponds to approximately 5 seconds of raw audio.

Implementation and Limitations

To facilitate the use of sparse attention, OpenAI open-sourced a set of block-sparse GPU kernels that efficiently perform the necessary slicing of query and key matrices.

Despite these advancements, the researchers noted that autoregressive sequence generation remains impractical for very high-resolution images or video. They suggest that these optimized attention operations could serve as primitives for multi-scale modeling approaches in the future.

Sources