Skip to content
AIAI Wranglers

5 Deep Learning Building Blocks

In Chapter 3 we formalised the generative modelling problem as finding a map gΞΈ:𝒡→𝒳 such that the push-forward gΞΈ#pZ approximates pdata. But what is gΞΈ? In virtually every modern generative model, it is a deep neural network, a composition of simple, parameterised, differentiable functions that can approximate extraordinarily complex mappings.

This chapter introduces the neural network building blocks that the rest of the book relies on: feedforward layers, activation functions, convolutional and transposed-convolutional layers, residual connections, normalisation layers, and the Transformer's attention mechanism. For each building block, we give the mathematical definition, state its key properties, and explain why it appears in specific generative model architectures.

Historical Note.

The idea of modelling computation with networks of artificial neurons dates to McCulloch and Pitts (1943) [1]. Rosenblatt's perceptron (1958) [2] could learn from data, but Minsky and Papert (1969) [3] showed it could not solve non-linearly-separable problems like XOR, triggering the first β€œAI winter.” The field revived in the 1980s with backpropagation [4] and again in the 2010s when GPUs, large datasets, and deeper architectures led to breakthroughs in vision [5] and language [6]. Today, deep networks with billions of parameters are the engine behind every generative model studied in this book.

Feedforward Neural Networks

Definition 1 (Feedforward Neural Network).

An L-layer feedforward neural network (also called a multilayer perceptron, MLP) computes a function fΞΈ:ℝd0⁑→ℝdL⁑ by alternating affine transformations and element-wise nonlinearities: (FFN)𝒛(β„“)=W(β„“)𝒂(β„“βˆ’1)+𝒃(β„“),𝒂(β„“)=Οƒ(𝒛(β„“)),β„“=1,…,L, where 𝒂(0)=𝒙 is the input, W(β„“)βˆˆβ„dβ„“Γ—dβ„“βˆ’1⁑ is a weight matrix, 𝒃(β„“)βˆˆβ„dℓ⁑ is a bias vector, and Οƒ is a nonlinear activation function. The output is 𝒂(L) (or 𝒛(L) if the last layer is linear). The collection of all weights and biases is ΞΈ={W(β„“),𝒃(β„“)}β„“=1L.

A 3-layer feedforward neural network (L=3) with layer widths d0=3, d1=d2=4, d3=2. Each arrow represents a learnable weight.

Activation Functions

The activation function Οƒ introduces nonlinearity. Without it, the composition of affine layers would itself be affine, and the network could only represent linear functions.

Definition 2 (Common Activation Functions).

(Sigmoid)Sigmoid:Οƒ(z)=11+eβˆ’z,Tanh:Οƒ(z)=ezβˆ’eβˆ’zez+eβˆ’z,ReLU:Οƒ(z)=max⁑(0,z),LeakyΒ ReLU:Οƒ(z)=max⁑(Ξ±z,z),Ξ±β‰ͺ1,GELU:Οƒ(z)=zΞ¦(z),SiLUΒ /Β Swish:Οƒ(z)=zβ‹…sigmoid(z)=z1+eβˆ’z, where Ξ¦ is the standard Gaussian CDF.

Remark 1 (Which Activation Where?).

  • ReLU and Leaky ReLU are standard in CNNs and GAN discriminators [7].

  • GELU and SiLU/Swish are used in Transformers (GPT, diffusion U-Nets) due to smoother gradients.

  • Sigmoid appears as the output activation for GAN discriminators (outputting a probability) and in LSTM gates.

  • Tanh is often used in the final layer of image generators (mapping to [βˆ’1,1] pixel range).

Three commonly used activation functions. Sigmoid saturates to [0,1]; tanh saturates to [βˆ’1,1]; ReLU is unbounded but zero for negative inputs.

The Universal Approximation Theorem

Why are neural networks so effective as the generator gΞΈ? The following foundational result answers this question.

Theorem 1 (Universal Approximation - Cybenko 1989, Hornik 1991).

Let Οƒ be any continuous, non-constant activation function. For any continuous function f:[0,1]d→ℝ⁑ and any Ο΅>0, there exists a single-hidden-layer network f^(𝒙)=βˆ‘j=1NΞ±jΟƒ(π’˜jβŠ€π’™+bj) with sufficiently large N such that supπ’™βˆˆ[0,1]d⁑|f(𝒙)βˆ’f^(𝒙)|<Ο΅.

Caution.

The universal approximation theorem guarantees existence but says nothing about how to find the weights, nor about the required width N (which can be exponential in d). In practice, depth is far more efficient than width: deep networks can represent functions with exponentially fewer parameters than shallow ones. This is why generative models use networks with dozens to hundreds of layers, not single-layer networks with billions of neurons.

Insight.

The universal approximation theorem tells us that a neural network gΞΈ is, in principle, flexible enough to represent any continuous push-forward map, from a Gaussian latent space to the space of natural images. The practical challenge is finding good weights ΞΈ via training, which is the subject of the next section.

Training Neural Networks

Loss Functions and Gradient Descent

Training a neural network means finding parameters ΞΈ that minimise a loss function β„’(ΞΈ) that measures discrepancy between the network's predictions and the training data: (LOSS)ΞΈ^=arg minθ⁑ℒ(ΞΈ)=arg minθ⁑1nβˆ‘i=1nβ„“(fΞΈ(xi),yi), where β„“ is a per-sample loss (e.g. mean squared error, cross-entropy).

Definition 3 (Gradient Descent).

Gradient descent iteratively updates parameters in the direction of steepest decrease of β„’: (GD)ΞΈt+1=ΞΈtβˆ’Ξ·βˆ‡ΞΈβ„’(ΞΈt), where Ξ·>0 is the learning rate.

Definition 4 (Stochastic Gradient Descent (SGD)).

When the dataset is large, computing the full gradient is expensive. SGD approximates it using a random mini-batch β„¬βŠ‚{1,…,n} of size B: (SGD)ΞΈt+1=ΞΈtβˆ’Ξ·1Bβˆ‘iβˆˆβ„¬tβˆ‡ΞΈβ„“(fΞΈ(xi),yi). The mini-batch gradient is an unbiased estimate of the full gradient, with variance proportional to 1/B.

Remark 2 (Adam and Modern Optimisers).

In practice, most generative models use Adam (Adaptive Moment Estimation), which maintains running averages of both the gradient and its square, adapting the learning rate per parameter. Variants like AdamW (with weight decay) are standard for Transformers, while SGD with momentum remains popular for CNNs.

Backpropagation

Computing βˆ‡ΞΈβ„’ for a deep network uses the chain rule of calculus, applied layer by layer from the output back to the input. This algorithm, called backpropagation [4], is the workhorse of deep learning.

Algorithm 1 (Backpropagation (Sketch)).

  1. Forward pass: compute 𝒛(β„“) and 𝒂(β„“) for β„“=1,…,L using (FFN).

  2. Compute loss: evaluate β„’ at the output.

  3. Backward pass: for β„“=L,Lβˆ’1,…,1, compute βˆ‚β„’βˆ‚W(β„“)=βˆ‚β„’βˆ‚π’›(β„“)β‹…(𝒂(β„“βˆ’1))⊀,βˆ‚β„’βˆ‚π’›(β„“βˆ’1)=(W(β„“))βŠ€βˆ‚β„’βˆ‚π’›(β„“)βŠ™Οƒβ€²(𝒛(β„“βˆ’1)).

  4. Update: apply gradient descent ((GD) or (SGD)).

Key Idea.

Backpropagation computes the gradient of a loss with respect to all parameters in O(forwardΒ pass) time, regardless of how many parameters there are. This efficiency is what makes training networks with millions (or billions) of parameters feasible. It is implemented automatically by frameworks like PyTorch via automatic differentiation.

Normalisation Layers

Deep networks are prone to β€œinternal covariate shift,” where the distribution of each layer's inputs shifts as preceding layers update. Normalisation layers stabilise training by re-centring and re-scaling activations.

Definition 5 (Batch Normalisation).

Given a mini-batch of activations {xi}i=1B at a particular layer, batch normalisation (BN) [8] computes (Batchnorm)x^i=xiβˆ’ΞΌβ„¬Οƒβ„¬2+Ο΅,yi=Ξ³x^i+Ξ², where μℬ and σℬ2 are the mean and variance over the batch, Ξ³ and Ξ² are learnable parameters, and Ο΅ is a small constant for numerical stability. BN normalises across the batch dimension.

