Attention Is All You Need — the transformer

Attention Is All You Need — the transformer

Vaswani et al. (2017) removed recurrence entirely. The transformer processes all positions in a sequence simultaneously, using self-attention as the sole mechanism for mixing information across positions. No hidden state passed from one token to the next. No sequential dependency during training. The paper's title — "Attention Is All You Need" — was a literal architectural claim, not a slogan.

The encoder-decoder architecture

The original transformer is an encoder-decoder model designed for machine translation. The encoder reads a source sentence and produces a sequence of representations; the decoder generates the target sentence one token at a time, attending to both its own prior outputs and the encoder's representations.

The encoder consists of 6 identical blocks stacked sequentially. Each block contains two sub-layers:

  • Multi-head self-attention: every position attends to every other position in the input sequence
  • Position-wise feed-forward network: a two-layer MLP applied independently to each position

Each sub-layer is wrapped in a residual connection followed by layer normalization: output = LayerNorm(x + sublayer(x)).

The decoder also has 6 blocks, but each block contains three sub-layers:

  • Masked multi-head self-attention: each position attends only to earlier positions (preventing the model from seeing future tokens during training)
  • Cross-attention: the decoder attends to the encoder's output representations
  • Position-wise feed-forward network

The same residual-connection-then-layer-norm pattern wraps every sub-layer.

Model dimensions and parameter count

The base transformer uses d_model=512 — every token is represented as a 512-dimensional vector throughout the network. Attention uses 8 heads, each with d_k = d_v = 512 / 8 = 64 dimensions. The feed-forward network's inner dimension is d_ff=2048 — 4x the model dimension, a ratio that became standard across nearly all later transformers.

Parameters per encoder block:

  • Self-attention: 4 projection matrices (Q, K, V, output) of size 512 × 512 each = 4 × 262,144 = 1,048,576
  • Feed-forward network: 512 × 2048 + 2048 × 512 = 2,097,152
  • Layer norms: 2 × 2 × 512 = 2,048
  • Total per block: approximately 3.15M

Six encoder blocks: ~18.9M parameters. The decoder blocks are slightly larger due to cross-attention (an additional ~1M per block), totaling ~24.8M. Adding embeddings (a shared source/target vocabulary of ~37,000 BPE tokens × 512 dims = ~19M) and the output projection brings the base model to approximately 65M parameters.

The big transformer scales d_model to 1024 with 16 heads and d_ff=4096, reaching ~213M parameters.

Positional encoding

Without recurrence, the model has no inherent notion of token order — it sees a set, not a sequence. The original transformer injects position information by adding sinusoidal vectors to the token embeddings before the first encoder block:

PE(pos, 2i) = sin(pos / 10000^(2i / d_model))

PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model))

Each dimension i of the positional encoding oscillates at a different frequency, from wavelength (dimension 0) to wavelength 10000 × 2π (dimension d_model-1). This gives the model a unique positional fingerprint for each position and allows it to learn relative distances — the dot product between two positional encodings depends only on their offset, not their absolute positions.

Later models replaced sinusoidal encodings with learned positional embeddings (GPT-2, BERT) or rotary position embeddings (RoPE, used in Llama and most modern models). The sinusoidal encoding is rarely used in practice today, but it introduced the concept that position must be explicitly encoded when recurrence is removed.

Training setup and results

The original transformer was trained on the WMT 2014 English-German dataset (~4.5M sentence pairs) and English-French dataset (~36M pairs). The base model trained for 100,000 steps on 8 NVIDIA P100 GPUs in approximately 12 hours. The big model trained for 300,000 steps on the same hardware in 3.5 days.

Optimization used Adam with a custom learning rate schedule — a linear warmup over 4,000 steps followed by inverse-square-root decay: lr = d_model^(-0.5) * min(step^(-0.5), step * warmup_steps^(-1.5)). This schedule became the default for transformer training for several years.

Results on WMT 2014 English-German: the big transformer achieved BLEU 28.4, surpassing all prior single models and ensemble systems. On English-French: BLEU 41.0, a new state of the art, at less than 1/4 the training cost of the previous best system. The combination of better quality and dramatically lower training cost made the result impossible to ignore.

Why this architecture changed everything

Three properties make the transformer fundamentally different from RNN-based architectures:

Parallel computation. In an RNN, token t cannot be processed until token t-1 finishes. In a transformer, all tokens are processed simultaneously — self-attention computes all pairwise interactions in a single batched matrix multiplication. A 512-token sequence that requires 512 sequential steps in an LSTM requires exactly one attention operation in a transformer, fully utilizing GPU parallelism.

Direct connections. In an RNN, information from token 1 must pass through every intermediate hidden state to reach token 500. In a transformer, token 1 and token 500 are connected directly through attention — the path length is O(1), not O(n). This eliminates the vanishing gradient problem for long-range dependencies entirely.

Scalability. The transformer's parallel computation makes it possible to scale to billions of parameters on clusters of GPUs or TPUs. Every major language model since 2018 — GPT-2 (Radford et al. 2019), BERT (Devlin et al. 2019), GPT-3 (Brown et al. 2020), PaLM (Chowdhery et al. 2022), Llama (Touvron et al. 2023) — is a transformer. As of 2026, the paper has approximately 130,000 citations.

