Attention Is All You Need — Explaining the Transformer

Why This Matters

Before 2017, sequence-to-sequence tasks (like machine translation) were dominated by Recurrent Neural Networks (RNNs), LSTMs, and Gated Recurrent Units (GRUs). These networks process tokens sequentially, meaning to compute the representation of the 10th word in a sentence, you must first compute the representations of the first 9 words.

This sequential bottleneck creates two massive problems:

Inability to Parallelize: Modern hardware (GPUs/TPUs) thrives on parallel execution. Sequential processing means training cannot be easily accelerated across long sequences.
Vanishing/Exploding Gradients: Signals struggle to travel across long temporal distances, making it extremely difficult for the model to learn dependencies between distant words.

The Big Idea

"We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely."

By replacing recurrence with self-attention, the Transformer connects all positions in a sequence with a constant number of sequentially executed operations ($O(1)$ path length). This allows the network to look at the entire sentence at once, capture long-range dependencies instantly, and train orders of magnitude faster.

How It Works

1. Scaled Dot-Product Attention

The core mathematical engine is Scaled Dot-Product Attention. The input consists of Queries ($Q$), Keys ($K$), and Values ($V$).

$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

Why scale by $\sqrt{d_k}$? As the dimensionality $d_k$ grows large, the dot products grow large in magnitude, pushing the softmax function into regions with extremely small gradients. Dividing by $\sqrt{d_k}$ counteracts this effect.

Interactive Scaling & Vanishing Gradients

Widget 1

Simulate how large key dimensions ($d_k$) push softmax values to extremes, and how the paper's scaling factor ($\frac{1}{\sqrt{d_k}}$) restores healthy gradients.

Dimension size ($d_k$): 64

Raw Dot Product Scale: 2.0

Apply Scaling ($\frac{1}{\sqrt{d_k}}$)

Softmax Distribution (Attention Weights)

Token 1 (High Match)0.0%

Token 2 (Medium Match)0.0%

Token 3 (Low Match)0.0%

2. Multi-Head Attention

Instead of performing a single attention function, the authors project the queries, keys, and values $h$ times with different, learned linear projections. This allows the model to jointly attend to information from different representation subspaces at different positions.

Multi-Head Subspace Visualizer

Widget 2

Select different attention heads to see how they capture different syntactic and semantic relationships for the word "it" in the sentence:

"The animal didn't cross the street because it was too tired."

Note: This mimics real layer 5 attention maps from the paper where different heads resolve anaphora.

3. Positional Encoding

Since the Transformer contains no recurrence or convolution, it has no inherent sense of word order. To fix this, the authors add positional encodings to the input embeddings:

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

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

Positional Encoding Wave Generator

Widget 3

Hover over the 2D Positional Encoding matrix below to see how different dimensions represent positions with varying sinusoidal frequencies.

Dimension Index ($i$) →

Sequence Position ($pos$) →

Hovered Position

Pos: 0

Dim: 0

Value: 0.0000

Lower dimensions (left) oscillate rapidly; higher dimensions (right) oscillate slowly. This allows the model to learn relative positions easily.

Results & Efficiency

On the WMT 2014 English-to-German translation task, the Transformer (big) achieved a state-of-the-art 28.4 BLEU, outperforming existing ensembles by over 2.0 BLEU.

Crucially, it achieved this while requiring significantly less training compute (FLOPs) than convolutional or recurrent models.

WMT'14 English-to-German Results vs. Compute Cost

Notice how both Transformer models achieve higher BLEU scores with orders of magnitude less training compute (FLOPs).

Limitations & Open Questions

Despite its revolutionary performance, the original Transformer has a few notable limitations:

Quadratic Complexity $O(n^2)$: Self-attention requires computing similarity between every pair of tokens. For extremely long sequences, this becomes computationally prohibitive.
Autoregressive Bottleneck: During inference, the model generates text token-by-token, which cannot be parallelized in the same way training is.
No Inductive Bias: Unlike CNNs (which assume local spatial patterns) or RNNs (which assume sequential order), Transformers make no assumptions. This means they require massive amounts of data to generalize well.

Glossary

Recurrent Neural Network (RNN)

A class of neural networks where connections between nodes form a directed graph along a temporal sequence. This allows it to exhibit temporal dynamic behavior, but forces sequential processing.

Self-Attention

An attention mechanism relating different positions of a single sequence in order to compute a representation of the sequence.

Positional Encoding

Static or learned vectors added to the input embeddings to inject information about the relative or absolute position of tokens in the sequence.

Citation

@inproceedings{vaswani2017attention,
  author    = {Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, {\L}ukasz and Polosukhin, Illia},
  title     = {Attention is all you need},
  booktitle = {Advances in Neural Information Processing Systems},
  pages     = {5998--6008},
  year      = {2017}
}