Module 6

The Transformer Architecture

Positional encoding, transformer blocks, residual connections, layer normalisation, and the full language model head.

Learning Objectives

  • Implement sinusoidal positional encoding
  • Build a full transformer block (attention + FFN + residuals)
  • Understand why Transformers revolutionised NLP
  • Assemble a complete language model

6.1 Positional Encoding

Unlike RNNs, Transformers have no built-in notion of position. "The dog bites man" and "Man bites the dog" look identical to attention. Positional encodings inject position information as an additive signal.

Sinusoidal Encoding

PE(pos, 2i) = sin(pos / 10000^(2i / d_model)) PE(pos, 2i+1) = cos(pos / 10000^(2i / d_model)) Lower dimensions change slowly → coarse position info Higher dimensions change rapidly → fine-grained position info
Haskell — positional encoding matrix
positionalEncoding
  :: Int           -- sequence length
  -> Int           -- model dimension
  -> Matrix Double
positionalEncoding seqLen d_model =
  [[pe_func pos dim | dim <- [0 .. d_model - 1]] | pos <- [0 .. seqLen - 1]]
  where
    pe_func pos dim
      | dim `mod` 2 == 0 =
          sin (fromIntegral pos / (10000 ** (fromIntegral dim / fromIntegral d_model)))
      | otherwise =
          cos (fromIntegral pos / (10000 ** (fromIntegral (dim - 1) / fromIntegral d_model)))

6.2 Transformer Block

Block Structure

Input ↓ LayerNorm ↓ Multi-Head Self-Attention ↓ + (residual connection) Output₁ = Input + Attention(LayerNorm(Input)) ↓ LayerNorm ↓ Feed-Forward Network (Linear → ReLU → Linear) ↓ + (residual connection) Output = Output₁ + FFN(LayerNorm(Output₁))

Layer Normalisation

Haskell — layer norm
layerNorm :: Matrix Double -> Double -> Matrix Double
layerNorm input eps =
  let mean       = average input
      variance   = average (map (^2) (map (subtract mean) input))
      normalized = map (\x -> (x - mean) / sqrt (variance + eps)) input
  in normalized

Feed-Forward Network

Haskell — FFN (expansion ratio typically 4×)
data FeedForward = FeedForward
  { firstLinear  :: (Matrix Double, Vector Double)
  , secondLinear :: (Matrix Double, Vector Double)
  }

feedForward :: FeedForward -> Matrix Double -> Matrix Double
feedForward ff input =
  let (w1, b1) = firstLinear ff
      (w2, b2) = secondLinear ff
      hidden   = fmap relu (matmul input w1 + b1)  -- 4× expansion
      output   = matmul hidden w2 + b2
  in output

Full Transformer Block

Haskell — transformer block
data TransformerBlock = TransformerBlock
  { mhAttention :: MultiHeadAttention
  , ffNetwork   :: FeedForward
  }

transformerBlock :: TransformerBlock -> Matrix Double -> Matrix Double
transformerBlock block input =
  -- Self-attention with residual
  let norm1        = layerNorm input 1e-6
      attended     = multiHeadAttention (mhAttention block) norm1 norm1 norm1
      afterAttn    = input + attended        -- Residual
      -- Feed-forward with residual
      norm2        = layerNorm afterAttn 1e-6
      ffOut        = feedForward (ffNetwork block) norm2
  in afterAttn + ffOut                       -- Residual
Residual connections (x + f(x)) allow gradients to flow unchanged through deep stacks. Without them, training networks deeper than ~10 layers is nearly impossible.

6.3 Complete Language Model

Haskell — language model forward pass
data LanguageModel = LanguageModel
  { embedding        :: Matrix Double
  , posEncoding      :: Matrix Double     -- Cached PE matrix
  , encoderBlocks    :: [TransformerBlock]
  , outputProjection :: (Matrix Double, Vector Double)
  , vocabSize        :: Int
  }

forward :: LanguageModel -> [Int] -> Matrix Double
forward model tokenIds =
  let embeds     = map (embedding model !!) tokenIds
      withPos    = zipWith (zipWith (+)) embeds (rows (posEncoding model))
      -- Process through transformer blocks
      encoded    = foldl (\x block -> transformerBlock block x) withPos
                         (encoderBlocks model)
      -- Project to vocabulary
      (w, b)     = outputProjection model
      logits     = matmul encoded w + b   -- (seq_len, vocab_size)
  in logits

-- Cross-entropy loss
crossEntropyLoss :: Matrix Double -> [Int] -> Double
crossEntropyLoss logits targets =
  let losses = [(-log (softmax (logits !! i) !! (targets !! i)))
                | i <- [0..length targets - 1]]
  in average losses

Reading & Viewing

Essential

Attention Is All You Need (2017) Vaswani et al. — NeurIPS Read the full paper again after Module 5 — the architecture section will now make complete sense. arxiv.org/abs/1706.03762

Blog (highly recommended)

The Illustrated Transformer Jay Alammar The best visual explanation of the Transformer available. Read before or alongside the paper. jalammar.github.io/illustrated-transformer/
Self-Attention from Scratch Sebastian Raschka sebastianraschka.com
The Annotated Transformer Harvard NLP Full PyTorch implementation side-by-side with the paper. ~2000 lines of annotated code. nlp.seas.harvard.edu

Video

Attention Is All You Need — Yannic Kilcher YouTube