Forward versus Backward Pass: Backpropagation in Transformers

Training a modern language model fundamentally relies on a repeating cycle of two consecutive phases: the forward pass and the backward pass. Although an LLM during normal inference only runs the forward computation to predict the next token, every training and fine-tuning process requires the full loop. In the forward phase, tokenized vectors flow through dozens of transformer layers to arrive at a probability distribution over the vocabulary. A loss function then calculates the discrepancy between the prediction and the actual target text. Immediately afterward, the backward pass begins: via the chain rule of differential calculus (backpropagation), error gradients travel back through the network to determine which weights need to be adjusted.

In this article, we break down the mathematical foundations, the computational asymmetry, and the enormous memory pressure that backpropagation exerts on modern GPU clusters. We follow a concrete tensor flow through a transformer block and analyze why storing intermediate activations forms the biggest bottleneck in large-scale LLM training.

The forward pass: from embedding to logits

During the forward pass, the model transforms a discrete sequence of input tokens step by step into a sequence of continuous representations. Consider a batch of tokens with dimensions [B, S], where B represents the batch size and S the sequence length. The first operation is the embedding lookup, in which each token ID is converted into a vector of hidden dimension D (the model dimension $d_{\text{model}}$), supplemented with positional information. This yields the initial activation tensor $X_0 \in \mathbb{R}^{B \times S \times D}$.

From this point on, the tensor successively passes through $L$ identical transformer layers. Within each layer $l$, two consecutive main transformations take place, each wrapped by a residual connection (skip connection) and a normalization step (RMSNorm or LayerNorm):

First, the normalized input undergoes a multi-head self-attention computation. The tensor is projected via linear weight matrices $W_Q, W_K, W_V$ into Queries, Keys, and Values. To understand the exact matrix multiplications and the scaling factor of this step, you can study the mathematical derivation of self-attention. The output of the attention heads is concatenated, linearly projected via $W_O$, and added to the original residual stream:

# Wiskundige stappen binnen de voorwaartse pass van laag l
X_norm1 = RMSNorm(X_{l-1})
X_attn  = MultiHeadAttention(Q=X_norm1 * W_Q, K=X_norm1 * W_K, V=X_norm1 * W_V) * W_O
X_mid   = X_{l-1} + X_attn

X_norm2 = RMSNorm(X_mid)
X_ffn   = FeedForward(X_norm2)
X_l     = X_mid + X_ffn

After the attention layer comes the feed-forward network (FFN or MLP). In modern architectures such as Llama, this network consists of paired linear projections with a non-linear gating function. For a deeper dive into why classic activations are no longer used here, read the analysis of SwiGLU and GELU activation functions. After the final transformer layer $L$, a final normalization and a linear projection to the vocabulary follow via the embedding matrix $W_{\text{vocab}}$, resulting in logits $Z \in \mathbb{R}^{B \times S \times V}$.

The loss function and the start of the backward pass

The transition from the forward to the backward pass takes place at the loss function. For autoregressive language models, this is almost always the cross-entropy loss over the next token. Given the logits $Z$ and the true target $y \in \{1, \dots, V\}$, the softmax function calculates the predicted probability distribution $\hat{p}$:

# Kansverdeling via Softmax voor token op positie t
p_{t, i} = exp(z_{t, i}) / sum_{j=1}^V exp(z_{t, j})

# Cross-entropy verlies voor positie t
L_t = -log(p_{t, y_t})

# Totale gemiddelde loss over de batch en sequentie
Loss = (1 / (B * S)) * sum_{b=1}^B sum_{t=1}^S L_{b, t}

The backward pass begins by determining the partial derivative of the total loss with respect to the unnormalized logits $z_{t, i}$. A well-known mathematical property of the combination of softmax and cross-entropy is that this derivative reduces remarkably simply:

dL / dz_{t, i} = p_{t, i} - 1   (als i gelijk is aan doeltarget y_t)
dL / dz_{t, i} = p_{t, i}       (als i NIET gelijk is aan target y_t)

Oftewel in vectorvorm:
dL / dZ = (P - Y_onehot) / (B * S)

This gradient tensor $\frac{\partial \mathcal{L}}{\partial Z}$, with dimensions identical to the logits [B, S, V], forms the starting point of the backward wave. From here, the error vector travels layer by layer back to the beginning of the network.

Backpropagation via the chain rule through the transformer block