Definition 6 (Layer Normalisation).

Layer normalisation (LN) [9] normalises across the feature dimension for each sample independently: (Layernorm)x^ij=xijβˆ’ΞΌiΟƒi2+Ο΅,yij=Ξ³jx^ij+Ξ²j, where ΞΌi=1Fβˆ‘jxij and Οƒi2=1Fβˆ‘j(xijβˆ’ΞΌi)2 are per-sample statistics (with F features).

Remark 3 (Which Normalisation Where?).

  • Batch normalisation: default in CNN-based discriminators and early GAN generators [7].

  • Layer normalisation: standard in Transformers [6] and diffusion-model U-Nets, because it does not depend on batch size and works with variable-length sequences.

  • Group normalisation and instance normalisation: used in style transfer and some GAN generators (e.g. StyleGAN [10]).

Convolutional Neural Networks

Images have spatial structure: nearby pixels are correlated, and patterns (edges, textures) can appear anywhere in the image. Convolutional neural networks (CNNs) [26] exploit this structure by replacing full matrix multiplications with local, weight-shared convolution operations.

Definition 7 (2D Convolution Layer).

Given an input feature map Xβˆˆβ„CinΓ—HΓ—W⁑ and a bank of Cout kernels Kkβˆˆβ„CinΓ—khΓ—kw⁑, the output of a 2D convolution is (Conv2d)Yk(i,j)=βˆ‘c=1Cinβˆ‘m=0khβˆ’1βˆ‘n=0kwβˆ’1Kk(c,m,n)X(c,iβ‹…s+m,jβ‹…s+n)+bk, where s is the stride and bk is a bias.

Example 1 (Convolution Arithmetic).

A 3Γ—3 convolution with stride 1 and β€œsame” padding applied to a 32Γ—32 image with 3 input channels and 64 output channels has 64Γ—3Γ—3Γ—3+64=1,792 parameters, far fewer than a fully connected layer (64Γ—3Γ—32Γ—32=196,608), yet it captures all local 3Γ—3 patterns.

Insight.

CNNs are the backbone of image-based generative models:

  • GAN discriminators are typically deep CNNs that progressively downsample the image (strided convolutions) to produce a real/fake score [7].

  • GAN generators use transposed convolutions (Definition 8) to progressively upsample a latent vector into a full-resolution image.

  • VAE encoders downsample with convolutions; VAE decoders upsample with transposed convolutions.

  • Diffusion-model denoisers use a U-Net (The U-Net Architecture), which combines downsampling and upsampling paths.

Transposed Convolutions

Generation requires going from a small latent representation to a large output image. Transposed convolutions (sometimes misleadingly called β€œdeconvolutions”) perform learned upsampling.

Definition 8 (Transposed Convolution).

A transposed convolution with kernel K, stride s, and padding p is the adjoint (transpose) of the corresponding forward convolution. Concretely, it inserts sβˆ’1 zeros between input elements, then applies a standard convolution with kernel K. If the forward convolution maps a HΓ—W input to ⌊(Hβˆ’k+2p)/sβŒ‹+1 outputs, the transposed convolution maps a smaller input back to HΓ—W.

Remark 4 (Checkerboard Artifacts).

Transposed convolutions are known to produce β€œcheckerboard artifacts,” regular grid patterns visible in generated images, when the kernel size is not divisible by the stride. Alternative upsampling strategies include nearest-neighbour interpolation followed by a regular convolution, which avoids this issue and is used in modern generators like StyleGAN [11].

Pooling

Pooling layers reduce spatial dimensions by summarising local regions. Max pooling takes the maximum value in each window; average pooling takes the mean. In modern generative models, pooling has largely been replaced by strided convolutions (which are learnable), but global average pooling remains common in discriminators.

Residual Networks and Skip Connections

Training very deep networks (50+ layers) is difficult because gradients can vanish or explode as they propagate through many layers. Residual connections (skip connections) [27] solve this by adding the input of a block directly to its output.

Definition 9 (Residual Block).

A residual block computes (Residual)𝒂out=𝒂in+β„±(𝒂in), where β„± is a sub-network (typically two convolutional layers with normalisation and activation). The β€œ+𝒂in” is the skip connection.

A residual block: the input 𝒂in is added to the output of two convolutional layers. Gradients can flow directly through the skip connection, enabling training of very deep networks.

Key Idea.

Why do residual connections matter for generative models?

  • They enable very deep generators and discriminators (100+ layers), which have greater representational power.

  • The identity mapping through the skip connection ensures that the gradient βˆ‚β„’/βˆ‚π’‚in always has a direct path, preventing vanishing gradients.

  • Residual blocks are the core building block of U-Nets (The U-Net Architecture), used in virtually all diffusion models.

The U-Net Architecture

The U-Net, originally proposed for biomedical image segmentation, has become the standard architecture for diffusion-model denoisers (Chapter 16). It consists of:

  1. an encoder (downsampling path) that progressively reduces spatial resolution while increasing channels,

  2. a decoder (upsampling path) that progressively restores resolution, and

  3. long skip connections that concatenate encoder features to the corresponding decoder level, preserving fine-grained spatial information.

The U-Net architecture. The encoder reduces spatial resolution (downsampling); the decoder restores it (upsampling). Long skip connections (dashed blue) concatenate encoder features to the decoder, preserving spatial detail.

Insight.

In diffusion models (Chapter 16), the U-Net takes as input a noisy image xt and a time step t, and predicts the noise ϡθ(xt,t) that was added. The skip connections are essential: low-level encoder features (edges, textures) help the decoder reconstruct fine details that would otherwise be lost at the bottleneck. The time step t is typically injected via learned embeddings added to each residual block, telling the network how much noise to expect.

Sequence Models: From RNNs to Attention

Many data types are sequential: text, audio, time series. Generative models for sequences need architectures that can process variable-length inputs and capture long-range dependencies.

Recurrent Neural Networks

Definition 10 (Recurrent Neural Network).

A recurrent neural network (RNN) processes a sequence (x1,…,xT) by maintaining a hidden state 𝒉t that is updated at each time step: (RNN)𝒉t=Οƒ(Wh𝒉tβˆ’1+Wx𝒙t+𝒃),π’št=Wy𝒉t+𝒃y, where Wh,Wx,Wy are weight matrices shared across all time steps.

Remark 5 (Vanishing and Exploding Gradients).

Training RNNs via backpropagation through time (BPTT) requires multiplying gradient matrices across T time steps. If the spectral radius of Wh is less than 1, gradients vanish exponentially; if greater than 1, they explode. This limits the ability of vanilla RNNs to learn long-range dependencies.

Long Short-Term Memory (LSTM)

Definition 11 (LSTM Cell).

The LSTM introduces a cell state 𝒄t and three gates to control information flow: (LSTM Forget)𝒇t=Οƒ(Wf[𝒉tβˆ’1,𝒙t]+𝒃f),(forgetΒ gate)π’Št=Οƒ(Wi[𝒉tβˆ’1,𝒙t]+𝒃i),(inputΒ gate)𝒄~t=tanh⁑(Wc[𝒉tβˆ’1,𝒙t]+𝒃c),(candidate)𝒄t=𝒇tβŠ™π’„tβˆ’1+π’ŠtβŠ™π’„~t,(cellΒ update)𝒐t=Οƒ(Wo[𝒉tβˆ’1,𝒙t]+𝒃o),(outputΒ gate)𝒉t=𝒐tβŠ™tanh⁑(𝒄t).(hiddenΒ state) The forget gate 𝒇t decides what to discard from the cell state; the input gate π’Št decides what new information to store.

Remark 6 (From LSTMs to Transformers).

LSTMs dominated sequence modelling from 2014–2017, powering early neural machine translation and text generation. However, their sequential nature (each step depends on the previous) prevents parallelisation across time steps. The Transformer architecture [6] replaced recurrence with attention, enabling massive parallelism and scaling to billions of parameters, a breakthrough that led to GPT, BERT, and modern large language models.

The Transformer

The Transformer [6], introduced in the 2017 paper β€œAttention Is All You Need,” has become the dominant architecture for both language and vision tasks. Its core innovation is the self-attention mechanism, which allows every position in a sequence to attend to every other position in constant depth.

Scaled Dot-Product Attention

Definition 12 (Scaled Dot-Product Attention).

