How Backpropagation Works: Tracing the Reverse Cascade of Signals
After completing this topic
You will have the answer to the question you postponed in Part 3: "How is the gradient calculated for each of the 33.62 million parameters in a practical amount of time?" The algorithm that makes this possible is backpropagation, and the underlying mathematics is the chain rule.
This part is a companion to Part 3. If Part 3 dealt with "how much to move," this part deals with "how to know which direction to move." Parts 5 (Transformers) and 12 (PyTorch in Practice) will build on what is presented here.
The Crushing Failure of the Naive Approach
In Part 3, we defined the partial derivative of a parameter ΞΈ_i as follows:
βL/βΞΈ_i = lim_{Ξ΅ β 0} [L(ΞΈ_1, ..., ΞΈ_i + Ξ΅, ..., ΞΈ_P) - L(ΞΈ)] / Ξ΅Let's try implementing this directly, as it appears in textbooks. For each parameter, nudge it slightly by Ξ΅, observe how much the loss changes, and divide by Ξ΅ to get the partial derivative. This is called the finite difference method.
The problem: If there are 33.62 million parameters, you need to perform a full forward pass of the neural network to obtain a single partial derivative. To obtain all 33.62 million partial derivatives, you need to perform 33.62 million forward passes of the neural network.
Let's get a sense of scale with a pathological example of a neural network. If a single forward pass takes 1 millisecond on a GPU, then calculating the gradient for one step would take 33.62 million milliseconds = about 9.3 hours. And training requires repeating this step tens of thousands or even millions of times. With this approach, training would take hundreds of years.
That doesn't mean the finite difference method is completely useless. As we'll see later, it is used as a verification tool for backpropagation implementations. When asking "Is our backpropagation code actually correct?", the finite difference method can be used to calculate the partial derivatives of a few parameters and compare them to the "ground truth." This is called gradient checking, and it is a crucial verification step that neural network library developers perform.
However, for actual training, a much faster method is needed. That method is backpropagation.
An Aberrant Signal Cascade: Where Does the Problem Lie?
To grasp the intuition behind backpropagation, let's consider a laboratory scenario.
Suppose you are conducting a cell signaling experiment. When you treat cells with a specific ligand, you expect downstream gene expression to be induced. The signal pathway looks something like this:
Ligand (external stimulus)
β
Receptor activation
β
Secondary signal (e.g., cAMP)
β
Kinase cascade (RAF β MEK β ERK, etc.)
β
Transcription factor phosphorylation
β
Induction of gene expression (measured outcome)You treat the cells with the ligand and measure the downstream gene expression, but it turns out to be 50% lower than expected. Your question: "Where did the problem occur?"
Perhaps the receptor is not being activated sufficiently, or perhaps the secondary signal is not being produced effectively, or perhaps the activity of a specific kinase is low, or perhaps the transcription factor is not being phosphorylated. To find the cause, you need to quantify how much each upstream node "contributes" to the 50% error downstream.
How would you approach this problem? You would quantitatively measure each step and trace it back along the relationships.
- Gene expression is 50% lower β What was the level of transcription factor phosphorylation? 30% lower β 30% contribution at the transcription factor stage, and 20% contribution at the mapping stage below (transcription factor β expression)
- Transcription factor phosphorylation is 30% lower β What is the upstream ERK activity? 15% lower β 15% contribution at the ERK stage
- ERK activity is 15% lower β MEK? 8% lower β 8% contribution from MEK
- ... continue upstream
At each step, you calculate "how much did this step contribute to the downstream anomaly." The method is as follows:
Downstream known contribution Γ Sensitivity of this step to downstream = Contribution of this step
Multiplying these two factors and propagating them upstream is exactly what backpropagation does. This is formalized mathematically as the chain rule.
Important observation: Before performing this backward tracing, you must first pass through the forward direction (ligand β expression) to know the normal values (baseline) of each step. In other words, one forward measurement + one backward tracing yields the contribution of each upstream node. This is why backpropagation is so much faster than the finite difference method (which requires repeated forward passes for each parameter).
The Chain Rule: The Calculus of Tracing Backwards
Let's understand the chain rule, the backbone of backpropagation, with the minimum amount of mathematics.
Suppose we have a composite function. When x is input, y = f(x) is first calculated, and then z = g(y). Finally, z = g(f(x)).
We want to know how much z changes when x changes slightly. In other words, we want to calculate dz/dx.
We could calculate this directly, but the chain rule breaks it down into two steps:
dz/dx = (dz/dy) Γ (dy/dx)In words: "How much does x change y?" multiplied by "How much does y change z?" gives "How much does x change z?"
Back to the bio-cascade analogy: "How much does the ligand concentration change gene expression?" is the product of several steps:
Ξ(expression)/Ξ(ligand) =
Ξ(expression)/Ξ(transcription_factor) Γ Ξ(transcription_factor)/Ξ(ERK) Γ Ξ(ERK)/Ξ(MEK) Γ ... Γ Ξ(receptor)/Ξ(ligand)At each arrow, the "sensitivity from this step to the next step" is multiplied, and finally, the overall sensitivity from ligand to expression is obtained.
Neural networks are exactly the same. The loss function L is a function of the output of the last layer, the output of the last layer is a function of the output of the previous layer, and so on, until it is a function of the parameters of the first layer.
The partial derivative of L with respect to the parameters of the first layer is, by the chain rule, the product of the sensitivities at each layer. Backpropagation is the process of calculating this product in reverse.
The Skeleton of the Backpropagation Algorithm
Now let's organize the algorithm. Suppose we have a neural network with L layers.
Step 1: Forward Pass
Pass the neural network from input to output to calculate and store the activation value of each layer.
a^(0) = input
a^(1) = Ο(W^(1) Β· a^(0) + b^(1))
a^(2) = Ο(W^(2) Β· a^(1) + b^(2))
...
a^(L) = Ο(W^(L) Β· a^(L-1) + b^(L))a^(l) is the activation vector of layer l, W^(l)Β·b^(l) is the weight matrixΒ·bias vector of layer l, and Ο is the activation function. The diagram from Part 2.
Here's the important part: store the a^(l) values for each layer in memory. This is needed in the backpropagation step. This is why training uses much more GPU memory than inference. Inference only requires a forward pass, but training must hold onto the activation values of all layers until the backpropagation step.
Step 2: Loss Calculation
Calculate the loss using the output of the last layer a^(L) and the ground truth y.
L = CrossEntropy(a^(L), y)This yields the partial derivative of the loss with respect to the activation value of the last layer: Ξ΄^(L) = βL/βa^(L).
In the case of the Softmax + Cross-Entropy combination, this partial derivative is surprisingly elegant:
Ξ΄^(L) = a^(L) - yIn other words, "predicted probability distribution minus ground truth one-hot vector." This elegance is why Softmax + CE are used together. Derived in Appendix A.4.
Step 3: Backpropagation
Starting from Ξ΄^(L), go through the layers in reverse and calculate Ξ΄^(l) for each layer.
Ξ΄^(l) = (W^(l+1))^T Β· Ξ΄^(l+1) β Ο'(z^(l))β is element-wise multiplication, and z^(l) = W^(l) Β· a^(l-1) + b^(l) is the value before the activation function. Ο' is the derivative of the activation function.
What does this equation do?
(W^(l+1))^T Β· Ξ΄^(l+1): This projects the error signal from the next layer back into the activation space of this layer. It passes the weights used to send the signal to the next layer in the reverse direction.β Ο'(z^(l)): This multiplies by the local gradient of the activation function at this layer. If the activation function is "dead" (e.g., a ReLU neuron wherez < 0), then Ο' is 0, and the error is blocked.
Repeat this from layer L to layer 1 to obtain Ξ΄^(l) for all layers.
Step 4: Calculate Parameter Partial Derivatives
Once Ξ΄^(l) is obtained for each layer, the partial derivatives with respect to the weights and biases of that layer are immediately obtained.
βL/βW^(l) = Ξ΄^(l) Β· (a^(l-1))^T
βL/βb^(l) = Ξ΄^(l)These are the actual gradients that were discussed in Part 3. Now, the optimizer (e.g., Adam) in Part 3 uses this gradient to update the parameters.
Why is it so fast?
The amazing thing about backpropagation is that the partial derivatives for all parameters are obtained with a single forward pass and a single backward pass. This is in contrast to the finite difference method, which requires repeated forward passes for each parameter.
Let's get a sense of scale. The computational complexity of backpropagation is approximately 2-3 times that of a forward pass. So, one training step for a neural network with 33.62 million parameters takes about 3-4 times the time of a forward pass. Compared to the 33.62 million forward passes required by the finite difference method, this is about 8 million times faster.
Without this algorithm, modern deep learning would not exist. In fact, it was after the publication of the Rumelhart-Hinton-Williams paper in 1986, which established the modern form of backpropagation, that the training of multi-layer neural networks became practical, and the AI boom that we see today began.
Back to the bio scenario: In the analysis of signal cascades, we do not perform a separate experiment for each upstream node. With one normal stimulus + one observation of an aberrant situation + one backward tracing, we obtain the contribution of each upstream node. The finite difference method, which is equivalent to individually knocking out or knocking down each node, would require an experiment for each node. The reason why backpropagation-style cause analysis is so much more efficient in terms of experimental cost is the same.
Computational Graphs: What PyTorch Actually Does
Up until now, our explanations have been limited to layer-by-layer neural networks. However, real-world deep learning frameworks (PyTorch, JAX) handle backpropagation in a much more general form. This form is the computational graph.
Here's the principle: When the framework executes the forward pass code, it records each operation (addition, multiplication, matrix multiplication, softmax, etc.) as a node in the graph. The edges of the graph represent the data flow.
x β (multiply by W_1) β z_1 β (ReLU) β a_1 β (multiply by W_2) β z_2 β ...Each operation node stores what operation it is and what inputs it received. Once this graph is complete, starting from the loss L, it traverses the graph in reverse order, multiplying the local gradients of each node to obtain the partial derivatives of all parameters.
This automated backpropagation is called automatic differentiation (autograd). Today, neural network developers no longer need to write the backpropagation code by hand. They only need to write the forward pass code, and the framework automatically creates the computational graph and handles the backpropagation.
In PyTorch, it looks like this:
import torch
x = torch.randn(32, 100) # Batch size 32, features 100W1 = torch.randn(100, 50, requires_grad=True)W2 = torch.randn(50, 3, requires_grad=True)y_true = torch.randint(0, 3, (32,))
# Forward pass - PyTorch automatically constructs the computational graphz1 = x @ W1a1 = torch.relu(z1)logits = a1 @ W2loss = torch.nn.functional.cross_entropy(logits, y_true)
# Backpropagation - this single line computes all partial derivativesloss.backward()
# The partial derivatives of each parameter are stored in W1.grad and W2.gradprint(W1.grad.shape, W2.grad.shape)The single line loss.backward() contains all four steps we described above. Thanks to this automation, developers can focus on the forward pass when experimenting with new architectures.
From a biological perspective: This computational graph is conceptually identical to the representation of a regulatory network in systems biology. Each node represents a reaction or regulatory relationship, and the edges represent the causal flow. The way a systems biologist analyzes the contribution of each gene to a specific downstream phenotype in a regulatory network is essentially the same logic as automatic differentiation in a computational graph.
Vanishing and Exploding Gradients: The Scourge of Backpropagation
The multiplicative structure of backpropagation can create unexpected problems.
Vanishing Gradient: If the local gradients in each layer are less than 1 (e.g., the derivative of the Sigmoid activation function is at most 0.25), the gradient will decrease exponentially as it is multiplied through multiple layers. In a 20-layer neural network, the gradient will be virtually zero when it reaches the first layer. This results in the first layer's parameters not being trained.
Exploding Gradient: Conversely, if the local gradients are greater than 1, they will grow exponentially. The gradient of the first layer's parameters will become an astronomically large value, causing the parameters to fluctuate wildly and the training to diverge.
Solutions (see Part #2 A.8 Initialization, more details in Part #12):
- ReLU-based activation functions: The derivative is exactly 1 in the region where the activation value is positive, so it doesn't vanish when multiplied.
- Xavier/He initialization: Adjust the variance of the initial weights based on the size of the layer to prevent the gradient from vanishing from the beginning.
- Batch Normalization/Layer Normalization: Adjust the distribution of the layer's activation values to maintain the gradient scale.
- Residual Connection (Part #5): Create a shortcut that skips layers, so the gradient doesn't have to be multiplied through multiple layers. This is the key to ResNet and Transformers.
- Gradient Clipping: Forcefully clip the exploding gradient to a specific value or below.
Biological Analogy: This is the same structure as each step in a signal cascade having an amplification or attenuation coefficient that is multiplied. In nature, cells have evolved to solve this problem. Each step has negative and positive feedback loops, MAPK scaffold proteins (which regulate signal strength), and a balance of phosphorylation and dephosphorylation. The initialization, normalization, and residual techniques used in neural networks conceptually correspond to these regulatory mechanisms in cellular signaling systems.
Biological Application Scenarios
Scenario 1: The Reality of Training a Cell Type Classifier
Training the scRNA-seq cell type classifier mentioned in Part #3. When you implement this in PyTorch, you only need to write the forward pass.
class CellTypeClassifier(torch.nn.Module): def __init__(self, n_genes, n_types, hidden=512): super().__init__() self.encoder = torch.nn.Sequential( torch.nn.Linear(n_genes, hidden), torch.nn.ReLU(), torch.nn.Dropout(0.3), torch.nn.Linear(hidden, hidden // 2), torch.nn.ReLU(), torch.nn.Linear(hidden // 2, n_types), )
def forward(self, expr_matrix): return self.encoder(expr_matrix)
model = CellTypeClassifier(n_genes=20000, n_types=25)optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
for cells, labels in train_loader: logits = model(cells) # Forward pass loss = torch.nn.functional.cross_entropy(logits, labels) optimizer.zero_grad() loss.backward() # Automatic backpropagation optimizer.step() # Adam updateloss.backward() automatically performs the four-step algorithm from this section. The developer only needs to focus on selecting the forward pass structure and the loss function necessary for cell type classification.
Scenario 2: Training Protein Sequence Embeddings (ESM)
Meta AI's ESM (Evolutionary Scale Modeling) is a model that learns embeddings of protein sequences. It uses hundreds of millions of protein sequences from UniProt as training data. This training also relies on the power of backpropagation. The number of parameters is thousands of times larger than the scale seen in Part #2 (on the order of 15 billion parameters). The fact that it is physically possible to train such a large neural network is due to the computational efficiency of backpropagation.
Once ESM is trained, each protein sequence can be represented as a vector, and this vector is widely used in AlphaFold, protein function prediction, and drug discovery.
Scenario 3: Parameter Estimation for Systems Biology Models
Some systems biology labs use automatic differentiation to fit the parameters of ordinary differential equation (ODE) models of cellular signaling cascades to data. The gradient of the loss with respect to each parameter is obtained using automatic differentiation, and the parameters are optimized using gradient descent. This is essentially the same procedure as neural network training. This is called differentiable programming or Scientific ML, and it is a new application area for frameworks such as JAX and PyTorch.
Key Takeaways
- Backpropagation is an efficient algorithm that obtains the partial derivatives of all parameters with a single forward pass plus a single backward pass.
- The underlying mathematics is the chain rule. The derivative of a composition of multiple functions is the product of the local derivatives of each step.
- The algorithm consists of four steps: forward pass (store activation values), loss calculation, backward propagation of the error signal (calculate Ξ΄), and calculation of the partial derivatives of each layer's parameters.
- Today's frameworks (PyTorch, JAX) completely automate this process with a computational graph + automatic differentiation. Developers only need to write the forward pass.
- The problems of vanishing and exploding gradients are mitigated by a combination of activation functions, initialization, normalization, residual connections, and gradient clipping.
- Automatic differentiation is being used as a parameter optimization tool in systems biology and scientific computing, even outside of neural networks.
π Appendix β Mathematical Formulas for Experts
Difficulty: Very Hard Target Audience: Readers with a graduate-level understanding of linear algebra, multivariable calculus, and optimization theory.
A.1 Scalar Chain Rule
The simplest form. If z = g(y), y = f(x), and both are scalar functions:
dz/dx = (dz/dy) Β· (dy/dx)For multiple steps:
dz/dx = (dz/dy_k) Β· (dy_k/dy_{k-1}) Β· ... Β· (dy_2/dy_1) Β· (dy_1/dx)A.2 Multivariate Chain Rule
If z = g(y_1, y_2, ..., y_m) and each y_i = f_i(x_1, ..., x_n):
βz/βx_j = Ξ£_{i=1}^{m} (βz/βy_i) Β· (βy_i/βx_j)This can be rewritten in vector/matrix form as the product of Jacobian matrices.
If y = f(x) and x β β^n, y β β^m, then the Jacobian is:
J_f = βy/βx = [ βy_i/βx_j ]_{i,j} (m Γ n matrix)The Jacobian of z = g(f(x)):
J_{gβf} = J_g Β· J_fBackpropagation is precisely the process of computing this matrix product from right to left (towards the input). Computing it from left to right results in forward-mode automatic differentiation. In cases like neural networks, where the input dimension is large and the output is a scalar (loss), the backward mode is overwhelmingly more efficient.
A.3 Layer-by-Layer Formula for Backpropagation (Full Form)
For an L-layer fully connected neural network:
Forward Propagation:
z^(l) = W^(l) a^(l-1) + b^(l)
a^(l) = Ο(z^(l))Loss:
L = β(a^(L), y)Backpropagation β Error signal for the last layer:
Ξ΄^(L) = βL/βz^(L) = (ββ/βa^(L)) β Ο'(z^(L))Backpropagation β Error signal for layer l (l = L-1, L-2, ..., 1):
Ξ΄^(l) = ((W^(l+1))^T Ξ΄^(l+1)) β Ο'(z^(l))Partial derivatives of each layer's parameters with respect to the loss:
βL/βW^(l) = Ξ΄^(l) (a^(l-1))^T
βL/βb^(l) = Ξ΄^(l)For a mini-batch, these partial derivatives are averaged over the examples in the batch.
A.4 Elegant Partial Derivatives for Softmax + Cross-Entropy
Softmax:
p_i = exp(z_i) / Ξ£_k exp(z_k)Cross-entropy (when the correct class is c):
L = -log(p_c)Calculate βL/βz_i.
Case 1: i = c (correct class):
βL/βz_c = -1 + p_c = p_c - 1Case 2: i β c (incorrect class):
βL/βz_i = p_iIf the correct answer is represented as a one-hot vector y, both cases can be combined into a single formula:
βL/βz = p - yThat is, the predicted probability distribution minus the one-hot encoded ground truth. Remarkably elegant, and the reason Softmax + Cross-Entropy are used together.
A.5 Derivatives of Activation Functions
Sigmoid: Ο(z) = 1/(1+e^{-z})
Ο'(z) = Ο(z)(1 - Ο(z))Maximum value of 0.25 (at z=0). Multiplying over multiple layers leads to exponential decay β the vanishing gradient problem.
Tanh: tanh(z) = (e^z - e^{-z})/(e^z + e^{-z})
tanh'(z) = 1 - tanh^2(z)Maximum value of 1. Better than Sigmoid, but still has a derivative of 0 at extremes.
ReLU: ReLU(z) = max(0, z)
ReLU'(z) = 1 if z > 0 else 0Exactly 1 in the positive region. Greatly mitigates the vanishing gradient problem. However, it suffers from the dying ReLU problem (neurons in the negative region stop training due to a gradient of 0).
Leaky ReLU: LeakyReLU(z) = z if z > 0 else Ξ±z (Ξ±=0.01)
LeakyReLU'(z) = 1 if z > 0 else Ξ±Maintains a small gradient in the negative region.
GELU: GELU(z) = z Β· Ξ¦(z) (Ξ¦: standard normal CDF)
GELU'(z) β Ξ¦(z) + z Β· Ο(z)Standard in Transformers. Natural smoothing.
A.6 Computational Graph and Reverse-Mode Automatic Differentiation
Computational graph G = (V, E):
- Nodes
V: Each atomic operation (addition, multiplication, matrix multiplication, activation function, etc.) - Edges
E: Data flow
For each node v:
- A function can be defined in advance that can compute the local partial derivative
βv/βparent_i(built into the framework).
Reverse-mode automatic differentiation:
- Start at the output node with
βL/βL = 1. - Traverse the nodes in reverse topological order of the graph.
- At each node, propagate the product of the upstream node's partial derivative and the local partial derivative.
- When a parameter node is reached, that is the final parameter partial derivative.
Time complexity: On the order of forward propagation (2-3 times the number of operations). Space complexity: Requires storing all activations, so the memory requirement is on the order of forward propagation.
A.7 Gradient Checkpointing
A trick to reduce training memory. During forward propagation, not all activations are stored, but only some are stored, and during backpropagation, the necessary activations are recomputed by running forward propagation again.
- Time: Forward propagation twice, backward propagation once β approximately 1.5 times slower.
- Space: Activation storage is reduced to the square root of L (for an L-layer neural network).
An essential technique for training large models, as mentioned in Part #14 (LLM Training in Practice).
A.8 Gradient Clipping
A technique to address exploding gradients. If the L2 norm of the gradient exceeds a threshold c, it is scaled down:
if ||β||_2 > c:
β β β Β· (c / ||β||_2)c is typically 1.0. Virtually standard in Transformer training.
A.9 Gradient Check (Implementation Verification)
Compare the backpropagation code with finite differences to check for accuracy:
βL/βΞΈ_i β [L(ΞΈ + Ξ΅ Β· e_i) - L(ΞΈ - Ξ΅ Β· e_i)] / (2Ξ΅)e_i is the i-th unit vector, and Ξ΅ β 1e-5.
Relative error metric:
relative_error = |grad_analytic - grad_numeric| / max(|grad_analytic|, |grad_numeric|)Generally, if the relative error is less than 1e-7, the implementation is considered accurate; if it is less than 1e-5, it is within the acceptable range of numerical error; if it exceeds 1e-3, a bug is suspected.
A.10 Time Complexity Summary for Backpropagation
- Forward propagation time:
O(F)β Total number of operations (FLOPs) in the neural network. - Backpropagation time:
O(F)β The same order as forward propagation (a constant factor of 2-3). - Full gradient computation time using finite differences:
O(P Β· F)β Number of parameters Γ Forward propagation.
Backpropagation obtains the entire gradient in time independent of the number of parameters. This is the key reason why training large neural networks is possible.
This elegance in time complexity is known as the Baur-Strassen theorem (1983) and forms the theoretical basis for automatic differentiation.
References
All the content, scenarios, analogies, and figures in this part are developed in-house by BioPlayground. The following are external references that can help in understanding the concepts.
- Original backpropagation paper: Rumelhart, Hinton, Williams, "Learning representations by back-propagating errors" (Nature 1986)
- Computational theory of automatic differentiation: Baur & Strassen, "The complexity of partial derivatives" (Theoretical Computer Science 1983)
- Standard deep learning textbook: Goodfellow et al., "Deep Learning" Chapter 6 (backpropagation)
- Practical guide to automatic differentiation: Baydin et al., "Automatic Differentiation in Machine Learning: a Survey" (JMLR 2018)
- PyTorch autograd documentation: pytorch.org autograd tutorial
- JAX automatic differentiation: jax.readthedocs.io β a modern implementation of functional automatic differentiation
- Deep learning visualization education: 3Blue1Brown "Deep Learning" Ch 3Β·4 (YouTube) β for pedagogical reference
- ESM paper: Rives et al., "Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences" (PNAS 2021)
With this part, both axes of neural network training (methods for exploring the loss landscape and methods for computing gradients) have been discussed. From Part #5, we will discuss how this training machine is reassembled to suit language and how it became a Transformer.
Next Concept
- Ep. #5
transformer-and-embeddingβ Why do the embedding, positional encoding, and residual connections of a Transformer stabilize learning? - Ep. #6
attention-mechanismβ Why is the gradient of the attention score the key to understanding context? - Ep. #12
pytorch-basicsβ Practice the theory from this episode with code.