Attention Is All You Need Paper Breakdown

A comprehensive engineering teardown of Attention Is All You Need (Vaswani et al., NIPS 2017), the seminal research paper that replaced recurrent neural networks (RNNs) and convolutional architectures with pure self-attention mechanisms.

1. Scaled Dot-Product Attention Formulation

The core innovation of the Transformer is Scaled Dot-Product Attention. Given input query matrix Q, key matrix K, and value matrix V:

Attention(Q, K, V) = softmax((Q * K^T) / sqrt(d_k)) * V

Mathematical Mechanics

1. Query-Key Dot Product: Computes unnormalized similarity scores between every query token and key token, producing an N x M matrix.
2. Scaling Factor: Prevents large projection dimensions from pushing the softmax function into regions with extremely small gradients, preserving gradient flow during backpropagation.
3. Softmax Normalization: Applies row-wise softmax to convert similarity scores into a valid probability distribution over keys.
4. Value Projection: Multiplies normalized attention weights by values to compute weighted context representations.

2. Multi-Head Attention (MHA) Architecture

Instead of performing a single attention function with d_model-dimensional queries, keys, and values, Multi-Head Attention projects Q, K, V into h lower-dimensional subspaces:

head_i = Attention(Q * W_i^Q, K * W_i^K, V * W_i^V)

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

Where projection matrices have dimensions:

• W_i^Q in R^(d_model x d_k)
• W_i^K in R^(d_model x d_k)
• W_i^V in R^(d_model x d_v)
• W^O in R^(h*d_v x d_model)

Multi-Head Attention allows the model to jointly attend to information from different representation subspaces at different positions.

3. Sinusoidal Positional Encoding

Because the Transformer contains no recurrence or convolution, it has no implicit sequence order awareness. To inject positional information, the authors add fixed sinusoidal encodings to input embeddings:

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

Where pos is the token position in the sequence and i is the dimension index. This formulation allows the model to easily learn relative position attention offsets since PE_(pos+k) can be represented as a linear function of PE_pos.

4. Engineering Tradeoffs & KV Cache Memory Footprint

While self-attention allows total parallelization during training, it introduces key engineering trade-offs for production serving:

| Dimension | Recurrent (RNN / LSTM) | Transformer (Self-Attention) | |---|---|---| | Training Time Complexity | O(N) sequential operations | O(1) parallel matrix multiplications | | Context Length Scaling | O(1) fixed hidden state | O(N^2) quadratic attention matrix | | Inference KV Cache VRAM | Zero KV cache (fixed state) | 2 * b * s * l * h * d bytes |

The quadratic O(N^2) compute complexity and linear KV cache memory expansion motivated subsequent optimizations like FlashAttention, PagedAttention, and Grouped-Query Attention (GQA).