Given queries Qβˆˆβ„nΓ—dk⁑, keys Kβˆˆβ„nΓ—dk⁑, and values Vβˆˆβ„nΓ—dv⁑, the scaled dot-product attention is (Attention)Attention(Q,K,V)=softmax⁑(QK⊀dk)V. The softmax⁑ is applied row-wise, so each row of the output is a weighted combination of the value vectors, where the weights reflect the similarity between the corresponding query and all keys.

Remark 7 (Why Scale by dk?).

When dk is large, the dot products QK⊀ tend to have large magnitude (variance β‰ˆdk if entries have unit variance), pushing the softmax into saturation where gradients are near zero. Dividing by dk keeps the variance β‰ˆ1, ensuring well-behaved gradients throughout training.

Scaled dot-product attention: queries and keys are compared via dot product, scaled, passed through softmax, and used to weight the values.

Multi-Head Attention

Definition 13 (Multi-Head Attention).

Multi-head attention runs h attention functions in parallel, each with its own learned projections: (Multihead)MultiHead(Q,K,V)=Concat(head1,…,headh)WO,headi=Attention(QWiQ,KWiK,VWiV), where WiQβˆˆβ„dΓ—dk⁑, WiKβˆˆβ„dΓ—dk⁑, WiVβˆˆβ„dΓ—dv⁑, and WOβˆˆβ„hdvΓ—d⁑ are learned projection matrices.

Each head can learn to attend to different aspects of the input (e.g. one head for syntactic relations, another for semantic similarity).

Positional Encoding

Since attention operates on sets (it is permutation-equivariant), we must inject position information explicitly.

Definition 14 (Sinusoidal Positional Encoding).

The original Transformer [6] uses fixed sinusoidal encodings: (Posenc)PE(pos,2i)=sin⁑(pos100002i/d),PE(pos,2i+1)=cos⁑(pos100002i/d), where pos is the position index and i is the dimension index. These encodings are added to the input embeddings.

Remark 8 (Learnable vs. Sinusoidal Encodings).

Modern models (GPT-2/3, Vision Transformers) often use learned positional embeddings instead. Rotary positional embeddings (RoPE) and ALiBi are recent alternatives that improve length generalisation. We develop RoPE in detail in Rotary Position Embeddings (RoPE) below.

Masked Self-Attention for Autoregressive Generation

In autoregressive generation, position i must not attend to any future position j>i (since those tokens have not been generated yet). This is enforced by adding a causal mask: (Causal MASK)Attention(Q,K,V)=softmax⁑(QK⊀dk+M)V,Mij={0ifΒ j≀i,βˆ’βˆžifΒ j>i. The βˆ’βˆž entries become zero after softmax, preventing information leakage from the future.

Insight.

Masked self-attention is the mechanism behind GPT-style autoregressive language models and autoregressive image models (e.g. PixelCNN uses masked convolutions for the same purpose). The chain-rule factorisation p(𝒙)=∏ip(xi|x<i) from Chapter 3 is implemented by a Transformer decoder with causal masking: each token's prediction depends only on previous tokens.

Cross-Attention for Conditional Generation

In cross-attention, the queries come from one sequence (e.g. the image being generated) while the keys and values come from another (e.g. a text prompt): (Cross ATTN)CrossAttn(Ximage,Xtext)=softmax⁑(QimgKtxt⊀dk)Vtxt. This is how text-to-image diffusion models (Stable Diffusion, DALL-E) inject the text conditioning signal into the U-Net denoiser.

The Transformer Block

A complete Transformer encoder block combines multi-head self-attention, a feedforward network, residual connections, and layer normalisation:

Algorithm 2 (Transformer Encoder Block).

  1. 𝒉′=LayerNorm(𝒙+MultiHead(𝒙,𝒙,𝒙)) (self-attention + residual + LN)

  2. 𝒉=LayerNorm(𝒉′+FFN(𝒉′)) (feedforward + residual + LN)

The FFN is typically a two-layer MLP: FFN(𝒙)=W2Οƒ(W1𝒙+𝒃1)+𝒃2 with GELU or SiLU activation (or SwiGLU, SwiGLU Feedforward Networks).

Remark 9 (Pre-Norm vs. Post-Norm).

The original Transformer applies layer normalisation after the residual addition (β€œpost-norm”). Most modern implementations use β€œpre-norm” (normalise before attention/FFN), which improves training stability for very deep models.

Rotary Position Embeddings (RoPE)

The sinusoidal positional encoding of Definition 14 is additive: it adds a fixed vector to each token embedding, encoding the absolute position of the token in the sequence. Rotary Position Embeddings (RoPE) [28] take a fundamentally different approach. Instead of adding a position signal to the embeddings, RoPE rotates the query and key vectors in a position-dependent manner before computing the dot-product attention score. The crucial consequence is that the inner product between a rotated query at position m and a rotated key at position n depends only on their relative distance mβˆ’n, not on the absolute positions themselves. This property makes RoPE a natural fit for language modelling and other sequence tasks where relative position carries far more information than absolute position.

Definition 15 (Rotary Position Embedding).

Let d be the head dimension (assumed even) and define the frequency schedule (ROPE FREQ)ΞΈi=10000βˆ’2(iβˆ’1)/d,i=1,…,d/2. For each position m∈{0,1,2,…}, define the rotation matrix (ROPE)RΘm=diag⁑(R(mΞΈ1),R(mΞΈ2),…,R(mΞΈd/2))βˆˆβ„dΓ—d⁑, where each 2Γ—2 block is the standard planar rotation R(Ξ±)=(cosβ‘Ξ±βˆ’sin⁑αsinβ‘Ξ±βˆ’cos⁑α). Given an input token 𝒙m at position m, the RoPE-augmented query and key are (ROPE QK)𝒒~m=RΘmWQ𝒙m,π’Œ~n=RΘnWK𝒙n, where WQ,WKβˆˆβ„dΓ—d⁑ are the usual query and key projection matrices. The attention score between positions m and n is then 𝒒~mβŠ€π’Œ~n/d.

Note that RoPE modifies only the queries and keys; the value projections and the rest of the Transformer block remain unchanged. The rotation angles grow linearly with position: low-frequency components (small ΞΈi) rotate slowly, encoding long-range positional information, while high-frequency components (large ΞΈi) rotate rapidly, capturing fine-grained local structure.

Proposition 1 (Relative Position Property of RoPE).

Let 𝒒m=WQ𝒙m and π’Œn=WK𝒙n denote the unrotated query and key vectors. Then for any positions m,n, (ROPE Relative)𝒒~mβŠ€π’Œ~n=𝒒m⊀RΘnβˆ’mπ’Œn. In particular, the attention score depends on the relative position nβˆ’m only (through the rotation matrix RΘnβˆ’m), not on the absolute positions m and n individually.

Proof.

The key observation is that every 2Γ—2 rotation block is an orthogonal matrix: R(Ξ±)⊀=R(βˆ’Ξ±). Since the full rotation matrix RΘm is block-diagonal, its transpose is obtained by transposing each block independently, giving (RΘm)⊀=diag⁑(R(βˆ’mΞΈ1),R(βˆ’mΞΈ2),…,R(βˆ’mΞΈd/2))=RΞ˜βˆ’m. Now expand the inner product using (ROPE QK): (ROPE Proof STEP)𝒒~mβŠ€π’Œ~n=(RΘm𝒒m)⊀(RΘnπ’Œn)=𝒒m⊀(RΘm)⊀RΘnπ’Œn=𝒒m⊀RΞ˜βˆ’mRΘnπ’Œn. Because each 2Γ—2 block satisfies R(βˆ’Ξ±)R(Ξ²)=R(Ξ²βˆ’Ξ±), the block-diagonal product simplifies: RΞ˜βˆ’mRΘn=diag⁑(R((nβˆ’m)ΞΈ1),…,R((nβˆ’m)ΞΈd/2))=RΘnβˆ’m. Substituting back into yields 𝒒~mβŠ€π’Œ~n=𝒒m⊀RΘnβˆ’mπ’Œn, which completes the proof.

∎
Left: A two-dimensional query vector is rotated by multiples of ΞΈ as the position index m increases. The angle between any two position-rotated vectors depends only on their relative distance. Right: The full rotation matrix RΘm for d=8. Each 2Γ—2 block along the diagonal applies a rotation at a different frequency ΞΈi; all off-diagonal entries are zero.

Remark 10 (RoPE in Practice).

The base frequency ΞΈbase=10000 in (ROPE FREQ) works well for moderate context lengths (up to roughly 4096 tokens). For longer contexts in the range of 32K to 128K tokens, practitioners increase the base frequency: for instance, Qwen 2.5 [12] sets ΞΈbase=106 (sometimes called Adjusted Base Frequency). Interpolation methods such as YaRN [13] take a complementary approach, rescaling the rotation frequencies so that a model trained on short contexts can be extended to much longer sequences without full retraining. In both cases, the key insight is that the geometric structure of RoPE (block-diagonal rotations) makes it straightforward to manipulate the effective context window by adjusting only the frequency schedule.

Insight.

RoPE has become the de facto positional encoding for modern large language models, including LLaMA [14], Qwen [12], DeepSeek [15], and Mistral [16]. Beyond language, it also appears in Diffusion Transformers (DiT) [17] for image generation, where 2-D positional information is encoded by applying independent RoPE rotations along each spatial axis. The relative position property (Proposition 1) means the model can generalise to sequence lengths longer than those seen during training, a critical capability for serving LLMs in production where input lengths vary widely.

Efficient Attention: MQA, GQA, and KV Caching

During autoregressive generation (recall Masked Self-Attention for Autoregressive Generation), each new token requires attending to all previous tokens. With standard multi-head attention, this naΓ―vely means recomputing the full key and value matrices at every generation step, even though the keys and values for earlier positions have not changed. The KV cache eliminates this redundancy by storing previously computed keys and values and appending to them incrementally. However, the cache itself can become a formidable memory bottleneck: for a model with L layers and h attention heads, each with head dimension dk, it grows as O(Lhdkn) where n is the sequence length generated so far. This motivates Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), which shrink the KV cache by reducing the number of key-value heads.

