Transformer Architecture

Summary: The Transformer (Vaswani et al., 2017) is a sequence transduction model based entirely on attention mechanisms, dispensing with recurrence and convolutions. It uses an encoder-decoder stack with multi-head self-attention, position-wise feed-forward networks, residual connections, layer normalization, and sinusoidal positional encodings, achieving SOTA on translation tasks with significantly less training time.

Architecture Overview

Diagram: (Mermaid diagram - view source for diagram code)

flowchart TB
    subgraph INPUT["Input Processing"]
        SRC[Source Tokens] --> SRC_EMB[Token Embedding]
        TGT[Target Tokens] --> TGT_EMB[Token Embedding]
        SRC_EMB --> POS_ENC_S[Positional Encoding]
        TGT_EMB --> POS_ENC_T[Positional Encoding]
    end
subgraph ENCODER["Encoder Stack (N=6)"]
    POS_ENC_S --> ENC_1[Encoder Layer 1]
    ENC_1 --> ENC_2[Encoder Layer 2]
    ENC_2 --> ENC_N[Encoder Layer N]
    
    subgraph ENC_LAYER["Encoder Layer"]
        ENC_IN[Input] --> MHSA[Multi-Head Self-Attention]
        MHSA --> ADD_NORM1[Add & Norm]
        ADD_NORM1 --> FFN[Position-wise FFN]
        FFN --> ADD_NORM2[Add & Norm]
        ADD_NORM2 --> ENC_OUT[Output]
    end
end

subgraph DECODER["Decoder Stack (N=6)"]
    POS_ENC_T --> DEC_1[Decoder Layer 1]
    DEC_1 --> DEC_2[Decoder Layer 2]
    DEC_2 --> DEC_N[Decoder Layer N]
    
    subgraph DEC_LAYER["Decoder Layer"]
        DEC_IN[Input] --> MMHSA[Masked Multi-Head Self-Attention]
        MMHSA --> ADD_NORM_D1[Add & Norm]
        ADD_NORM_D1 --> ENC_DEC_ATT[Encoder-Decoder Attention]
        ENC_N -->|K, V| ENC_DEC_ATT
        ADD_NORM_D1 -->|Q| ENC_DEC_ATT
        ENC_DEC_ATT --> ADD_NORM_D2[Add & Norm]
        ADD_NORM_D2 --> FFN_D[Position-wise FFN]
        FFN_D --> ADD_NORM_D3[Add & Norm]
        ADD_NORM_D3 --> DEC_OUT[Output]
    end
end

DEC_N --> LINEAR[Linear + Softmax]
LINEAR --> OUTPUT[Output Probabilities]

classDef input fill:#e3f2fd,stroke:#90caf9;
classDef enc fill:#e8f5e9,stroke:#c8e6c9;
classDef dec fill:#fff3e0,stroke:#ffcc02;
classDef out fill:#fce4ec,stroke:#f8bbd0;
class SRC,TGT,SRC_EMB,TGT_EMB,POS_ENC_S,POS_ENC_T input;
class ENC_1,ENC_2,ENC_N,MHSA,ADD_NORM1,FFN,ADD_NORM2 enc;
class DEC_1,DEC_2,DEC_N,MMHSA,ADD_NORM_D1,ENC_DEC_ATT,ADD_NORM_D2,FFN_D,ADD_NORM_D3 dec;
class LINEAR,OUTPUT out;

Overview

The Transformer architecture replaced the dominant recurrent/convolutional sequence models with a fully attention-based approach. Key innovations:

  • Parallelization: Self-attention connects all positions in O(1) sequential operations vs O(n) for RNNs
  • Long-range dependencies: Constant path length between any two positions
  • Multi-head attention: Allows joint attention to different representation subspaces
  • Positional encoding: Sinusoidal encodings inject sequence order information

Key Components

Encoder Stack

  • N=6 identical layers
  • Each layer: Multi-head self-attention → Position-wise FFN
  • Residual connections + LayerNorm around each sub-layer
  • Output dimension: d_model=512

Decoder Stack

  • N=6 identical layers
  • Each layer: Masked multi-head self-attention → Encoder-decoder attention → Position-wise FFN
  • Residual connections + LayerNorm
  • Masking prevents attending to future positions (auto-regressive property)

Attention Mechanism

See attention-mechanism and multi-head-attention for detailed breakdown.

Positional Encoding

See positional-encoding for sinusoidal encoding details.

Feed-Forward Networks

Position-wise FFN: FFN(x) = max(0, xW₁ + b₁)W₂ + b₂

  • Inner dimension d_ff = 2048
  • Applied identically to each position

Mathematical Formulation

In plain English: The Transformer processes sequences by computing attention weights between all token pairs, then mixing representations through feed-forward layers. Multi-head attention runs multiple attention computations in parallel on projected subspaces.

Attention(Q, K, V) = softmax((QK^T)/(sqrt(d_k)))V
MultiHead(Q, K, V) = Concat(head_1, ..., head_h)W^O

where head_i = Attention(QW^Q_i, KW^K_i, VW^V_i)

Positional encoding:

PE_{(pos, 2i)} = sin(pos / 10000^{2i/d_{model}})
PE_{(pos, 2i+1)} = cos(pos / 10000^{2i/d_{model}})

Variants / Extensions

  • BERT (Devlin et al., 2018): Encoder-only, bidirectional, masked language modeling
  • GPT (Radford et al., 2018): Decoder-only, causal LM, left-to-right generation
  • T5 (Raffel et al., 2019): Encoder-decoder, unified text-to-text framework
  • Vision Transformer (ViT) (Dosovitskiy et al., 2020): Image patches as tokens
  • Swin Transformer (Liu et al., 2021): Hierarchical with shifted windows
  • FlashAttention (Dao et al., 2022): IO-aware exact attention with tiling

Applications

  • Machine translation (original task)
  • Language modeling (BERT, GPT, T5, LLaMA)
  • Computer vision (ViT, Swin, DETR)
  • Speech recognition (Speech-Transformer, Conformer)
  • Multimodal (CLIP, BLIP, Flamingo)
  • Code generation (Codex, CodeLlama)

Historical Context

Year Work Contribution
2017 Transformer (Vaswani et al.) Original architecture, SOTA translation
2018 BERT Bidirectional encoder, pre-training paradigm
2018 GPT Decoder-only, generative pre-training
2019 Transformer-XL Segment-level recurrence for longer context
2020 ViT Pure transformer for image classification
2022 FlashAttention IO-aware kernel for memory-efficient attention

Related Concepts

Sources

  • arxiv-1706.03762: Sections 1, 3.1–3.7, 4, 5.1 — Transformer architecture, attention, training details