Transformers from Scratch

Introduction

The Transformer architecture, introduced in "Attention Is All You Need" (Vaswani et al., 2017), has become the backbone of modern deep learning. This post walks through the key components from scratch.

Self-Attention Mechanism

The core of the Transformer is the scaled dot-product attention:

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

where:

Multi-Head Attention

Instead of a single attention function, we use multiple heads:

MultiHead(Q,K,V)=Concat(head1,...,headh)WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, ..., \text{head}_h)W^O

where each head is:

headi=Attention(QWiQ,KWiK,VWiV)\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)

This allows the model to attend to information from different representation subspaces.

Positional Encoding

Since attention is permutation invariant, we add positional information:

PE(pos,2i)=sin(pos100002i/dmodel)PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i/d_{model}}}\right)

PE(pos,2i+1)=cos(pos100002i/dmodel)PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i/d_{model}}}\right)

Feed-Forward Network

Each layer contains a position-wise feed-forward network:

FFN(x)=max(0,xW1+b1)W2+b2\text{FFN}(x) = \max(0, xW_1 + b_1)W_2 + b_2

This is applied identically to each position.

Layer Normalization & Residual Connections

Each sub-layer uses a residual connection followed by layer normalization:

output=LayerNorm(x+SubLayer(x))\text{output} = \text{LayerNorm}(x + \text{SubLayer}(x))

Encoder-Decoder Architecture

Why Transformers Work

  1. Parallelization: Unlike RNNs, all positions are processed simultaneously
  2. Long-range dependencies: Direct connections between any two positions
  3. Scalability: Efficiently scales with data and compute

Conclusion

The Transformer has enabled GPT, BERT, Vision Transformers, and countless other breakthroughs. Understanding its components from scratch is essential for modern ML research.

References

  1. Vaswani et al., "Attention Is All You Need," NeurIPS 2017.
  2. Devlin et al., "BERT: Pre-training of Deep Bidirectional Transformers," NAACL 2019.