Definition 16 (KV Cache).

At generation step t, instead of recomputing keys and values for the entire prefix (𝒙1,…,𝒙t), the KV cache proceeds as follows:

  1. Compute only the new key and value vectors: π’Œt=WK𝒙t, 𝒗t=WV𝒙t.

  2. Append them to the cache: (KV Cache)Kt=[Ktβˆ’1π’Œt⊀],Vt=[Vtβˆ’1𝒗t⊀].

  3. Compute the query for the new position: 𝒒t=WQ𝒙t.

  4. Compute attention for position t: outputt=softmax⁑(𝒒tKt⊀dk)Vt.

Since only a single query row is computed against the cached keys, the per-step complexity drops from O(n2d) (recomputing from scratch) to O(nd) (with cache), yielding a substantial speedup during generation.

Attention computation during autoregressive generation. Left: Without a KV cache, the full lower-triangular attention matrix is recomputed at each step. Right: With a KV cache, previous keys and values are stored, and only the new row (blue) is computed. The cache grows by one row per step.

Definition 17 (Multi-Query and Grouped-Query Attention).

Standard multi-head attention (Definition 13) maintains separate key and value projections for every head, so the KV cache stores 2hdk values per token per layer. Two variants reduce this cost:

  • Multi-Query Attention (MQA) [18]: All h query heads share a single set of key and value projections. Head i computes its own query Qi=XWiQ but uses shared keys K=XWK and values V=XWV. This reduces the KV cache by a factor of h.

  • Grouped-Query Attention (GQA) [19]: A middle ground between MHA and MQA. The h query heads are divided into G groups, each group sharing one key-value head: (GQA)headi=Attention(QWiQ,KWg(i)K,VWg(i)V),g(i)=⌊iGhβŒ‹. When G=h, GQA reduces to standard MHA; when G=1, it becomes MQA.

Comparison of key-value sharing strategies for h=8 query heads. Left: Standard multi-head attention (MHA) uses 8 independent KV heads. Middle: Grouped-query attention (GQA) with G=2 groups shares each KV head among 4 query heads. Right: Multi-query attention (MQA) uses a single KV head for all 8 query heads. Blue circles are query heads; orange circles are KV heads.

Remark 11 (KV Cache Memory Budget).

For a model with L layers, h attention heads, head dimension dk, and context length n, the KV cache stores 2Lhdkn values (the factor of 2 accounts for keys and values). Consider LLaMA-70B with L=80, h=64, dk=128, and n=4096: the KV cache requires 2Γ—80Γ—64Γ—128Γ—4096=5.4Γ—109 values, or approximately 10,GB in FP16 (2 bytes per value). GQA with G=8 groups instead of h=64 independent KV heads reduces the cache by a factor of 8Γ— to roughly 1.25,GB, making long-context serving practical on a single GPU. This saving is the primary reason that LLaMA-2 and most modern large language models adopt GQA.

Flash Attention

Computing the attention matrix S=QK⊀/dkβˆˆβ„NΓ—N⁑ and storing it in GPU memory is the principal bottleneck for long sequences. For N=8192 and d=128, the NΓ—N matrix alone requires 256,MB in FP32, and this is per head, per layer. The issue is not computation (modern GPUs have ample FLOPs) but memory bandwidth: reading and writing this matrix to GPU high-bandwidth memory (HBM) is slow relative to the arithmetic. Flash Attention [29] resolves this by computing exact attention through a tiling scheme that never materialises the full NΓ—N matrix in HBM.

Key Idea.

Modern GPUs have a memory hierarchy with two levels relevant to attention:

  • SRAM (on-chip, shared memory): very fast (∼19,TB/s bandwidth on an A100) but small (∼20,MB).

  • HBM (off-chip, high-bandwidth memory): large (∼80,GB) but comparatively slow (∼2,TB/s).

Standard attention writes the full NΓ—N score matrix to HBM, incurring O(N2) memory accesses. Flash Attention partitions Q, K, V into blocks that fit in SRAM, computes partial attention scores entirely on-chip using an online softmax trick, and writes only the final O(Nd) output back to HBM. The result is bit-for-bit identical to standard attention but uses far fewer HBM accesses.

Algorithm 3 (Flash Attention (Simplified)).

  1. Divide Q into Tr=⌈N/BrβŒ‰ row-blocks and K,V into Tc=⌈N/BcβŒ‰ column-blocks, where Br,Bc are chosen so that the working set fits in SRAM.

  2. Initialise the output O=πŸŽβˆˆβ„NΓ—d⁑, running row-wise maxima π’Ž=βˆ’βˆžβˆˆβ„N⁑, and running row-wise sums β„“=πŸŽβˆˆβ„N⁑.

  3. For each query block Qi,; i=1,…,Tr: enumerate[label=()]

  4. Load Qi from HBM to SRAM.

  5. For each key-value block (Kj,Vj),; j=1,…,Tc: itemize

  6. Load Kj,Vj from HBM to SRAM.

  7. Compute block scores on-chip: Sij=QiKj⊀/dk.

  8. Update running max: π’Žnew=max⁑(π’Ž,rowmax(Sij)).

  9. Rescale previous partial sums by exp⁑(π’Žoldβˆ’π’Žnew) and add new contributions exp⁑(Sijβˆ’π’Žnew)Vj.

  10. Update the running denominator β„“ analogously. itemize

  11. Write the final output block Oi back to HBM. enumerate

The crucial insight is that softmax can be computed incrementally: when a new block of scores arrives, we rescale the previous partial sums by the ratio exp⁑(moldβˆ’mnew) and accumulate the new block's contribution. No NΓ—N matrix is ever stored.

Flash Attention avoids materialising the NΓ—N attention matrix. Left: The attention matrix is logically divided into blocks; at any moment, only one block (amber) resides in fast on-chip SRAM. The outer loop iterates over query blocks; the inner loop over key-value blocks. Right: Memory hierarchy: blocks are loaded from slow HBM into fast SRAM for computation, and only the final output Oi is written back.

Proposition 2 (IO Complexity of Flash Attention).

Standard attention requires Θ(Nd+N2) HBM accesses: O(Nd) to read Q,K,V and write O, plus O(N2) to write and read the attention matrix S. Flash Attention requires Θ(N2d2M) HBM accesses, where M is the SRAM size (in elements). For typical GPU configurations where M=Θ(Nd), this simplifies to Θ(Nd), matching the cost of simply reading the inputs and writing the outputs.

Proof sketch.

The algorithm processes Trβ‹…Tc=O(N2/(BrBc)) block pairs. For each pair, it loads O(Brd+Bcd) elements from HBM (a block of queries and a block of key-value pairs). To maximise the block sizes subject to the SRAM constraint, we set Br=Bc=Θ(M/d) (so that Qi, Kj, Vj, and the intermediate block all fit). The total number of HBM accesses is then N2(M/d)2β‹…Mdβ‹…d=N2d2M. When Mβ‰₯Nd (which holds for typical sequence lengths on modern GPUs), this becomes O(Nd).

