Attention Is All You Need

Goal

Translate a sequence of tokens into another sequence using only attention, without any recurrence or convolution.

Transformer Structure

Encoder

• Stack of identical layers (usually 6)
• Each layer has:
  1. Multi-Head Self-Attention
  2. Position-wise Feed-Forward Network

Decoder

• Stack of identical layers
• Each layer has:
  1. Masked Self-Attention (to block future tokens)
  2. Cross-Attention over encoder outputs
  3. Feed-Forward Network

Each sublayer is wrapped with residual connections + layer normalization.

Scaled Dot-Product Attention

Given queries ( Q ), keys ( K ), and values ( V ):

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

• $d_k$ is the key dimension
• The scaling by $\sqrt{d_k}$ stabilizes gradients

Multi-Head Attention

Instead of a single attention, we use multiple “heads”:

$$\text{head}\_i = \mathrm{Attention}(QW_i^Q,\ KW_i^K,\ VW_i^V)$$

$$\mathrm{MHA}(Q,K,V) = \mathrm{Concat}(\text{head}\_1,\dots,\text{head}\_h)W^O$$

Each head can focus on different relations in the sentence.

Positional Encoding

Since there’s no recurrence, position info is added using sine and cosine functions:

$$\mathrm{PE}(\text{pos},2i) = \sin\left(\frac{\text{pos}}{10000^{2i/d*{\text{model}}}}\right)$$

$$\mathrm{PE}(\text{pos},2i+1) = \cos\left(\frac{\text{pos}}{10000^{2i/d*{\text{model}}}}\right)$$

This is added to token embeddings at the input.

Feed-Forward Network

Each position passes through the same small neural network:

$$\mathrm{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2$$

Masking in Decoder

To prevent seeing future tokens, apply a causal mask ( M ):

$$\mathrm{softmax}\left(\frac{QK^\top}{\sqrt{d_k}} + M\right)$$

where $M_{ij} = -\infty$ if $j > i$, otherwise 0.

Training

• Teacher forcing during training
• Cross-entropy loss for next-token prediction
• Label smoothing for better generalization
• Adam optimizer with warmup and inverse-sqrt learning rate schedule

Why It Works Well

• Fully parallelizable, unlike RNNs
• Handles long-range dependencies directly
• Multiple heads learn different linguistic patterns

Complexity

For sequence length ( n ) and model size ( d ):

• Time: $O(n^2 d)$
• Memory: $O(n^2)$ This comes from the attention matrix of size ( n \times n ).

Quick Algorithm Summary

1. Embed tokens and add positional encodings.
2. Encoder: repeat [MHA + AddNorm, FFN + AddNorm] ( L ) times.
3. Decoder: repeat [Masked MHA, Cross MHA, FFN] with AddNorm each.
4. Final linear + softmax for output probabilities.

Key Formulas

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

$$\text{Positional Encoding: sin/cos with geometric wavelengths}$$