# The Annotated Transformer

> Line-by-line PyTorch implementation of the Transformer architecture

- Type: Blog
- By: Sasha Rush
- Year: 2018
- Topics: Deep learning, Transformers, Attention, NLP
- Page: https://aipedia.org/annotated-transformer

## Resources

- [The Annotated Transformer](https://nlp.seas.harvard.edu/annotated-transformer/): Article, Harvard NLP, 2018
- [harvardnlp/annotated-transformer](https://github.com/harvardnlp/annotated-transformer): Code, Notebook source
- [Attention Is All You Need](https://arxiv.org/abs/1706.03762): Paper, Vaswani et al., 2017

**The Annotated Transformer** is a line-by-line guide to implementing the Transformer architecture in PyTorch. Created by Harvard NLP, it makes the seminal "Attention Is All You Need" paper concrete and reproducible: every equation in the paper is paired with the few lines of code that compute it.

> This page is about a **tutorial implementation**, not a new model. If you have not met the architecture yet, read [Transformer](https://aipedia.org/transformer) first; for the paper itself, see [Attention Is All You Need](https://aipedia.org/attention-is-all-you-need).

## Why Students Should Care

- Reading a paper and *implementing* it are very different skills — this resource bridges the two.
- The full working Transformer fits in roughly 400 lines of PyTorch. Seeing that demystifies the architecture behind modern LLMs.
- The habits it teaches (map each equation to a small module, test pieces independently) transfer to implementing any paper.

## Why Code Matters

The original Transformer paper describes the architecture mathematically. The Annotated Transformer shows exactly how those equations become working code. For example, the attention equation

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

becomes just three lines:

```python
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
p_attn = scores.softmax(dim=-1)
return torch.matmul(p_attn, value)
```

If the math felt abstract, the code makes it concrete: a matrix multiply to compare queries with keys, a softmax to turn scores into weights, and another matrix multiply to mix the values.

## Core Components

The full Transformer in ~400 lines breaks down into six key pieces:

1. **Embeddings + Positional Encoding** — token lookup + position information
2. **Multi-Head Attention** — parallel attention heads
3. **Feed-Forward Network** — position-wise MLP
4. **Encoder Layer** — self-attention + FFN with residuals
5. **Decoder Layer** — masked self-attention + cross-attention + FFN
6. **Generator** — project to vocabulary for prediction

Each piece is a small, independent module — none of them is individually complicated.

## Interactive Code Explorer

Click through the architecture to see the implementation of each component:

*Interactive demo: https://aipedia.org/annotated-transformer*

## Key Implementation Details

### Scaled Dot-Product Attention

The complete attention function, including the mask used by the decoder to hide future tokens:

```python
def attention(query, key, value, mask=None, dropout=None):
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    p_attn = scores.softmax(dim=-1)
    return torch.matmul(p_attn, value), p_attn
```

Note the trick: masked positions get a score of `-1e9`, so after the softmax their weight is effectively zero.

### Sublayer Connection (Residual + LayerNorm)

Every attention and feed-forward block is wrapped in the same pattern — normalize, apply the sublayer, add the result back to the input:

```python
class SublayerConnection(nn.Module):
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))
```

The residual `x +` is what lets gradients flow through many stacked layers.

## Learning Path

The Annotated Transformer is best approached in order:

1. **Embeddings** — how tokens become vectors
2. **Attention** — the core mechanism
3. **Multi-Head** — parallel attention
4. **Encoder/Decoder** — full architecture
5. **Training** — label smoothing, optimizer

## Common Confusion

- **The Annotated Transformer is a tutorial, not a model.** It introduces no new ideas — it faithfully implements the 2017 paper.
- **Transformer vs. Attention Is All You Need vs. this page.** [Transformer](https://aipedia.org/transformer) is the architecture, [Attention Is All You Need](https://aipedia.org/attention-is-all-you-need) is the paper, and the Annotated Transformer is the walkthrough that turns the paper into code.
- **~400 lines does not include training infrastructure.** Data loading, batching, and large-scale training are separate engineering concerns.

## Where To Go Next

- Read [Transformer](https://aipedia.org/transformer) for the conceptual overview of the architecture
- Read [Attention Is All You Need](https://aipedia.org/attention-is-all-you-need) for the original paper and results
- Read [Layer Normalization](https://aipedia.org/layer-normalization) for the normalization used in every sublayer
- Read [BERT](https://aipedia.org/bert) and [GPT](https://aipedia.org/gpt) for what people built once the architecture was implementable