∎

Remark 12 (FlashAttention-2 and Beyond).

FlashAttention-2 [20] further improves performance by reducing non-matmul FLOPs (which are slow on modern GPUs whose tensor cores are optimised for matrix multiplications) and by parallelising across the sequence length dimension rather than only the batch and head dimensions. On an A100 GPU, FlashAttention-2 achieves up to 230,TFLOPs/s, reaching 50–73% of the theoretical peak. These advances make it feasible to train Transformers with context windows of 128k tokens or more without running into the memory wall, and Flash Attention has become a standard component in virtually all modern large-model training and inference frameworks.

Gated Attention

The scaled dot-product attention output (Definition 12) is a convex combination of value vectors: the softmax normalisation forces the attention weights to sum to one along each row, so every output row lies in the convex hull of {V1,…,Vn}. This is elegant, but it carries a subtle limitation. For fixed attention weights, the output is a linear function of V, and the sum-to-one constraint compels every query to β€œspend” its entire attention budget, even when the context contains mostly uninformative positions. Gated Attention [30] addresses both issues by applying a learnable element-wise gate to the attention output.

Definition 18 (Gated Attention).

Given the standard attention output Y=Attention(Q,K,V)βˆˆβ„nΓ—d⁑ and the sublayer input Xβˆˆβ„nΓ—d⁑, gated attention computes (Gated ATTN)Yβ€²=YβŠ™Οƒ(XWΞΈ), where WΞΈβˆˆβ„dΓ—d⁑ is a learned projection matrix (one per head) and Οƒ denotes the sigmoid function applied element-wise. The gate G=Οƒ(XWΞΈ)∈[0,1]nΓ—d acts as a soft binary mask on every dimension of the attention output.

Remark 13 (Why Gating Helps).

The gate provides three interconnected benefits:

  • Non-linearity. Standard attention is piecewise-linear in V: once the softmax weights are determined by Q and K, the output is merely a weighted sum. The sigmoid gate introduces a genuine non-linearity that breaks the low-rank bottleneck inherent in linear mixing.

  • Sparsity. The sigmoid function naturally pushes gate entries towards 0 or 1. In practice many entries saturate near zero, effectively zeroing out uninformative dimensions. This learnable sparsity is more flexible than fixed sparse attention patterns such as local windows or strided masks.

  • Attention-sink elimination. In standard transformers, the softmax constraint βˆ‘jΞ±ij=1 forces every query to distribute its attention mass somewhere. This often causes the first token, or a special separator token, to accumulate disproportionate attention mass as a β€œsink” [21]. The gated mechanism sidesteps this problem: even if the attention weights assign large mass to a particular position, the gate can suppress the result, removing the need for artificial sinks entirely.

Insight.

Gated Attention won the Best Paper award at NeurIPS 2025. It is a drop-in replacement for standard scaled dot-product attention in any transformer architecture: the only additional parameters are WΞΈβˆˆβ„dΓ—d⁑ per head, negligible compared to the existing WQ,WK,WV,WO projections. Experiments at the 15,B parameter scale showed consistent perplexity improvements on language modelling, improved few-shot accuracy, and faster convergence during pretraining. Several frontier model teams adopted gated attention within months of publication, making it one of the fastest-adopted architectural changes in the history of large language models.

Modern Transformer Components

The original Transformer block (Algorithm 2) used post-norm LayerNorm and a two-layer ReLU feedforward network. Over six years of scaling experiments, the community converged on a set of improved components that individually seem minor but collectively enable models an order of magnitude larger to train stably and efficiently. We catalogue the three most important changes below.

RMSNorm

Definition 19 (Root Mean Square Layer Normalisation).

RMSNorm [22] normalises a vector π’™βˆˆβ„d⁑ by its root mean square: (Rmsnorm)RMSNorm(𝒙)=𝒙RMS(𝒙)βŠ™πœΈ,RMS(𝒙)=1dβˆ‘i=1dxi2, where πœΈβˆˆβ„d⁑ is a learned element-wise scale parameter.

Remark 14 (RMSNorm vs. LayerNorm).

Unlike LayerNorm, RMSNorm omits the mean-centring step (π’™β†π’™βˆ’xβ€Ύ). Zhang and Sennrich [22] showed that this simplification preserves model quality while reducing normalisation computation by 10–15%, since computing the mean and the associated subtraction are no longer needed. RMSNorm is used in virtually all modern large language models, including LLaMA [14], Qwen [12], DeepSeek [15], and Mistral, and is always applied in pre-norm position (before the attention or FFN sublayer, not after).

SwiGLU Feedforward Networks

Definition 20 (SwiGLU).

The SwiGLU feedforward layer [23] replaces the standard two-matrix FFN with a gated variant using three weight matrices: (Swiglu)SwiGLU(𝒙)=(Swish(𝒙W1)βŠ™(𝒙V))W2, where W1,Vβˆˆβ„dΓ—dff⁑, W2βˆˆβ„dffΓ—d⁑, and Swish(z)=zβ‹…Οƒ(z) is the Swish activation (also known as SiLU). The element-wise product βŠ™ is the β€œgating” mechanism: the branch through V modulates (gates) the activated branch through W1.

Remark 15 (SwiGLU in Practice).

SwiGLU was proposed by Shazeer [23] as one of several GLU (Gated Linear Unit) variants. Because it uses three weight matrices instead of two, the hidden dimension is typically set to dff=⌊23β‹…4dβŒ‹, rounded to the nearest multiple of 256 for hardware efficiency, so that the total parameter count remains comparable to the original ReLU FFN. SwiGLU consistently outperforms both ReLU and GELU in language model perplexity at matched parameter counts. It is the default FFN in LLaMA [14], PaLM [24], Qwen [12], DeepSeek [15], and Mistral.

Parallel Attention and FFN

Definition 21 (Parallel Transformer Block).

Instead of the standard sequential block, where the FFN takes as input the output of the attention sublayer, the parallel formulation feeds both branches from the same normalised input and sums their outputs: (Parallel Block)𝒙′=𝒙+Attn(Norm(𝒙))+FFN(Norm(𝒙)). Compare the sequential form, where the FFN receives the attention-updated representation: 𝒙′=𝒙+FFN(Norm(𝒙+Attn(Norm(𝒙)))). In the parallel block, both the attention and FFN branches see the same normalised input Norm(𝒙).

Remark 16 (Why Parallel Blocks?).

The parallel formulation was popularised by GPT-J [25] and adopted in PaLM [24]. Its main advantage is reduced sequential computation: on hardware with sufficient parallelism (modern GPUs and TPUs), the attention and FFN kernels can execute concurrently, reducing wall-clock latency by up to 15%. The trade-off is a slight quality degradation at small model sizes, because the two sublayers no longer interact within a single block. However, this gap vanishes at scale (≳,8,B parameters), which is precisely the regime where the latency savings matter most.

ComponentOriginal (2017)Modern (2024+)
NormalisationLayerNorm (post-norm)RMSNorm (pre-norm)
FFN activationReLUSwiGLU
Positional encodingSinusoidal (additive)RoPE (multiplicative)
Attention headsMHA (G=h)GQA (Gβ‰ͺh)
Block structureSequentialParallel (optional)
Evolution of the Transformer block. Each row represents an independent design choice; modern LLMs typically adopt all five changes simultaneously.

Architecture Spotlight: Recent Large Language Models

The building blocks from the preceding subsections (RoPE, Rotary Position Embeddings (RoPE); grouped-query attention, Efficient Attention: MQA, GQA, and KV Caching; SwiGLU, SwiGLU Feedforward Networks; RMSNorm, RMSNorm; and optionally gated attention, Gated Attention) are not used in isolation. Frontier large language models combine them into carefully engineered stacks, and the differences between architectures often reduce to a handful of design choices. We briefly survey three representative architectures to illustrate how the components compose.

DeepSeek and Multi-Head Latent Attention

Grouped-query attention (Efficient Attention: MQA, GQA, and KV Caching) reduces the KV cache by sharing key–value heads across groups of query heads. DeepSeek-V2 takes a more aggressive approach: instead of sharing heads, it compresses the entire key–value representation through a low-rank bottleneck.

Definition 22 (Multi-Head Latent Attention (MLA)).

