Attention Is All You Need

The definitive deep-dive into the 2017 paper by Vaswani et al. that introduced the Transformer architecture.

1. The Problem with Prior Architectures

Before the Transformer, the dominant sequence transduction models for Natural Language Processing (NLP) tasks like machine translation were based on complex recurrent or convolutional neural networks (RNNs, LSTMs, GRUs) often arranged in an encoder-decoder configuration.

The Transformer's Solution: The Transformer eschews recurrence and convolutions entirely, relying solely on an attention mechanism to draw global dependencies between input and output. This allows for significantly more parallelization and reduces the number of sequential operations to a constant $O(1)$.

2. Model Architecture overview

The Transformer follows the standard encoder-decoder architecture, but replaces the recurrent layers with stacked self-attention and point-wise, fully connected layers.

The Encoder

The encoder is composed of a stack of $N = 6$ identical layers. Each layer has two sub-layers:

  1. Multi-Head Self-Attention Mechanism.
  2. Position-wise Fully Connected Feed-Forward Network.

A residual connection is employed around each of the two sub-layers, followed by Layer Normalization. That is, the output of each sub-layer is $LayerNorm(x + Sublayer(x))$, where $Sublayer(x)$ is the function implemented by the sub-layer itself. To facilitate these residual connections, all sub-layers in the model, as well as the embedding layers, produce outputs of dimension $d_{model} = 512$.

The Decoder

The decoder is also composed of a stack of $N = 6$ identical layers. In addition to the two sub-layers in each encoder layer, the decoder inserts a third sub-layer:

  1. Masked Multi-Head Self-Attention Mechanism: The self-attention sub-layer is modified to prevent positions from attending to subsequent positions (looking into the future). This masking, combined with the fact that output embeddings are offset by one position, ensures that the predictions for position $i$ can depend only on the known outputs at positions less than $i$.
  2. Multi-Head Attention over Encoder Output: This layer performs attention over the output of the encoder stack (queries come from previous decoder layer, keys and values come from the encoder output).
  3. Position-wise Feed-Forward Network.

Like the encoder, residual connections and Layer Normalization are applied around each sub-layer.

3. Deep Dive into Attention

An attention function maps a Query (Q) and a set of Key (K) - Value (V) pairs to an output. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

3.1 Scaled Dot-Product Attention

The specific attention mechanism used is "Scaled Dot-Product Attention."

Formula: $Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V$

Step-by-step computation:

  1. Compute the dot products of the query with all keys: $QK^T$. This results in a score that indicates how much focus to place on other parts of the sequence.
  2. Divide each result by $\sqrt{d_k}$.
  3. Apply a softmax function to obtain the weights (probabilities summing to 1) on the values.
  4. Multiply the softmax weights by the Value matrix $V$.

Why scale by $\frac{1}{\sqrt{d_k}}$?
For large values of $d_k$ (the dimension of keys/queries), the dot products grow very large in magnitude. This pushes the softmax function into regions where it has extremely small (vanishing) gradients, severely slowing down learning. Scaling by the inverse square root of $d_k$ pulls the dot products back to a variance of $1$, keeping the gradients stable.

3.2 Multi-Head Attention

Instead of performing a single attention function with $d_{model}$-dimensional keys, values, and queries, the authors found it beneficial to linearly project the queries, keys, and values $h$ times with different, learned linear projections to $d_k$, $d_k$, and $d_v$ dimensions, respectively.

Attention is applied in parallel to each of these $h$ projected versions, yielding $d_v$-dimensional output values. These are concatenated and once again projected, resulting in the final values.

Formula:

$MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O$

where $head_i = Attention(QW_i^Q, KW_i^K, VW_i^V)$

Dimensions used in the paper: $h = 8$ parallel attention layers (heads). For each of these, $d_k = d_v = d_{model} / h = 512 / 8 = 64$. Because the dimension of each head is reduced, the total computational cost is similar to that of single-head attention with full dimensionality.

Why Multi-Head? It allows the model to jointly attend to information from different representation subspaces at different positions. A single attention head would average this information out.

3.3 How Attention is Used in the Model

4. Position-wise Feed-Forward Networks (FFN)

In addition to attention sub-layers, each of the layers in the encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically.

Formula: $FFN(x) = max(0, xW_1 + b_1)W_2 + b_2$

This consists of two linear transformations with a ReLU activation in between. While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way to describe this is as two convolutions with kernel size 1.

Dimensions: The input and output dimensionality is $d_{model} = 512$, and the inner layer has dimensionality $d_{ff} = 2048$.

5. Embeddings and Softmax

Similarly to other sequence transduction models, learned embeddings are used to convert the input tokens and output tokens to vectors of dimension $d_{model}$. A learned linear transformation and softmax function are used to convert the decoder output to predicted next-token probabilities.

