Concepts
Self-attention and the Transformer architecture
Intuition
Imagine reading a long sentence in two ways: word by word, only understanding the next one after finishing the previous one, or looking at the whole sentence at once and letting each word directly consult all the others to figure out which ones matter for its meaning. The first way is faithful to the order, but it is slow and makes information from a distant word pass through many intermediaries. The second connects any pair of words in a single step, and since during training all the words in the sentence are handled at the same time, this work can be split across many processors in parallel.
How it works
Recurrent networks (such as LSTM and GRU) compute the state at position t from the state at the previous position and the input at t, which, according to the authors, prevents parallelization within a single training example. The Transformer dispenses with recurrence and convolutions and uses attention only: it maps a query and a set of key-value pairs to a weighted sum of the values, and the weight comes from the compatibility between the query and each key. In the paper this is scaled dot-product attention, in which the product of the query with all the keys is divided by the square root of the key dimension and passed through a softmax; in self-attention, queries, keys, and values come from the output of the previous layer, so each position attends to all positions. Instead of a single attention, multi-head projects queries, keys, and values h times with learned linear projections, applies attention in parallel, concatenates, and projects again (h = 8 heads of dimension 64 in the base configuration), which lets it attend to different representation subspaces at different positions. Since there is no recurrence or convolution, order enters through sine and cosine positional encodings added to the input embeddings; learned positional embeddings gave nearly identical results. In the decoder, self-attention is masked so that position i depends only on earlier positions. Table 1 compares the layers, with n the sequence length and d the representation dimension: self-attention costs O(n² · d) per layer, with O(1) sequential operations and a maximum path of O(1) between two positions; the recurrent layer costs O(n · d²), with O(n) sequential operations and a path of O(n). That is why self-attention is faster than recurrent when n is smaller than d, but the n² term makes the cost grow quickly on long sequences.
In practice
In the paper, this architecture was evaluated on machine translation: the large model reached 28.4 BLEU on WMT 2014 English-to-German, more than 2 BLEU above the best previous results, and the base model took about 12 hours to train on 8 P100 GPUs. The authors point out that n smaller than d is common in the sentence representations used in translation (word-piece and byte-pair), and for very long inputs they propose, as future work, restricting attention to a neighborhood of size r, with cost O(r · n · d), at the price of a longer maximum path, O(n/r). The paper does not describe later uses of the architecture.
Use it in a design
Practice the cases that rely on this concept in the System Design school.