Let 𝒙tβˆˆβ„d⁑ be the hidden state at position t. Multi-head latent attention [15] first projects 𝒙t to a shared low-dimensional latent vector, then reconstructs the per-head keys and values from it: (MLA DOWN)𝒄tKV=WDKV𝒙t,π’Œt=WUK𝒄tKV,𝒗t=WUV𝒄tKV, where WDKVβˆˆβ„dΓ—dc⁑ projects to a latent dimension dcβ‰ͺd, and WUK,WUVβˆˆβ„dcΓ—dk⁑ reconstruct the per-head keys and values. Only the compressed vector 𝒄tKVβˆˆβ„dc⁑ needs to be cached during inference, a reduction factor of roughly d/dc compared to standard multi-head caching.

Remark 17 (MLA in Practice).

DeepSeek-V2 uses dc=512 with d=5120, achieving a roughly 10Γ— KV cache reduction relative to standard multi-head attention. A subtlety: RoPE cannot be absorbed into the latent because the position-dependent rotation would break the low-rank structure. DeepSeek solves this by applying RoPE to a small separate portion of each key (decoupled from the latent), preserving the relative-position property while still compressing the bulk of the KV representation.

Qwen 2.5

Qwen 2.5 [12] is a representative β€œrecipe” model that assembles standard best practices without introducing a single novel component. Its 72B-parameter variant stacks 80 Transformer layers with model dimension d=8192, grouped-query attention with 64 query heads and 8 KV heads (G=8), RoPE with a base frequency ΞΈbase=106 (tuned for 128K context), SwiGLU feedforward blocks with intermediate dimension dff=29568, and RMSNorm in the pre-norm position throughout. Context extension to 128K tokens is achieved via YaRN [13] interpolation of the RoPE frequencies. The model is trained on approximately 18 trillion tokens. The lesson is instructive: much of the power of modern LLMs comes from careful engineering at scale: the right combination of well-understood components, clean data, and long training, not from architectural novelty.

GLM and Bidirectional Attention

Most contemporary LLMs adopt a decoder-only (fully autoregressive) architecture. GLM [31] takes a different path: it partitions the input into a β€œcontext” prefix (Part A) and β€œblanks” to be filled (Part B). Part A uses bidirectional attention, where every token in A can attend to every other token in A, while Part B uses autoregressive attention, where each token attends to all of A and to the preceding tokens in B. A two-dimensional positional encoding distinguishes inter-part and intra-part positions. This design unifies the strengths of BERT-style bidirectional understanding and GPT-style sequential generation within a single model, and it forms the backbone of the ChatGLM family of conversational models.

Remark 18 (Mixture-of-Experts).

Both DeepSeek-V2 and Qwen 2.5 offer Mixture-of-Experts (MoE) variants. In an MoE Transformer, the dense SwiGLU feedforward layer is replaced by E β€œexpert” feedforward sub-networks, and a lightweight learned router selects the top-k experts for each token. This decouples model capacity (total parameter count) from model cost (compute per token): the network can store far more knowledge in its experts without proportionally increasing the FLOPs required at inference time. DeepSeekMoE further innovates with fine-grained expert segmentation (splitting standard-sized experts into many smaller ones for more flexible routing) and shared β€œalways-on” experts that capture common knowledge activated for every token regardless of the router's decision.

Historical Note.

The Transformer was designed in 2017 for a single task: English-to-German machine translation. Within seven years, the same architecture, with the refinements surveyed in this section, became the backbone of models with trillions of parameters, powering not just language tasks but image generation (Stable Diffusion 3, DALL-E 3), video synthesis (Sora), music (Jukebox), protein structure prediction (AlphaFold 2), and weather forecasting (GenCast). Few architectures in the history of computing have proved so universally applicable. The lesson, perhaps, is that a simple, scalable primitive, computing pairwise interactions in parallel, can be adapted to almost any domain where data has structure.

The Complete Modern Transformer

We close The Transformer with a comprehensive architectural diagram that assembles every component discussed in this section into a single modern decoder-only Transformer block, as found in models such as LLaMA, Qwen, and DeepSeek.

A complete modern decoder-only Transformer block as used in LLaMA, Qwen, and similar models. The dashed box is repeated N times (e.g. N=80 for LLaMA-70B). Pre-norm RMSNorm precedes both the attention and feedforward sub-layers; residual connections (grey paths via βŠ•οΈŽ) bypass each sub-layer. Compare with the original Transformer block in Algorithm 2.

The block just described is the fundamental computational unit in every architecture listed in the summary table of Architectures for Generative Models. In Parts II–V we will see how this block is adapted to specific generative model families: autoregressive language models stack it as a causal decoder, diffusion models interleave it with U-Net residual blocks and cross-attention, and vision transformers partition images into patch sequences to feed through the same architecture.

Architectures for Generative Models

We close this chapter by connecting the building blocks to the specific generative model families studied in Parts II–V.

Model FamilyArchitectureKey Blocks
GAN generatorCNN (upsample)Transposed conv, BN, ReLU/Tanh
GAN discriminatorCNN (downsample)Strided conv, BN/SN, LeakyReLU
VAE encoderCNN (downsample)Strided conv, BN, ReLU
VAE decoderCNN (upsample)Transposed conv, BN, ReLU/Sigmoid
Autoregressive (text)Transformer decoderMasked self-attn, FFN, LN
Autoregressive (image)Masked CNNMasked conv, gated activation
Diffusion denoiserU-Net + attentionResBlocks, skip conn, cross-attn, LN
Normalising flowInvertible networkCoupling layers, 1Γ—1 conv
How each generative model family uses neural network building blocks. β€œSN” = spectral normalisation; β€œLN” = layer normalisation.

Key Idea.

Every generative model is, at its core, a clever arrangement of the building blocks introduced in this chapter. Understanding what each block does, and why it is used in a particular model, will make the detailed treatments in Parts II–V much more accessible. When you encounter a new generative model, ask yourself: What is the architecture of gΞΈ? What is it trained to predict? What loss function (divergence) is being minimised?

Exercises

Exercise 1 (Forward Pass).

Consider a 2-layer feedforward network with input 𝒙=[1,βˆ’1]⊀, weight matrices W(1)=[2βˆ’111], W(2)=[1,βˆ’1], zero biases, and ReLU activation. Compute the output by performing the full forward pass. Then compute βˆ‚y/βˆ‚W(2) using backpropagation.

Exercise 2 (XOR and Depth).

Show that no single-layer perceptron can compute XOR (f(x1,x2)=x1βŠ•οΈŽx2) by proving that the four input-output pairs are not linearly separable. Construct a two-layer network with ReLU activation that computes XOR exactly.

Exercise 3 (Universal Approximation and Depth).

The universal approximation theorem guarantees that a single hidden layer suffices to approximate any continuous function. Explain why, despite this, deep networks are preferred in practice. Give an example of a function class that requires exponential width for a shallow network but only polynomial width for a deep network. Hint: consider parity functions or polynomials.

Exercise 4 (Batch Norm vs. Layer Norm).

For a hidden activation tensor of shape (B,F) with B=4 samples and F=3 features: [123456789101112], compute the normalised output for both batch normalisation and layer normalisation (with Ξ³=1, Ξ²=0, Ο΅=0). Explain why layer normalisation is preferred in Transformers.

Exercise 5 (Convolution Arithmetic).

A 2D convolution with Cin=3, Cout=64, kernel size 3Γ—3, stride 2, and no bias is applied to a 64Γ—64 input. (a) How many parameters does this layer have? (b) What is the spatial size of the output? (c) How many parameters would a fully connected layer require for the same input/output sizes?

Exercise 6 (Transposed Convolution Output Size).

A transposed convolution with kernel size 4Γ—4, stride 2, and padding 1 is applied to a 16Γ—16 feature map. What is the output size? Explain how this is used in a GAN generator to progressively upsample from a small latent map to a full image.

Exercise 7 (Attention Complexity).

Show that the time and memory complexity of computing scaled dot-product attention ((Attention)) for a sequence of length n with embedding dimension d is O(n2d). Why does this become a bottleneck for long sequences? Name one approach that reduces this to O(nnd) or better.

Exercise 8 (Causal Mask and Autoregressive Sampling).

Explain why removing the causal mask from a GPT-style model during training would make it useless for autoregressive generation at inference time. Specifically, what information leakage occurs, and why can't the model generate tokens one-by-one without the mask?

Exercise 9 (Residual Gradients).