Backpropagation is, in essence, the systematic application of the multivariable chain rule from calculus. Consider a generic linear layer $Y = X \cdot W$, where $X \in \mathbb{R}^{N \times D_{\text{in}}}$ is the input activation and $W \in \mathbb{R}^{D_{\text{in}} \times D_{\text{out}}}$ is the weight matrix. During the backward pass, this layer receives the upstream gradient $\frac{\partial \mathcal{L}}{\partial Y}$. The layer must now perform two separate calculations:

Here we directly see the central mechanism of deep learning: to calculate the weight gradient $\frac{\partial \mathcal{L}}{\partial W}$, the system needs the original input activation $X$ that was generated during the forward pass. If $X$ has not been retained in GPU memory, the multiplication $X^T \cdot \frac{\partial \mathcal{L}}{\partial Y}$ cannot take place.

When the gradient flows through a residual connection $X_{l} = X_{l-1} + F(X_{l-1})$, the addition rule dictates that the gradient splits:

dL / dX_{l-1} = (dL / dX_l) + (dL / dF) * (dF / dX_{l-1})

The term $\frac{\partial \mathcal{L}}{\partial X_l}$ flows unimpeded via the residual stream directly through to earlier layers. This phenomenon prevents the notorious problem of vanishing gradients in deep architectures with hundreds of layers.

Activation memory: the hidden bottleneck

When training large models, VRAM memory usage does not depend only on the number of parameters, but especially on the retained intermediate activations. While weights and optimizer states have a fixed size, activation memory scales linearly with batch size, linearly with the number of layers, and quadratically (or linearly with FlashAttention) with sequence length.

For a standard transformer layer, the following tensors must be stored during the forward pass to make the backward pass possible:

• Input for LayerNorm / RMSNorm: Needed to determine the normalization gradient.
• Queries, Keys, and Values: Needed for the gradients of the projection weights $W_Q, W_K, W_V$.
• Attention matrix (before and after softmax): Needed to calculate the Value weight and Query/Key gradients ($S \times S$ elements per head).
• Input for the MLP layer: Needed for the gradients of the up-projection and gate matrices.
• Pre-activation states: The values before the SwiGLU or GELU operation, used to calculate the derivative of the activation function.

For a model with 70 billion parameters and a context length of 8,192 tokens, the memory needed for these intermediate activations quickly exceeds the memory of the model parameters themselves. This forces engineers to apply advanced techniques, such as Activation Checkpointing (also known as gradient checkpointing). Here, activations are discarded during the forward pass and recomputed locally during the backward pass, which costs about 30% extra compute time but reduces activation memory by as much as 70% to 80%.

Computational comparison: FLOPs and compute time

There is a rule of thumb within LLM engineering: the backward pass costs roughly twice as many FLOPs (floating-point operations) as the forward pass. This brings the total training cost per token to approximately $6N$ FLOPs, where $N$ represents the number of active model parameters ($2N$ for forward, $4N$ for backward).

Why does the backward pass cost exactly twice as much compute as the forward pass? As shown earlier with the linear layer $Y = X \cdot W$: in the forward pass there is only one matrix multiplication ($X \cdot W$). In the backward pass, however, two matrix multiplications must be performed: one to propagate the gradient to the input ($\frac{\partial \mathcal{L}}{\partial Y} \cdot W^T$) and one to determine the gradient for the parameters ($X^T \cdot \frac{\partial \mathcal{L}}{\partial Y}$).

In production systems, we see that the cost of these computational steps directly affects infrastructure choices; anyone who wants to see how this computational split affects server costs and throughput speeds can read the analysis of batch versus real-time processing is worth consulting.

Step-by-step gradient flow in a PyTorch-like implementation

To visualize how a transformer block computationally switches between the forward and backward flow, below is a modular representation in Python pseudocode:

import torch

