Module 12

Advanced Topics & Extensions

Ongoing exploration: improved normalisation, activation functions, positional encodings, sparse attention, and quantisation.

This module is a starting point for continued research. There are no fixed implementation tasks — pick the topics most relevant to your experiments.

12.1 Normalisation Variants

RMSNorm

A simpler alternative to LayerNorm that only normalises by the root mean square, without mean-centring. Used in LLaMA and other modern architectures.

RMSNorm(x) = x / RMS(x) · γ RMS(x) = √(1/n · Σ xᵢ²)

GroupNorm

Normalises within groups of channels — useful in settings where batch size is too small for BatchNorm to be stable.

Root Mean Square Layer Normalisation (2019) Zhang & Sennrich arxiv.org/abs/1910.07468

12.2 Improved Activation Functions

GELU

Gaussian Error Linear Units — a smoother alternative to ReLU, used in GPT-2 and BERT.

GELU(x) = 0.5 · x · (1 + tanh(√(2/π) · (x + 0.044715 · x³)))

SwiGLU

Swish-Gated Linear Unit — used in PaLM, LLaMA, and many other recent LLMs. Replaces the FFN's ReLU with a gated variant.

SwiGLU(x, W, V, b, c) = Swish(xW + b) ⊙ (xV + c)
GLU Variants Improve Transformer (2020) Noam Shazeer arxiv.org/abs/2002.05202

12.3 Positional Encoding Improvements

RoPE — Rotary Position Embeddings

Rather than adding positional encodings to embeddings, RoPE encodes position by rotating the query and key vectors. Extrapolates better to sequences longer than those seen in training and is used in LLaMA, Mistral, and most modern open models.

q̃ₘ = qₘ · e^(imθ) (complex rotation by position m) k̃ₙ = kₙ · e^(inθ) Attention score ∝ Re(q̃ₘ · k̃*ₙ) depends only on (m-n)

Relative Position Biases

An alternative approach: bias the attention logits by a learned function of relative distance rather than absolute position (e.g., ALiBi, T5's relative position bias).

RoFormer: Enhanced Transformer with Rotary Position Embedding (2021) Su, Lu, Pan et al. arxiv.org/abs/2104.09864

12.4 Sparse Attention

Full self-attention is O(n²) in sequence length. For 1B+ models on long documents this becomes a bottleneck. Sparse attention variants restrict which positions attend to which:

  • Local attention — each token attends only to a fixed window of nearby tokens
  • Strided attention — attend to every k-th token in addition to nearby ones
  • BigBird / Longformer — combine local + global + random attention patterns
Longformer: The Long-Document Transformer (2020) Beltagy, Peters, Cohan arxiv.org/abs/2004.05150
Big Bird: Transformers for Longer Sequences (2020) arxiv.org/abs/2007.14062

12.5 Quantisation

Reduce model size for inference by storing weights in lower precision (8-bit or 4-bit integers). Dramatically reduces memory — a 7B model that needs 14GB at FP16 fits in ~4GB at 4-bit.

LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (2022) Dettmers, Lewis, Belkada, Zettlemoyer arxiv.org/abs/2208.07339

Where to Go Next

Next Challenges

  • Instruction Fine-Tuning — train the model to follow instructions (RLHF/DPO)
  • Multimodal Models — add a vision encoder (ViT) alongside the text model
  • Code Generation — pretrain or fine-tune on code data
  • Context Extension — use RoPE or ALiBi to extend the context window
  • Efficiency — flash attention, quantisation-aware training

Essential Papers to Read Next

Language Models are Few-Shot Learners (GPT-3, 2020) arxiv.org/abs/2005.14165
Training Language Models to Follow Instructions with Human Feedback (InstructGPT/RLHF, 2022) arxiv.org/abs/2203.02155
Constitutional AI: Harmlessness from AI Feedback (2022) arxiv.org/abs/2212.08073
Efficient Transformers: A Survey (2020) arxiv.org/abs/2009.06732

Haskell Libraries Worth Exploring

  • Grenade — neural network library for Haskell
  • HMatrix — numeric linear algebra backed by LAPACK/BLAS

Rust Crates Worth Exploring

  • tch-rs — Rust bindings to PyTorch's libtorch
  • burn — pure-Rust deep learning framework with ROCm/CUDA/WGPU backends
  • candle — HuggingFace's minimal Rust inference framework