Consider a chain of L residual blocks: 𝒂(β„“)=𝒂(β„“βˆ’1)+β„±β„“(𝒂(β„“βˆ’1)). Show that βˆ‚π’‚(L)βˆ‚π’‚(0)=I+βˆ‘(crossΒ terms), and explain why the identity term I prevents gradient vanishing regardless of the depth L.

Exercise 10 (U-Net for Diffusion).

In a diffusion model, the U-Net takes noisy image xt and time step t as input and predicts the noise Ο΅. Explain: (a) why skip connections are important for this task, (b) how the time step t is typically injected into the network, and (c) how text conditioning is incorporated via cross-attention.

Exercise 11 (Scaling Factor in Attention).

In the self-attention formula, the dot products are scaled by 1/dk. Assume Q and K have entries drawn i.i.d. from 𝒩(0,1). Show that the variance of each entry of QK⊀ is dk. Explain why this motivates the 1/dk scaling and what happens to the softmax gradients without it.

Exercise 12 (RoPE Rotation Matrix).

Consider a model with head dimension dk=4 and RoPE base frequency ΞΈbase=10000. Compute the four rotation frequencies ΞΈ1,ΞΈ2 and write out the full 4Γ—4 rotation matrix RΘm for position m=3. Verify that RΘ3⊀RΘ5=RΘ2.

Exercise 13 (KV Cache Memory).

A Transformer language model has L=32 layers, h=32 attention heads, head dimension dk=128, and serves sequences of length n=8192 in FP16 (2 bytes per value). (a) Compute the total KV cache size in GB. (b) If the model switches to GQA with G=4 groups, what is the new cache size? (c) If it further adopts MLA with latent dimension dc=512 (replacing all per-head KV), what is the cache size?

Exercise 14 (Flash Attention IO Complexity).

Consider standard attention with sequence length N=4096 and head dimension d=128. (a) Compute the number of elements in the full NΓ—N attention matrix. (b) If the GPU SRAM can hold M=105 elements, compute the Flash Attention block sizes Br=Bc=⌊M/(4d)βŒ‹ and the number of block iterations TrΓ—Tc. (c) Compare the total HBM reads/writes (in number of elements) for standard vs. Flash Attention.

