The 2017 paper “Attention Is All You Need” (Vaswani et al.) is perhaps the most influential paper in modern AI, as it replaced recurrent (LSTM/GRU) and convolutional layers with the Transformer architecture.
The core principle is Self-Attention, which allows the model to process all words in a sequence simultaneously, enabling massive parallelization and capturing long-range dependencies better than any previous model.
1. Scaled Dot-Product Attention
This is the “engine” of the Transformer. It calculates how much focus each word should put on every other word in the sentence.
$$Attention(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$
Python
import torch
import torch.nn as nn
import torch.nn.functional as F
class ScaledDotProductAttention(nn.Module):
def __init__(self):
super().__init__()
def forward(self, q, k, v, mask=None):
# q, k, v: [batch_size, n_heads, seq_len, d_k]
d_k = q.size(-1)
# Calculate scores
scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k, dtype=torch.float32))
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn = F.softmax(scores, dim=-1)
output = torch.matmul(attn, v)
return output, attn
2. Multi-Head Attention
Instead of one expensive attention pass, the model splits the embedding into multiple “heads” that attend to different types of information (e.g., one head for syntax, another for pronouns) in parallel.
Python
class MultiHeadAttention(nn.Module):
def __init__(self, d_model, n_heads):
super().__init__()
self.n_heads = n_heads
self.d_model = d_model
self.d_k = d_model // n_heads
self.w_q = nn.Linear(d_model, d_model)
self.w_k = nn.Linear(d_model, d_model)
self.w_v = nn.Linear(d_model, d_model)
self.w_concat = nn.Linear(d_model, d_model)
self.attention = ScaledDotProductAttention()
def forward(self, q, k, v, mask=None):
batch_size = q.size(0)
# Linear projections and split into heads
q = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
k = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
v = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_k).transpose(1, 2)
# Scaled Dot-Product Attention
out, attn = self.attention(q, k, v, mask=mask)
# Concatenate heads and final linear layer
out = out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
return self.w_concat(out)
3. Positional Encoding
Because there are no RNNs to track order, the Transformer adds sinusoidal waves to the embeddings to give the model a sense of where each word is in the sentence.
Python
import numpy as np
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, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.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))
def forward(self, x):
# x: [batch_size, seq_len, d_model]
return x + self.pe[:, :x.size(1)]
4. The Transformer Block (Encoder)
Each layer consists of Multi-Head Attention, followed by a Position-wise Feed-Forward Network (FFN), tied together with Residual Connections and Layer Normalization.
Python
class EncoderLayer(nn.Module):
def __init__(self, d_model, n_heads, d_ff, dropout):
super().__init__()
self.self_attn = MultiHeadAttention(d_model, n_heads)
self.feed_forward = 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, mask):
# Sub-layer 1: Attention + Residual + Norm
attn_out = self.self_attn(x, x, x, mask)
x = self.norm1(x + self.dropout(attn_out))
# Sub-layer 2: FFN + Residual + Norm
ff_out = self.feed_forward(x)
x = self.norm2(x + self.dropout(ff_out))
return x
Why it Changed Everything
- Parallelization: Unlike LSTMs, which must wait for word $t-1$ to compute word $t$, the Transformer computes all words at once.
- Constant Path Length: Any word can “talk” to any other word in a single step, regardless of distance, solving the long-term dependency problem perfectly.
- Foundation for LLMs: This architecture is the direct ancestor of GPT, BERT, and Claude.
Since you’ve been exploring the evolution of sequence models, this video provides a detailed walkthrough of the Transformer architecture from scratch. It is highly relevant because it visually maps each mathematical component of the “Attention Is All You Need” paper to its corresponding line of PyTorch code.