Generative modeling with sparse transformers

OpenAI

We’ve developed the Sparse Transformer, a deep neural network which sets new records at predicting what comes next in a sequence—whether text, images, or sound. It uses an algorithmic improvement of the attention mechanism to extract patterns from sequences 30x longer than possible previously.

One existing challenge in AI research is modeling long-range, subtle interdependencies in complex data like images, videos, or sounds. The Sparse Transformer incorporates an O(NN) O(N \sqrt{N}) reformulation of the O(N2) O(N^2) Transformer⁠(opens in a new window) self-attention mechanism, along with several other improvements, to apply it directly to these rich data types. Previously, models used on these data were specifically crafted for one domain or difficult to scale to sequences more than a few thousand elements long. In contrast, our model can model sequences with tens of thousands of elements using hundreds of layers, achieving state-of-the-art performance across multiple domains. At OpenAI, we’re using it to help us build AI systems that possess a greater ability to understand the world.

Deep attention

In Transformers, every output element is connected to every input element, and the weightings between them are dynamically calculated based upon the circumstances, a process called attention. While it is believed that this allows Transformers to be more flexible than models with fixed connectivity patterns, in practice it requires the creation of an N×N N\times N  attention matrix for every layer and attention head, which can consume large amounts of memory when applied to data types with many elements, like images or raw audio.

Data type	Stored	Recomputed
1024 text tokens (several paragraphs)	1.0 GB	16 MB
32x32x3 pixels (CIFAR-10 image)	9.6 GB	151 MB
64x64x3 pixels (Imagenet 64 image)	154 GB	2.4 GB
24,000 samples (~2 seconds of 12 kHz audio)	590 GB	9.2GB

Sparse attention

Experimental results

Generating images

Model samples

Real data