The architectural skeleton in PyTorch

The following implements the encoder-decoder structure of the original transformer using PyTorch's nn.MultiheadAttention. Attention itself is not implemented from scratch here — see the Attention course for the full derivation and implementation of scaled dot-product and multi-head attention.

import torch
import torch.nn as nn
import math

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super().__init__()
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(
            torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
        )
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe.unsqueeze(0))  # (1, max_len, d_model)

    def forward(self, x):
        return x + self.pe[:, :x.size(1)]


class EncoderBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads,
                                               dropout=dropout, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_mask=None):
        attn_out, _ = self.self_attn(x, x, x, attn_mask=src_mask)
        x = self.norm1(x + self.dropout(attn_out))
        x = self.norm2(x + self.dropout(self.ffn(x)))
        return x


class DecoderBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(d_model, n_heads,
                                               dropout=dropout, batch_first=True)
        self.cross_attn = nn.MultiheadAttention(d_model, n_heads,
                                                dropout=dropout, batch_first=True)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, d_ff),
            nn.ReLU(),
            nn.Linear(d_ff, d_model),
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, enc_out, tgt_mask=None):
        attn_out, _ = self.self_attn(x, x, x, attn_mask=tgt_mask)
        x = self.norm1(x + self.dropout(attn_out))
        cross_out, _ = self.cross_attn(x, enc_out, enc_out)
        x = self.norm2(x + self.dropout(cross_out))
        x = self.norm3(x + self.dropout(self.ffn(x)))
        return x


class Transformer(nn.Module):
    def __init__(self, src_vocab, tgt_vocab, d_model=512,
                 n_heads=8, d_ff=2048, n_layers=6, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.src_embed = nn.Embedding(src_vocab, d_model)
        self.tgt_embed = nn.Embedding(tgt_vocab, d_model)
        self.pos_enc = PositionalEncoding(d_model)

        self.encoder = nn.ModuleList(
            [EncoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)]
        )
        self.decoder = nn.ModuleList(
            [DecoderBlock(d_model, n_heads, d_ff, dropout) for _ in range(n_layers)]
        )
        self.output_proj = nn.Linear(d_model, tgt_vocab, bias=False)

    def encode(self, src, src_mask=None):
        x = self.pos_enc(self.src_embed(src) * math.sqrt(self.d_model))
        for layer in self.encoder:
            x = layer(x, src_mask)
        return x

    def decode(self, tgt, enc_out, tgt_mask=None):
        x = self.pos_enc(self.tgt_embed(tgt) * math.sqrt(self.d_model))
        for layer in self.decoder:
            x = layer(x, enc_out, tgt_mask)
        return x

    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        enc_out = self.encode(src, src_mask)
        dec_out = self.decode(tgt, enc_out, tgt_mask)
        return self.output_proj(dec_out)


model = Transformer(src_vocab=37000, tgt_vocab=37000)
total_params = sum(p.numel() for p in model.parameters())
print(f"Total parameters: {total_params:,}")  # ~65M

src = torch.randint(0, 37000, (2, 30))  # batch=2, src_len=30
tgt = torch.randint(0, 37000, (2, 25))  # batch=2, tgt_len=25

# Causal mask for decoder self-attention
tgt_mask = nn.Transformer.generate_square_subsequent_mask(25)
logits = model(src, tgt, tgt_mask=tgt_mask)
print(f"Output shape: {logits.shape}")  # (2, 25, 37000)

The generate_square_subsequent_mask produces the causal mask that prevents each decoder position from attending to future positions — this is what makes autoregressive generation possible during training with teacher forcing.

What came next

The original transformer was an encoder-decoder for machine translation. Within two years, three variants emerged that defined the next era of NLP:

Encoder-only (BERT, Devlin et al. 2019): Use only the encoder stack. Train with masked language modeling — mask 15% of tokens, predict them from context. BERT-base: 12 layers, d_model=768, 110M parameters. Dominated NLU benchmarks (GLUE, SQuAD) and became the standard for classification, NER, and extractive QA.

Decoder-only (GPT, Radford et al. 2018): Use only the decoder stack (without cross-attention). Train with autoregressive language modeling — predict the next token given all previous tokens. GPT-1: 12 layers, d_model=768, 117M parameters. GPT-2 (Radford et al. 2019): 48 layers, d_model=1600, 1.5B parameters. GPT-3 (Brown et al. 2020): 96 layers, d_model=12288, 175B parameters. This lineage — decoder-only transformers trained as language models — is the architecture behind every major LLM today.

Encoder-decoder (T5, Raffel et al. 2020): Keep the original structure but frame every NLP task as text-to-text. T5-11B: 11B parameters. Still used in some production systems but the decoder-only paradigm dominates.

The encoder-decoder architecture from Vaswani et al. (2017) was a starting point. The decoder-only variant — simpler, easier to scale, and capable of zero-shot generalization via in-context learning — is what the field converged on.

← Previous