The Transformer shares the same weight matrix between the two embedding layers (input and output) and the pre-softmax linear transformation. In the embedding layers, those weights are multiplied by $\sqrt{d_{model}}$.

6. Positional Encoding

Since the model contains absolutely no recurrence and no convolution, it has no inherent sense of sequence order. If the words in a sentence were shuffled, a purely attention-based model would process them identically.

To inject information about the relative or absolute position of the tokens, "positional encodings" are added to the input embeddings at the bottoms of the encoder and decoder stacks. The positional encodings have the same dimension $d_{model}$ as the embeddings, so they can be summed.

The paper uses sine and cosine functions of different frequencies:

$PE_{(pos, 2i)} = sin(pos / 10000^{2i/d_{model}})$

$PE_{(pos, 2i+1)} = cos(pos / 10000^{2i/d_{model}})$

Where $pos$ is the position and $i$ is the dimension. Each dimension of the positional encoding corresponds to a sinusoid. The wavelengths form a geometric progression from $2\pi$ to $10000 \cdot 2\pi$.

Why this function? The authors hypothesized it would allow the model to easily learn to attend by relative positions, since for any fixed offset $k$, $PE_{pos+k}$ can be represented as a linear function of $PE_{pos}$. Furthermore, it allows the model to extrapolate to sequence lengths longer than the ones encountered during training.

7. Why Self-Attention? (Complexity Comparison)

The paper evaluates self-attention against recurrent and convolutional layers based on three desiderata:

  1. Total computational complexity per layer.
  2. Amount of computation that can be parallelized (measured by the minimum number of sequential operations required).
  3. Path length between long-range dependencies in the network.
Layer Type Complexity per Layer Sequential Operations Maximum Path Length
Self-Attention $O(n^2 \cdot d)$ $O(1)$ $O(1)$
Recurrent $O(n \cdot d^2)$ $O(n)$ $O(n)$
Convolutional $O(k \cdot n \cdot d^2)$ $O(1)$ $O(log_k(n))$

(where $n$ is sequence length, $d$ is representation dimension, $k$ is kernel size)

When sequence length $n$ is smaller than the representation dimensionality $d$ (which is typical in state-of-the-art models like Word-Piece translations), self-attention is computationally faster than RNNs. Moreover, self-attention requires a constant $O(1)$ sequential operations, making it highly parallelizable, and connects all positions with an $O(1)$ path length, making it exceptionally good at learning long-range dependencies.

8. Training Regimen and Regularization

8.1 Optimizer

The Adam optimizer was used with $\beta_1 = 0.9, \beta_2 = 0.98$, and $\epsilon = 10^{-9}$. They varied the learning rate over the course of training, according to the formula:

$lrate = d_{model}^{-0.5} \cdot \min(step\_num^{-0.5}, step\_num \cdot warmup\_steps^{-1.5})$

This corresponds to increasing the learning rate linearly for the first $warmup\_steps$ (set to 4000), and decreasing it thereafter proportionally to the inverse square root of the step number.

8.2 Regularization

Three types of regularization were employed:

  1. Residual Dropout: Dropout ($P_{drop} = 0.1$) was applied to the output of each sub-layer, before it is added to the sub-layer input and normalized. It was also applied to the sums of the embeddings and the positional encodings in both the encoder and decoder stacks.
  2. Label Smoothing: During training, label smoothing of value $\epsilon_{ls} = 0.1$ was employed. This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.

8.3 Hardware and Schedule

The base model was trained on 8 NVIDIA P100 GPUs for 100,000 steps or 12 hours. The big model ($d_{model} = 1024, h=16, d_{ff}=4096$) was trained for 300,000 steps (3.5 days).

9. Results

On the WMT 2014 English-to-German translation task, the big Transformer model achieved a new state-of-the-art BLEU score of 28.4, outperforming the best previously reported models (including ensembles) by over 2.0 BLEU. The base model surpassed all previously published models and ensembles, at a fraction of the training cost of any of the competitive models.

On the WMT 2014 English-to-French translation task, the big model achieved a BLEU score of 41.0, outperforming all previously published single models, at less than 1/4 the training cost of the previous state-of-the-art model.

Summary: Why it Changed AI Forever

By discarding sequential recurrence in favor of an architecture built entirely out of Multi-Head Self-Attention, "Attention Is All You Need" broke the sequential bottleneck of RNNs. This paradigm shift unlocked unprecedented scalability. Hardware (GPUs/TPUs) could now efficiently train massive models on gargantuan datasets because the computation was highly parallelizable.

This direct scalability gave birth to the era of LLMs (Large Language Models) - including BERT, GPT-2, GPT-3, GPT-4, Llama, and beyond. The Transformer is arguably the most impactful neural network architecture developed in the 21st century.