References

  1. A Logical Calculus of Ideas Immanent in Nervous Activity

    Warren McCulloch, Walter Pitts

    The Bulletin of Mathematical Biophysics, vol. 5, no. 4, pp. 115-133 Β· 1943

    BibTeX 
    @article{mcculloch1943logical,
  title={A Logical Calculus of Ideas Immanent in Nervous Activity},
  author={McCulloch, Warren and Pitts, Walter},
  journal={The Bulletin of Mathematical Biophysics},
  volume={5},
  number={4},
  pages={115--133},
  year={1943},
  publisher={Springer}
}

    Journal article

  2. The Perceptron: A Probabilistic Model for Information Storage and Organization in the Brain

    Frank Rosenblatt

    Psychological Review, vol. 65, no. 6, pp. 386-408 Β· 1958

    BibTeX 
    @article{rosenblatt1958perceptron,
  title={The Perceptron: A Probabilistic Model for Information Storage and Organization in the Brain},
  author={Rosenblatt, Frank},
  journal={Psychological Review},
  volume={65},
  number={6},
  pages={386--408},
  year={1958},
  publisher={American Psychological Association}
}

    Journal article

  3. Perceptrons: An Introduction to Computational Geometry

    Marvin Minsky, Seymour Papert

    MIT Press Β· 1969

    BibTeX 
    @book{minsky1969perceptrons,
  title={Perceptrons: An Introduction to Computational Geometry},
  author={Minsky, Marvin and Papert, Seymour},
  year={1969},
  publisher={MIT Press}
}

    Book

  4. Learning representations by back-propagating errors

    David E. Rumelhart, Geoffrey E. Hinton, Ronald J. Williams

    Nature, vol. 323, no. 6088, pp. 533-536 Β· 1986

    BibTeX 
    @article{rumelhart1986learning,
  title={Learning representations by back-propagating errors},
  author={Rumelhart, David E. and Hinton, Geoffrey E. and Williams, Ronald J.},
  journal={Nature},
  volume={323},
  number={6088},
  pages={533--536},
  year={1986},
  publisher={Nature Publishing Group}
}

    Journal article

  5. Deep Learning

    Yann LeCun, Yoshua Bengio, Geoffrey Hinton

    Nature, vol. 521, no. 7553, pp. 436-444 Β· 2015

    BibTeX 
    @article{lecun2015deep,
  title={Deep Learning},
  author={LeCun, Yann and Bengio, Yoshua and Hinton, Geoffrey},
  journal={Nature},
  volume={521},
  number={7553},
  pages={436--444},
  year={2015},
  publisher={Nature Publishing Group}
}

    Journal article

  6. Attention is all you need

    Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, Illia Polosukhin

    Advances in neural information processing systems, pp. 5998-6008 Β· 2017

    BibTeX 
    @inproceedings{vaswani2017attention,
title={Attention is all you need},
author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, {\L}ukasz and Polosukhin, Illia},
booktitle={Advances in neural information processing systems},
pages={5998--6008},
year={2017}
}

    Conference paper

  7. Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks

    Alec Radford, Luke Metz, Soumith Chintala

    International Conference on Learning Representations Β· 2016

    BibTeX 
    @inproceedings{radford2015unsupervised,
  title={Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks},
  author={Radford, Alec and Metz, Luke and Chintala, Soumith},
  booktitle={International Conference on Learning Representations},
  year={2016}
}

    Conference paper

  8. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift

    Sergey Ioffe, Christian Szegedy

    International Conference on Machine Learning (ICML), pp. 448-456 Β· 2015

    BibTeX 
    @inproceedings{ioffe2015batch,
  title={Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift},
  author={Ioffe, Sergey and Szegedy, Christian},
  booktitle={International Conference on Machine Learning (ICML)},
  pages={448--456},
  year={2015}
}

    Conference paper

  9. Layer Normalization

    Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E Hinton

    arXiv preprint arXiv:1607.06450 Β· 2016

    BibTeX 
    @article{ba2016layer,
  title={Layer Normalization},
  author={Ba, Jimmy Lei and Kiros, Jamie Ryan and Hinton, Geoffrey E},
  journal={arXiv preprint arXiv:1607.06450},
  year={2016}
}

    Journal article

  10. A Style-Based Generator Architecture for Generative Adversarial Networks

    Tero Karras, Samuli Laine, Timo Aila

    Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 4401-4410 Β· 2019

    BibTeX 
    @inproceedings{karras2019style,
  title={A Style-Based Generator Architecture for Generative Adversarial Networks},
  author={Karras, Tero and Laine, Samuli and Aila, Timo},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={4401--4410},
  year={2019}
}

    Conference paper

  11. Analyzing and Improving the Image Quality of StyleGAN

    Tero Karras, Samuli Laine, Miika Aittala, Janne Hellsten, Jaakko Lehtinen, Timo Aila

    Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 8110-8119 Β· 2020

    BibTeX 
    @inproceedings{karras2020analyzing,
  title={Analyzing and Improving the Image Quality of {StyleGAN}},
  author={Karras, Tero and Laine, Samuli and Aittala, Miika and Hellsten, Janne and Lehtinen, Jaakko and Aila, Timo},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={8110--8119},
  year={2020}
}

    Conference paper

  12. Qwen2.5 Technical Report

    Qwen Team

    arXiv preprint arXiv:2412.15115 Β· 2024

    BibTeX 
    @article{qwen25,
  title={{Qwen2.5} Technical Report},
  author={{Qwen Team}},
  journal={arXiv preprint arXiv:2412.15115},
  year={2024}
}

    Journal article

  13. YaRN: Efficient Context Window Extension of Large Language Models

    Bowen Peng, Jeffrey Quesnelle, Honglu Fan, Enrico Shippole

    arXiv preprint arXiv:2309.00071 Β· 2024

    BibTeX 
    @article{peng2024yarn,
  title={{YaRN}: Efficient Context Window Extension of Large Language Models},
  author={Peng, Bowen and Quesnelle, Jeffrey and Fan, Honglu and Shippole, Enrico},
  journal={arXiv preprint arXiv:2309.00071},
  year={2024}
}

    Journal article

  14. LLaMA: Open and Efficient Foundation Language Models

    Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, others

    arXiv preprint arXiv:2302.13971 Β· 2023

    BibTeX 
    @article{touvron2023llama,
  author    = {Touvron, Hugo and Lavril, Thibaut and Izacard, Gautier and Martinet, Xavier and Lachaux, Marie-Anne and others},
  title     = {{LLaMA}: Open and Efficient Foundation Language Models},
  journal   = {arXiv preprint arXiv:2302.13971},
  year      = {2023}
}

    Journal article

  15. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

    DeepSeek-AI

    arXiv preprint arXiv:2405.04434 Β· 2024

    BibTeX 
    @article{deepseekv2,
  title={{DeepSeek-V2}: A Strong, Economical, and Efficient Mixture-of-Experts Language Model},
  author={{DeepSeek-AI}},
  journal={arXiv preprint arXiv:2405.04434},
  year={2024}
}

    Journal article

  16. Mistral 7B

    Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, et al.

    arXiv preprint arXiv:2310.06825 Β· 2023

    BibTeX 
    @article{jiang2023mistral,
  title={{Mistral 7B}},
  author={Jiang, Albert Q. and Sablayrolles, Alexandre and Mensch, Arthur and Bamford, Chris and Chaplot, Devendra Singh and de las Casas, Diego and Bressand, Florian and Lengyel, Gianna and Lample, Guillaume and Saulnier, Lucile and others},
  journal={arXiv preprint arXiv:2310.06825},
  year={2023}
}

    Journal article

  17. Scalable Diffusion Models with Transformers

    William Peebles, Saining Xie

    Proceedings of the IEEE/CVF International Conference on Computer Vision Β· 2023

    BibTeX 
    @inproceedings{peebles2023scalable,
  title={Scalable Diffusion Models with Transformers},
  author={Peebles, William and Xie, Saining},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
  year={2023}
}

    Conference paper

  18. Fast Transformer Decoding: One Write-Head is All You Need

    Noam Shazeer

    arXiv preprint arXiv:1911.02150 Β· 2019

    BibTeX 
    @article{shazeer2019fast,
  title={Fast Transformer Decoding: One Write-Head is All You Need},
  author={Shazeer, Noam},
  journal={arXiv preprint arXiv:1911.02150},
  year={2019}
}

    Journal article

  19. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

    Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, Sumit Sanghai

    Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing Β· 2023

    BibTeX 
    @inproceedings{ainslie2023gqa,
  title={{GQA}: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints},
  author={Ainslie, Joshua and Lee-Thorp, James and de Jong, Michiel and Zemlyanskiy, Yury and Lebr{\'o}n, Federico and Sanghai, Sumit},
  booktitle={Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing},
  year={2023}
}

    Conference paper

  20. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning

    Tri Dao

    International Conference on Learning Representations Β· 2024

    BibTeX 
    @inproceedings{dao2024flashattention2,
  title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},
  author={Dao, Tri},
  booktitle={International Conference on Learning Representations},
  year={2024}
}

    Conference paper

  21. Efficient Streaming Language Models with Attention Sinks

    Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis

    International Conference on Learning Representations Β· 2024

    BibTeX 
    @inproceedings{xiao2024efficient,
  title={Efficient Streaming Language Models with Attention Sinks},
  author={Xiao, Guangxuan and Tian, Yuandong and Chen, Beidi and Han, Song and Lewis, Mike},
  booktitle={International Conference on Learning Representations},
  year={2024}
}

    Conference paper

  22. Root Mean Square Layer Normalization

    Biao Zhang, Rico Sennrich

    Advances in Neural Information Processing Systems, vol. 32 Β· 2019

    BibTeX 
    @inproceedings{zhang2019rmsnorm,
  title={Root Mean Square Layer Normalization},
  author={Zhang, Biao and Sennrich, Rico},
  booktitle={Advances in Neural Information Processing Systems},
  volume={32},
  year={2019}
}

    Conference paper

  23. GLU Variants Improve Transformer

    Noam Shazeer

    arXiv preprint arXiv:2002.05202 Β· 2020

    BibTeX 
    @article{shazeer2020glu,
  title={{GLU} Variants Improve Transformer},
  author={Shazeer, Noam},
  journal={arXiv preprint arXiv:2002.05202},
  year={2020}
}

    Journal article

  24. PaLM: Scaling Language Modeling with Pathways

    Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, et al.

    Journal of Machine Learning Research, vol. 24, no. 240, pp. 1-113 Β· 2023

    BibTeX 
    @article{chowdhery2023palm,
  title={{PaLM}: Scaling Language Modeling with Pathways},
  author={Chowdhery, Aakanksha and Narang, Sharan and Devlin, Jacob and Bosma, Maarten and Mishra, Gaurav and Roberts, Adam and Barham, Paul and Chung, Hyung Won and Sutton, Charles and Gehrmann, Sebastian and others},
  journal={Journal of Machine Learning Research},
  volume={24},
  number={240},
  pages={1--113},
  year={2023}
}

    Journal article

  25. GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model

    Ben Wang, Aran Komatsuzaki

    2021

    BibTeX 
    @article{wang2021gptj,
  title={{GPT-J-6B}: A 6 Billion Parameter Autoregressive Language Model},
  author={Wang, Ben and Komatsuzaki, Aran},
  year={2021},
  note={\url{https://github.com/kingoflolz/mesh-transformer-jax}}
}

    Journal article

  26. Gradient-based learning applied to document recognition

    Yann LeCun, LΓ©on Bottou, Yoshua Bengio, Patrick Haffner

    Proceedings of the IEEE, vol. 86, no. 11, pp. 2278-2324 Β· 1998

    BibTeX 
    @article{lecun1998gradient,
  title={Gradient-based learning applied to document recognition},
  author={LeCun, Yann and Bottou, LΓ©on and Bengio, Yoshua and Haffner, Patrick},
  journal={Proceedings of the IEEE},
  volume={86},
  number={11},
  pages={2278--2324},
  year={1998},
  publisher={IEEE}
}

    Journal article

  27. Deep Residual Learning for Image Recognition

    Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian Sun

    IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778 Β· 2016

    BibTeX 
    @inproceedings{he2016deep,
  title={Deep Residual Learning for Image Recognition},
  author={He, Kaiming and Zhang, Xiangyu and Ren, Shaoqing and Sun, Jian},
  booktitle={IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
  pages={770--778},
  year={2016}
}

    Conference paper

  28. RoFormer: Enhanced Transformer with Rotary Position Embedding

    Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, Yunfeng Liu

    Neurocomputing, vol. 568, pp. 127063 Β· 2024

    BibTeX 
    @article{su2024roformer,
  title={{RoFormer}: Enhanced Transformer with Rotary Position Embedding},
  author={Su, Jianlin and Ahmed, Murtadha and Lu, Yu and Pan, Shengfeng and Bo, Wen and Liu, Yunfeng},
  journal={Neurocomputing},
  volume={568},
  pages={127063},
  year={2024},
  publisher={Elsevier}
}

    Journal article

  29. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

    Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Re

    Advances in Neural Information Processing Systems, vol. 35 Β· 2022

    BibTeX 
    @inproceedings{dao2022flashattention,
  title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
  author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
  booktitle={Advances in Neural Information Processing Systems},
  volume={35},
  year={2022}
}

    Conference paper

  30. Gated Attention for Large Language Models

    Zihan Qiu, Zekun Wang, Bo Zheng, Zeyu Huang, Kaiyue Wen, Songlin Yang, Rui Men, Le Yu, et al.

    Advances in Neural Information Processing Systems, vol. 38 Β· 2025

    BibTeX 
    @inproceedings{qiu2025gated,
  title={Gated Attention for Large Language Models},
  author={Qiu, Zihan and Wang, Zekun and Zheng, Bo and Huang, Zeyu and Wen, Kaiyue and Yang, Songlin and Men, Rui and Yu, Le and Huang, Fei and Huang, Suozhi and Liu, Dayiheng and Zhou, Jingren and Lin, Junyang},
  booktitle={Advances in Neural Information Processing Systems},
  volume={38},
  year={2025}
}

    Conference paper

  31. GLM: General Language Model Pretraining with Autoregressive Blank Infilling

    Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, Jie Tang

    Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics Β· 2022

    BibTeX 
    @inproceedings{du2022glm,
  title={{GLM}: General Language Model Pretraining with Autoregressive Blank Infilling},
  author={Du, Zhengxiao and Qian, Yujie and Liu, Xiao and Ding, Ming and Qiu, Jiezhong and Yang, Zhilin and Tang, Jie},
  booktitle={Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics},
  year={2022}
}

    Conference paper

← 4 Divergence Measures and Information Theory6 Probabilistic Graphical Models β†’