class TransformerBlockBackprop:
    def __init__(self, d_model):
        # Initialisatie van projectiegewichten
        self.W_q = torch.randn(d_model, d_model, requires_grad=True)
        self.W_k = torch.randn(d_model, d_model, requires_grad=True)
        self.W_v = torch.randn(d_model, d_model, requires_grad=True)
        self.W_out = torch.randn(d_model, d_model, requires_grad=True)

    def forward(self, x):
        # Opslaan van invoeractivatie voor gebruik in backward pass
        self.saved_x = x
        
        # 1. Lineaire projecties
        self.q = torch.matmul(x, self.W_q)
        self.k = torch.matmul(x, self.W_k)
        self.v = torch.matmul(x, self.W_v)
        
        # 2. Scaled Dot-Product Attention
        d_k = self.q.shape[-1]
        scores = torch.matmul(self.q, self.k.transpose(-2, -1)) / (d_k ** 0.5)
        self.attn_weights = torch.softmax(scores, dim=-1)
        self.context = torch.matmul(self.attn_weights, self.v)
        
        # 3. Output projectie en residuale optelling
        out = torch.matmul(self.context, self.W_out)
        return x + out

    def backward(self, grad_output):
        # grad_output is dL/d(out), stroomopwaarts ontvangen
        
        # Gradiënt door de residuale verbinding splitst lineair
        grad_residual = grad_output.clone()
        
        # Gradiënten voor W_out en de context-tensor
        grad_W_out = torch.matmul(self.context.transpose(-2, -1), grad_output)
        grad_context = torch.matmul(grad_output, self.W_out.t())
        
        # Gradiënten door attention softmax en V-matrix
        grad_v = torch.matmul(self.attn_weights.transpose(-2, -1), grad_context)
        grad_attn_weights = torch.matmul(grad_context, self.v.transpose(-2, -1))
        
        # Softmax backward transformatie
        # (vereenvoudigde representatie van de Jacobiaan-vermenigvuldiging)
        s = self.attn_weights
        grad_scores = s * (grad_attn_weights - (grad_attn_weights * s).sum(dim=-1, keepdim=True))
        grad_scores = grad_scores / (self.q.shape[-1] ** 0.5)
        
        # Gradiënten voor Q en K projecties
        grad_q = torch.matmul(grad_scores, self.k)
        grad_k = torch.matmul(grad_scores.transpose(-2, -1), self.q)
        
        # Gewichtsgradiënten berekenen met opgeslagen invoeractivatie self.saved_x
        grad_W_q = torch.matmul(self.saved_x.transpose(-2, -1), grad_q)
        grad_W_k = torch.matmul(self.saved_x.transpose(-2, -1), grad_k)
        grad_W_v = torch.matmul(self.saved_x.transpose(-2, -1), grad_v)
        
        # Totale activeringsgradiënt naar vorige laag propageren
        grad_x = grad_residual + (
            torch.matmul(grad_q, self.W_q.t()) +
            torch.matmul(grad_k, self.W_k.t()) +
            torch.matmul(grad_v, self.W_v.t())
        )
        
        return grad_x, (grad_W_q, grad_W_k, grad_W_v, grad_W_out)

In the code example above, it is visible that `self.saved_x` must be explicitly retained until the `backward` function is called. Once the model processes millions of tokens per iteration across dozens of layers, this forms the root cause of out-of-memory (OOM) errors on training hardware.

Challenges and numerical instability in backpropagation

During the backward pass, gradients move across hundreds of matrix operations. This carries significant technical and numerical risks:

1. Gradient Explosion (Exploding Gradients): When matrix weights are larger than 1, or with repeated accumulation in the residual stream, gradients can grow exponentially as they return to the first layers. This leads to NaN-values (Not a Number) in the weights. The standard remedy in transformer training is Gradient Clipping, in which the global norm of all gradient vectors is capped at a fixed threshold value (for example, 1.0):

if ||g|| > max_norm:
    g = g * (max_norm / ||g||)

2. Underflow in mixed-precision training (FP16/BF16): Gradients are often extremely small (for example, $10^{-6}$ or smaller). In 16-bit floating point representations (particularly standard IEEE FP16 with only 5 exponent bits), this quickly results in underflow to zero. This causes the model to stop learning. BFloat16 largely solves this by using 8 exponent bits (at the cost of precision in the mantissa), while FP16 requires advanced Loss Scalingalgorithms that multiply the loss by a factor before backpropagation and then scale the calculated gradients back down afterward.

3. Asynchronous communication in Distributed Data Parallel (DDP): When training across multiple GPUs, gradient synchronization (AllReduce) starts asynchronously as soon as an individual layer completes its backward pass. If the compute time of the backward pass does not perfectly overlap with the network bandwidth between GPUs (NVLink/Infiniband), significant GPU idle time results.

Optimization strategies for the backward pass

To curb the immense memory pressure and compute load of the backward pass, the AI field has achieved several breakthroughs:

• FlashAttention: By cleverly recomputing intermediate steps in the GPU's fast SRAM memory, the gigantic $S \times S$ attention matrix never has to be written to the slower HBM memory, either in the forward pass or the backward pass.
• ZeRO (Zero Redundancy Optimizer): Splits optimizer states, gradients, and model parameters across all participating GPUs, drastically reducing the memory overhead of the backward pass.
• Reversible Layers: Architectures in which intermediate activations can be mathematically reconstructed exactly from the output of the next layer, meaning that in theory zero activations need to be stored during the forward pass.