Backpropagation: Efficiently Assigning the Responsibility of One Result Back to All Parameters

This section addresses a core question: since finite differences can approximate derivatives, why is backpropagation still needed for training neural networks?

Finite differences use a purely numerical method. For a parameter θ_i, keep the other parameters fixed, move it in the positive and negative directions by a very small step ε, then compare the two losses:

This formula is called the central difference. Its inputs are the current parameter value and a manually chosen perturbation step size ε, and its output is the approximate slope of the loss along the i-th parameter direction.

Its calculation process can be divided into three steps:

1. Change θ_i to θ_i+ε, perform a complete forward computation to obtain the loss L(θ_i+ε).
2. Change θ_i to θ_i-ε, perform the complete forward computation again to obtain the loss L(θ_i-ε).
3. Divide the difference between the two losses by the distance between the two parameter values 2ε, to estimate the change in loss caused by a unit change in the parameter.

The central difference uses information from both sides of the perturbation point and is usually more accurate than the forward difference computed from only one side. However, it is still an approximate method rather than an analytical gradient obtained directly from the computation process.

Finite differences first face a computational cost problem. Suppose the model has P parameters, checking the central difference for each parameter requires approximately 2P complete forward computations. One million parameters means about two million forward computations, and one billion parameters means about two billion forward computations. Each time only one parameter is changed, but the entire model must be rerun; most shared computation is repeatedly redone, making it unaffordable for everyday training.

Second, the step size ε has a numerical dilemma.

When ε is too large, the loss function may already be significantly curved in this range. The difference quotient measures the average slope over an interval, not the instantaneous slope at the current point, hence producing truncation error.

When ε is too small, the two loss values become extremely close. In finite-precision floating-point arithmetic, subtracting two nearly equal numbers can easily lose significant digits; the resulting tiny error is then divided by a very small 2ε, and may be further amplified. This is rounding error or cancellation error.

Therefore, ε is not the smaller the better. In practical checks, one often tries step sizes across multiple orders of magnitude and observes whether the numerical gradient remains stable within a certain interval. Usually, as ε decreases, the truncation error first decreases; when ε is too small, the floating-point error rises again.

These limitations determine the appropriate use of finite differences: they are suitable for gradient checking and not for directly training models. One can randomly select a small number of parameters and compare the finite-difference results with the gradients obtained by automatic differentiation to detect errors in the backward formula, path accumulation, or broadcast handling.

The efficiency of backpropagation comes from a different approach. All parameters of a neural network ultimately affect the same scalar loss, and they share a large amount of intermediate computation on the way to the loss. Finite differences treat each parameter as an independent experiment, thus repeatedly redoing these computations; backpropagation saves the structure of the forward computation and the necessary intermediate values, then starts from the loss and traverses the computation graph once in reverse according to the chain rule.

In a single backward traversal, the sensitivity of a shared node to the loss only needs to be computed once and can then be reused for all upstream parameters. Therefore, backpropagation can obtain the gradients of all parameters at a cost comparable to a few forward computations.

The difference between the two methods can be summarized as:

• Finite differences ask one by one: “If I nudge this parameter alone, how will the loss change?”
• Backpropagation computes uniformly: “How is the loss sensitivity propagated back to all parameters along shared computation paths?”

So the problem with finite differences is not that they cannot compute gradients, but that they cannot efficiently and stably compute gradients repeatedly for a massive number of parameters. Backpropagation truly leverages the shared structure in the computation graph, turning the trial-and-error that originally needed to be repeated parameter by parameter into a single systematic propagation of sensitivity.

This section explains the organizational basis of backpropagation: how the computational graph breaks a complex function into many simple, reusable local differentiation problems.

Take

as an example. Although the loss can be written directly as:

and differentiate the entire composite formula, but real neural networks contain numerous matrix multiplications, biases, normalizations, and nonlinear functions. If the entire model is expanded into a single enormous formula, derivation, implementation, and maintenance all become difficult.

The computational graph adopts a different representation: it breaks the composite function into a sequence of elementary operations. For example:

Here, u is the predicted value ŷ,e is the prediction error. Multiplication, subtraction, and the half-square each constitute a local operation node, and edges between nodes represent data dependencies. Since an operation must wait for its inputs to be produced, ordinary forward computation forms a directed acyclic graph.

Forward propagation computes values along the dependency direction:

Backpropagation propagates sensitivity in the opposite direction:

What is passed here is not the forward values, but the gradient of the final loss with respect to the output of the current node, that is, "how will the loss change if the current variable changes slightly?"

Backpropagation starts from:

Then, each node only needs to combine the received upstream gradient with its own local derivative.

For the half-square node:

The local derivative is:

For the subtraction node:

The local derivative is:

For the multiplication node:

The local derivative is:

After assembling them according to the chain rule, we obtain:

and:

Throughout the entire process, no node needs to know the complete model. The multiplication node is only responsible for the backward rule of multiplication, the subtraction node only for the backward rule of subtraction, and the square node only for the backward rule of the square. The gradients of a complex network are the result of combining these local rules according to the graph structure.

Therefore, the computational graph is not just a visualization of formulas; it is also an executable differentiation plan. It needs to record at least two types of information.

The first type is the connection relationships between operations. For example, u is produced by w and x by multiplication, e in turn depends on u and y. These relationships determine in what order backpropagation should traverse which nodes.

The second type is the forward values needed by the local backward formulas. For example:

• Computing the gradient of w in the multiplication node requires the forward value x;
• Computing the gradient of x requires the forward value w;
• Computing the backward pass of the half-square loss requires the error e=ŷ-y;
• the activation function may need to know which region the forward input falls into.

This shows that "the graph structure still exists" does not equal "backpropagation can definitely be executed." If the framework still knows that a node performed multiplication, but has already lost or overwritten the input value required for the backward pass of the multiplication, it cannot correctly compute the local gradient.

The relationship between the computational graph and intermediate states also explains why training memory usage is higher than inference memory usage. After inference completes a layer, if its activations no longer participate in subsequent forward computations, they can usually be freed; training, however, must retain intermediate values needed for the backward pass until the corresponding node completes its backward computation.

It also explains why in-place modification is dangerous. If a tensor is directly overwritten after the forward pass, while the backward formula needs the original value, what the backward pass reads is no longer the state that produced the current loss. Even if the forward result looks normal, the gradient may error out or be distorted.

When the computational graph has branches, the same variable may affect the loss along multiple paths. Each path will pass back its gradient contribution, and they are summed at shared nodes. This is the manifestation of the multivariate chain rule in the graph structure, and it is also the basis for the computational graph's ability to handle residual connections, weight sharing, and recurrent structures.

So, the way the computational graph solves complex differentiation problems is not by deriving an even larger overall formula, but by establishing a clear local collaboration mechanism:

It transforms global differentiation into many small local problems, while also establishing a hard boundary: the correctness of backpropagation depends on ensuring that both the forward computational graph and the necessary intermediate values remain intact.

This section zooms into a single operation node in backpropagation and explains exactly how gradients are passed between nodes.

Suppose a node performs:

The final loss L in turn depends on z. When backpropagation passes through this node, it uses the chain rule:

These three quantities each have a clear meaning:

• ∂ L/∂ z is the upstream gradient received by the node, indicating that z when it changes slightly, how the final loss changes;
• ∂ z/∂ u is the local derivative of the node, indicating that u when it changes slightly, the node output z how it changes;
• ∂ L/∂ u is the gradient passed to the input u, indicating that u how a small change in it ultimately affects the loss.

The reason these two sensitivities are multiplied can be understood from small changes:

Substituting the first expression into the second:

Therefore:

The first term describes the influence from z to the loss, and the second term describes the influence from u to z influence. After connecting the two segments end to end, we get the complete local influence from u to the loss.

The common overbar notation can make the formula more compact:

Thus the backward rule for a single node can be written as:

In other words: first receive the upstream gradient, then transform it using the local derivative of this node, and finally pass the result to the input.

For addition:

The local derivatives are:

So:

Both inputs receive the full upstream gradient. This is not averaging the gradient across the two inputs; rather, if either input increases a little, the output increases by the same amount.

For multiplication:

The local derivatives are:

Therefore:

Computing a gradient requires the forward-time b, and computing b gradient requires the forward-time a. This explains why backpropagation must save certain forward intermediate values.

For the square:

The local derivative is:

So:

When the upstream gradient passes through the square node, it is scaled according to the current input u value and sign.

For ReLU:

When u>0 , the local derivative is 1; when u<0 , the local derivative is 0. Therefore:

ReLU can be understood as a gradient gate: the positive region lets the upstream gradient pass through unchanged, while the negative region blocks it.

At u=0 , the left and right derivatives of ReLU differ, and the classical derivative does not exist. Autodiff frameworks must define a value, typically choosing some subgradient. This value is not uniquely determined by the ordinary chain rule but is an implementation convention of the operator.

For matrix multiplication:

Let the upstream gradient be:

Then the input gradients are:

Assume:

Then:

The resulting gradients respectively have the same shape as W and X . Checking gradient shapes is a basic method for determining whether the transpose positions in matrix backward formulas are correct.

In scalar examples, “multiplying the upstream gradient by the local derivative” looks like ordinary multiplication; in vector and tensor cases, the local derivative may be a Jacobian matrix. Strictly speaking, the node performs the product of the upstream gradient and the local Jacobian. Frameworks usually do not explicitly construct the full Jacobian but instead directly implement this product.

If a node has multiple inputs, it needs to compute gradients for each input separately. If the same input affects the loss through multiple paths, the contributions returned by each path must also be added. For example, u simultaneously affects z₁ and z₂, then:

Therefore, a single path follows “upstream gradient times local derivative,” while multiple paths at a shared node must also perform “summing the contributions from each path.”

Backpropagation does not create new calculus rules. It merely organizes the chain rule into an efficient execution process: each node implements only its own local backward rule, and the framework combines these rules in reverse order of the computation graph.

This mechanism also has clear prerequisites: local derivatives must exist or have clear conventions, and the backward rules must be correctly implemented. If a custom operator's forward computation is correct but its backward formula is wrong, the model can still produce normal forward values, but incorrect gradients will continue to propagate from that node to all upstream variables, contaminating the entire training chain.

This section uses a complete manual calculation to show: the gradient -4 is not something that appears suddenly, but is the result of the seed gradient from the loss end being gradually scaled through multiple local nodes.

Given:

The model and loss are:

Together they represent the composite function:

The calculation does not directly differentiate the entire composite formula; instead, it first completes the forward propagation, saves the intermediate values, and then performs backpropagation in reverse order.

First, compute the forward pass.

Multiplication node:

Here u is the model prediction.

Subtraction node:

Here e=-2 means the prediction is 2 lower than the target.

Half-squared loss node:

The forward propagation finally yields:

Here u and e are not just temporary calculation results; they are also the raw materials needed for backpropagation. The local derivative of the square node depends on e, and the local derivative of the multiplication node depends on w and x.

Backpropagation starts from the final loss. Its seed gradient is:

This is because the derivative of any variable with respect to itself equals 1. This 1 is not a learning rate, nor a guessed gradient magnitude, but a one-unit sensitivity that starts backpropagation.

Next, pass through the half-squared node:

For input e the local derivative of the node is:

Substitute the forward-saved e=-2 and multiply by the upstream gradient 1:

The negative sign in this step comes from the forward error e=-2. Although the square operation makes the final loss positive, its derivative still retains the directional information of the error.

Then pass through the subtraction node:

The local derivative of the node with respect to u is:

Therefore:

The gradient, when passing through the subtraction path to u does not change in magnitude or sign.

Finally, pass through the multiplication node:

The local derivative with respect to the parameter w is:

Since in the forward pass x=2, so:

Therefore,-4 the complete source can be written as:

The same multiplication node can also pass the gradient to the input x. Since:

and w=2, so:

If we also treat the target value y as a variable that requires differentiation, then the subtraction node on the path to y has local derivative:

Therefore:

However, in ordinary supervised learning,x and y are usually data, and the optimizer only updates the trainable parameter w.

The gradient:

means that near the current point w=2, if we let w increase by a very small amount Δ w, the change in loss is approximately:

Therefore, slightly increasing w will decrease the loss. This is consistent with the current prediction being too low: since x=2>0, increasing w will increase the prediction u=wx, moving it from 4 toward the target value 6.

But -4 does not mean that we should directly let w increase by 4. If using the simplest gradient descent, the update formula is:

Substitute the current gradient:

The actual update magnitude is determined by the learning rate η and the optimizer rule. Backpropagation only provides the direction and rate of change near the current point; it does not determine the actual step size.

Finally, these forward values and gradients only hold for:

this set of values. After the parameter update, the prediction u, the error e, the loss L and the various local gradients will change. Therefore, after each parameter update, you need to re-run the forward pass and then perform backpropagation based on the new intermediate state.

This section explains a key rule in the multivariable chain rule: when the same variable affects the final result through multiple computation paths, its total gradient equals the sum of all path contributions.

Take:

as an example. The parameter w has two paths in the computation graph leading to u:

and:

The first path goes through the squaring operation, and the second path goes directly into the addition. As long as w changes by a small amount, the outputs of these two paths will change simultaneously, so u's total change must include both parts.

The derivative contribution of the squaring path is:

The derivative contribution of the direct path is:

Therefore:

The plus sign here is not an incidental result of algebraic expansion, but the general rule for backpropagating through forked paths in a computation graph: if a variable affects the output through multiple downstream paths, the gradient contributions returned by all paths must be summed at that variable.

If there is also a loss L that depends on u, let the upstream gradient passed from downstream to u be:

The contribution returned by the squaring path to w is:

The contribution returned by the direct path to w is:

Therefore,w's complete gradient is:

That is:

This is the path-summation form of the multivariable chain rule in a computation graph.

The same conclusion can also be reached from the perspective of small changes. Suppose w changes by a very small amount Δ w. The change in the squared branch is approximately:

The change in the direct branch is:

Because the addition node adds the results of the two branches,u's total change is:

Therefore:

If we backpropagate only along the squaring path, we would incorrectly get:

This misses the contribution of 1 from the direct path. The incorrect result may still have the correct shape and finite values, appearing to be a reasonable gradient, and is therefore more insidious than an outright error.

Missing a path does not necessarily always make the gradient smaller. Contributions from different paths may have the same sign or opposite signs. For example, if one path contributes 5 and another path contributes -3, the correct total gradient is 2. If one of the paths is missed, the result may become 5 or -3, not only changing the magnitude but possibly also flipping the sign. Therefore, missing a path may slow down training or cause parameters to update in the wrong direction.

In automatic differentiation implementations, the gradient buffer for a variable must not be directly overwritten each time a new contribution arrives:

Instead, accumulation should be performed:

If overwrite by assignment is used, the path that arrives later will erase the path that arrived earlier, and in the end only part of the derivative will be retained. Correct backward traversal must wait for or aggregate the contributions from all downstream paths to obtain the complete total gradient.

This kind of path accumulation is very common in neural networks.

A residual connection is usually written as:

The input x takes one path through the transformation F, and the other path reaches the addition node directly through the identity connection. Therefore:

where I is the local derivative of the identity path. Backpropagation must preserve the contributions of both the transformation branch and the shortcut branch.

Weight sharing also produces path accumulation. If the same parameter w is used multiple times during the forward pass, for example:

Then:

Although they refer to the same parameter object, each use forms an independent downstream path. The gradient contributions generated at all usage positions must be accumulated into the same parameter.

In recurrent neural networks, the same set of weights is used repeatedly across multiple time steps. After unfolding the recurrence over time, each time step forms a path that uses the shared weights, so the final gradient for those weights is the sum of contributions from all time steps. This is a fundamental mechanism in backpropagation through time.

It is also necessary to distinguish between in-graph path summation and cross-mini-batch gradient accumulation.

In-graph path summation occurs within the same forward and backward pass. It comes from the dependency structure of the function itself and is an indispensable part of computing the correct total derivative. Whenever the computation graph contains branching, sharing, or repeated use, the automatic differentiation system must perform this accumulation.

Cross-mini-batch gradient accumulation, on the other hand, is a training strategy. After multiple independent forward-backward passes, it temporarily does not clear the parameter gradient buffers, thereby adding the gradients from multiple batches to simulate a larger effective batch size. For example, after continuously accumulating K micro-batches, the parameters are then updated.

These two types of accumulation both use addition, but for different reasons:

• In-graph fork summation is a requirement of the chain rule and cannot be omitted;
• Cross-batch gradient accumulation is a training configuration and can be chosen whether to use;
• In-graph summation is completed automatically by the computation graph;
• Cross-batch accumulation requires the training code to explicitly decide when to retain gradients, when to update, and when to zero them.

Therefore, the key to understanding forked paths is that a variable's gradient is not a local result obtained from arbitrarily choosing one path, but the total effect of that variable on the final loss through all downstream paths. Only by completely adding up the contributions from every path can the resulting gradient be the true total derivative.

This section explains why reverse-mode automatic differentiation can handle high-dimensional tensors: it usually does not construct the full Jacobian matrix, but directly computes vector–Jacobian products.

Suppose a layer is:

where:

The full Jacobian matrix of this function is:

The elements of the matrix are:

It records the local influence of each input component on each output component. If both the input and output are large, explicitly storing J would be very expensive. For example m=n=10 000 the Jacobian contains one hundred million elements; larger tensors would produce even more unacceptable storage and computational costs.

But what neural network training really needs is usually not J itself. The final goal is a scalar loss:

When backpropagation reaches the current layer, the downstream has already given the gradient of the loss with respect to the layer output:

According to the chain rule, the gradient of the loss with respect to the layer input is:

This operation is called the vector–Jacobian product, or VJP.

If gradients are uniformly represented as column vectors, the same relationship can also be written as:

The two forms merely adopt different row/column vector conventions and have exactly the same meaning.

The key point is that the final result is only a vector of the same shape as the input x, containing n elements. Since only v^TJ, there is no need to first create the mn elements' complete Jacobian, then perform the multiplication. The framework can leverage the structure of each operation to directly compute the input gradient from the upstream gradient.

Take the linear transformation as an example:

The gradient of the loss with respect to the output is:

Then the input gradient can be directly computed as:

If the gradient of parameter W is also needed, then:

Both results can be obtained directly through matrix multiplication, without creating the complete array of partial derivatives of the output with respect to each element of x or W.

For an element-wise activation function:

Its Jacobian is theoretically a diagonal matrix:

Explicitly creating this diagonal matrix would store a large number of useless zeros. Backpropagation only needs to compute:

where ⊙ denotes element-wise multiplication. The whole process only handles tensors of the same scale as the input.

Convolution also does not need to be expanded into a huge matrix. The backward pass of convolution can directly obtain the gradients of the input and the convolution kernel using convolution or correlation operations. Attention, normalization, and other complex operators likewise implement specialized backward rules to directly complete the required VJP.

Therefore, the interface that an automatic differentiation framework implements for operators can be understood as:

The full Jacobian exists only in the mathematical description and does not need to appear as a real tensor in GPU memory. This approach simultaneously saves storage, the time to construct the Jacobian, and the cost of subsequently performing large-scale multiplication.

Reverse mode is particularly suited to neural network training because training problems usually have many inputs and few outputs. If the model has P parameters:

And the final loss is only one scalar:

Then the full derivative is a gradient containing P elements:

Reverse mode starts from:

Starting from there, one backward traversal yields the gradient of the loss with respect to all parameters. Its cost grows mainly with the size of the computation graph, rather than performing one derivation separately for each parameter.

Forward-mode automatic differentiation propagates another kind of product. Given an input direction r, forward mode computes:

This is called the Jacobian–vector product, or JVP. It answers: if the input changes slightly along direction r changes slightly, how will all outputs change.

The difference between the two modes can be summarized as:

• Reverse mode computes VJP and is suitable for many inputs and few outputs;
• Forward mode computes JVP and is suitable for few inputs and many outputs.

If a function has only a few inputs but many outputs, forward mode can obtain the changes of all outputs at once for one input direction. If reverse mode is used to obtain the full Jacobian, it usually needs to perform backward repeatedly for multiple output directions.

Conversely, neural network training has millions or even billions of parameters, but usually only one scalar loss. If forward mode were used to compute the derivative for each parameter direction, many propagations would be needed; reverse mode only needs to propagate backward one time from the scalar loss, so it is more suitable.

When the output being differentiated is not a scalar, reverse mode must provide a seed vector with the same shape as the output v. At this point what is computed is:

rather than the entire J. It is equivalent to first weighting the vector output according to v into a scalar, and then computing the gradient of this scalar with respect to the input. In standard training, the loss is already a scalar, so the initial seed is naturally 1.

The full Jacobian is not absolutely impossible to obtain. Some automatic differentiation frameworks provide an explicit Jacobian interface, and its implementation usually combines results from different directions into a full matrix through multiple VJP or JVP calls. But when both input and output dimensions are large, the time and space cost of this computation is still high, so it is not suitable as the default method for ordinary training.

Therefore, the efficiency of backpropagation lies not only in applying the chain rule, but also in computing only what training really needs: along the direction given by the scalar loss, directly transform the upstream gradient into the input gradient without explicitly storing all partial derivatives of each output with respect to each input.

This section explains the fundamental reason why training memory usage is higher than pure forward inference: backpropagation needs the intermediate states from the forward phase, so many tensors that can be released immediately during inference must be retained during training.

Although the number of model parameters does not change, training and inference have different requirements for the lifetime of intermediate results.

In pure forward inference, after a layer's output is passed to the next layer, if the previous layer's activations no longer participate in subsequent computation, their memory can usually be freed or reused. The system mainly needs to keep the tensors still in use on the forward path, so intermediate states can be produced and discarded incrementally, like a pipeline.

In addition to the forward computation, training must run backpropagation in reverse order along the computation graph. Many local backward formulas depend on the forward-time inputs, outputs, or statistics, and cannot be computed from the operator name and the upstream gradient alone.

For example, for ReLU:

The backward rule is:

Therefore, the backward phase must know which positions of the forward input were greater than 0. The framework can save the input, the output, or a compressed mask, but it must always retain enough information to recover this decision.

For a matrix layer:

Its backward formulas are:

Computing the weight gradient requires the forward input activation X, and computing the input gradient requires the parameter W. Therefore, the input activation cannot be discarded immediately after the forward pass completes.

The backward pass of a normalization layer may require the input, mean, variance, or normalized result from the forward computation. Other activation functions, attention, and convolution operators also have their own intermediate states that must be saved.

The training forward pass therefore has two tasks:

• Compute predictions and loss;
• Save the necessary materials for future backpropagation.

These activations waiting to be used by the backward pass are the first major category of additional memory usage in training compared with inference.

The main memory items during training typically include:

• Model parameters;
• Forward activations and intermediate states needed for local backward passes;
• Parameter gradients;
• Optimizer states;
• Temporary workspace used during operator execution.

Among these, parameters are a basic memory cost required by both inference and training. The additional costs in training are mainly activations, gradients, and optimizer states.

Parameter gradients typically have the same shape as the corresponding parameters. For each trainable parameter θ, backpropagation needs to save:

Therefore, parameters and gradients alone may require two equally sized storage areas, with exact sizes depending on their respective data types.

Optimizers may also maintain, for each parameter, state that persists across training steps. For example, momentum-based optimizers save historical gradient statistics. These states are not temporary data that can be released after a single backward pass; they are training state that must continue to be used in the next update.

Activation memory is typically affected by the following factors:

• Batch size;
• Sequence length;
• Feature dimension;
• Network depth;
• Types of intermediate states that operators need to save.

The larger the batch, the more samples each layer's activations contain. The longer the sequence, the larger the intermediate tensors in attention and sequence models typically are. The deeper the network, the more layers are simultaneously waiting for use by the backward pass. Therefore, even if the model parameters are completely unchanged, increasing the batch size or context length can significantly increase training memory usage.

The activation lifecycle in a single training iteration is roughly:

At the very beginning of the backward pass, activations from many layers are still simultaneously held in GPU memory, so peak memory usage is usually much higher than forward-only inference.

Gradient checkpointing changes the strategy of 'saving all necessary activations'. It divides the network into several computation segments and, during the forward pass, saves only the activations at a few boundary positions, called checkpoints. Other intermediate values within a segment can be released after they are used in the forward pass.

When backpropagation reaches a segment, the system starts from the nearest checkpoint, re-runs the forward pass for that segment, temporarily restores the intermediate values needed for the backward pass, and then completes the backward pass for that segment.

The core trade-off of this method is:

The fewer checkpoints saved, the more activations can usually be released, but the more content must be recomputed during the backward pass. Therefore, gradient checkpointing trades extra computation time for lower memory usage.

Under ideal conditions, gradient checkpointing does not change the mathematical gradients. If the first forward pass and the recomputed forward pass produce exactly the same intermediate values, the subsequent execution still uses the same set of local derivatives and the chain rule.

But this requires the recomputation process to reproduce the original forward results.

For example, dropout generates a random mask during the forward pass. If the first forward pass uses mask M₁, but the recomputation uses a different mask M₂, then the restored intermediate values do not belong to the computation that produced the current loss. The gradients subsequently computed by backpropagation also no longer correspond to the original loss. Therefore, checkpointing mechanisms usually need to save or restore the random number state so that the recomputation gets the same mask.

In-place modification is equally dangerous. If the input saved by a checkpoint or a tensor that recomputation depends on is later overwritten directly, then the recomputation no longer reproduces the original forward process.

Operators with side effects can also break reproducibility. For example, an operator might modify external state, update counters, read mutable data that is changing, or depend on unrecoverable random behavior. Even if the inputs look the same, a second execution may produce different results.

Therefore, when using gradient checkpointing, pay particular attention to:

• Whether random layers can restore the same random state;
• Whether in-place operations overwrite the inputs needed for recomputation;
• Whether operators modify external state;
• Whether the two forward passes read the same data and model state;
• Whether the implementation contains non-deterministic behavior that cannot be reproduced.

Also distinguish between 'evaluation mode' and 'disabling gradient recording'. Switching the model to evaluation mode usually only changes the behavior of layers such as dropout and normalization; it does not necessarily stop the construction of the computation graph. If the program still requests gradients, the framework may still save intermediate states. To obtain the low-memory characteristics of pure inference, it is usually necessary to explicitly disable gradient recording.

Therefore, high training memory usage is not because model parameters suddenly become more numerous, but because the system simultaneously takes on more responsibilities: saving forward history, storing backward gradients, and maintaining optimizer states. Gradient checkpointing rearranges the trade-off between computation and storage through selective forgetting and on-demand recomputation.

This section discusses the engineering boundaries of automatic differentiation: even when the mathematical formulas are perfectly correct, the tensor operations actually executed by the program can still sever or contaminate the computation graph, causing gradients to become None, all zeros, NaN, Inf, or produce results with correct shape but incorrect values.

When troubleshooting these problems, first distinguish the meanings that different symptoms may express:

• A gradient of None, usually indicates that there is no trackable computation path between the parameter and the loss, or the parameter did not participate in this forward computation;
• A constant zero gradient may come from a true zero derivative, or from gradient blocking, activation saturation, non-differentiable operations, or low-precision underflow;
• A gradient containing NaN or Inf, is usually related to numerical overflow, illegal operations, or unstable computation;
• Gradient direction seems plausible but magnitude is abnormal; this may come from missing paths, incorrect broadcasting, stale gradient accumulation, or scaling errors;
• Gradient increases continuously with training steps; this may be due to forgetting to zero the gradient, or the model itself may be experiencing exploding gradients.

Symptoms can only narrow the scope of troubleshooting, not determine the root cause alone. Reliable diagnosis must return to the tensor operations actually executed by the code and independent gradient checking results.

The first type of failure is a severed computation graph.

Let:

If a detach operation is subsequently performed:

Then compute:

Although b has the same numerical value as a, b no longer retains a 's computation history. The loss L cannot trace back along:

this path to the parameter w, therefore w 's gradient may be None.

Converting a tensor to a plain array, a plain numeric value, or passing it to an external program not tracked by the automatic differentiation system can also produce the same effect. Even if the result is later converted back to a tensor, the new tensor is usually just an independent object without the original computation history; the original dependency relationships are not automatically restored.

When troubleshooting, check:

• whether the parameter has gradient tracking enabled;
• whether the loss actually depends on that parameter;
• whether detach or an equivalent operation was used;
• whether the tensor was converted to a plain array or scalar and then reconstructed back into a tensor;
• whether the relevant computation occurred inside a scope where gradient recording is disabled;
• whether the parameter exists only in a conditional branch that was not actually executed;
• whether the model contains parameters that were not used by this sample.

Note that a gradient of None is not the same as a zero gradient.None usually indicates that no differentiable path was formed; a zero gradient indicates that the path may exist, but the current local derivative or total path contribution is zero.

The second type of failure is in-place modification of saved activations.

Backpropagation often needs to read tensors saved during the forward phase. If the code directly modifies that tensor in the original memory, what the backward phase sees may no longer be the value that produced the current loss.

For example, the backward formula for a node may require the original input u, but after the forward pass the code executes:

If this is an in-place overwrite, then the u has already changed. The local derivative may not be computable, or may be computed based on an incorrect state.

Mature frameworks usually detect some dangerous modifications through tensor version information and report an error, but the absence of an error does not guarantee that all cases are safe. Shared storage, views, custom operators, and operations that bypass framework checks can still cause hidden contamination.

Troubleshooting methods include:

• temporarily change in-place operations to non-in-place forms that produce new tensors;
• enable automatic differentiation anomaly detection;
• check whether multiple tensors share underlying storage;
• look at the forward node pointed to by the error, not just the location where the error occurred during backward;
• check whether a custom operator saved and then modified the same tensor.

The third type of failure is forgetting to zero the parameter gradients.

Many automatic differentiation frameworks by default accumulate new gradients into the parameter's existing gradient buffer:

This is to support path accumulation in the same computation graph and gradient accumulation across mini-batches. However, if the training code expects each batch to be updated independently but does not zero the gradients at the correct time, then subsequent steps will see not the current batch's gradient, but the sum of contributions from multiple batches.

A typical training sequence is:

If gradient accumulation is intentional, you should specify:

• how many micro-batches to accumulate;
• whether the loss or gradient should be scaled by the number of accumulation steps;
• when to perform the parameter update;
• when to zero the gradients after the update.

When troubleshooting, you can record the gradient value or gradient norm after each backward step. If the gradient persistently contains historical contributions without a mathematical reason, check the timing of zeroing.

The fourth type of failure comes from underflow and overflow in mixed precision.

Low-precision floating-point formats have limited representable range and significant digits. Very small gradients may underflow to 0, and very large activations or gradients may overflow to Inf, and subsequent operations may produce NaN.

The basic approach of loss scaling is to first multiply the loss by a scale factor s:

Then:

Gradients that were originally too small are amplified, reducing underflow risk. Before the parameter update, divide the gradients by s, restoring the original scale.

When diagnosing, check:

• the proportion of finite values among gradients and activations;
• the number of zero values, NaN and Inf;
• whether the problem appears only in low-precision training;
• whether the problem disappears after temporarily switching to higher precision;
• whether the order of loss scaling and unscaling is correct;
• whether gradient clipping or threshold judgment was incorrectly performed before unscaling.

If dynamic loss scaling detects non-finite gradients, it usually skips the current parameter update and adjusts the scaling factor. This skipping itself may be normal protective behavior, but if it occurs frequently, it indicates that the numerical range still has problems.

The fifth type of problem is non-differentiable discrete operations.

Rounding, hard thresholding, discrete sampling, taking category indices, and certain selection operations may cause the output to remain unchanged over a large region of the input. In this case the local derivative is usually zero; at jump positions, the classical derivative may not exist.

This is not damage to the computation graph, but rather the chosen operation itself lacks a continuous derivative suitable for ordinary gradient descent. Solutions depend on the task and may include:

• smooth approximation;
• differentiable relaxation;
• surrogate gradient;
• specialized gradient estimator;
• redesigning the loss or model objective.

What these methods obtain may be approximate gradients, surrogate gradients, or stochastic estimators, and should not be mistaken for the exact ordinary derivative of the original discrete function.

The sixth type of problem is incorrect broadcasting.

Broadcasting allows tensors of different shapes to perform legal operations, but the fact that the forward pass runs does not mean the axis semantics are correct.

For example:

If b is reused along the batch dimension, then when backpropagation computes the gradient of b when computing its gradient, it must sum the contributions from all batch positions back to b 's original shape:

If the tensor axes are arranged incorrectly, the program may still broadcast successfully but duplicate the parameter along the wrong dimension. Backpropagation will then also sum along the wrong dimension. The final gradient shape may even still be correct, but the actual meaning of each value will be misaligned.

Therefore, check:

• the complete shape of each tensor;
• whether each axis represents batch, time, channel, or feature;
• on which axes broadcasting actually occurs;
• whether the backward reduction corresponds to the expanded axes;
• whether it can be verified by hand after changing to a small example with dimensions 1, 2, 3.

A practical troubleshooting order is:

1. Confirm whether a complete computation graph path exists between the parameters and the loss.
2. Check whether detach operations, plain array conversions, or no-gradient scopes were used.
3. Check whether the activations needed for backward were modified in place.
4. Clarify the timing for zeroing parameter gradients and intentional accumulation.
5. Check for zero, NaN, Inf and other abnormal values.
6. Use higher precision for comparison to determine whether a mixed-precision problem exists.
7. Check for discrete or non-differentiable nodes.
8. Check broadcasting, reduction, and tensor axis semantics.
9. In a small deterministic example, compare automatic gradients with hand calculations or finite differences.

When performing gradient checks, you should minimize randomness, fix data and model states, use double precision, and check only a small number of parameters at a time. If numerical gradients are stable but automatic gradients clearly disagree, prioritize suspecting local backward rules, graph paths, branch accumulation, and broadcast reductions.

The same symptom may be caused by multiple failures together. For example, forgetting to zero gradients may coexist with low-precision underflow; after fixing the zeroing problem, zero-gradient phenomena may still continue. Incorrect broadcasting may also appear together with a detached branch.

Therefore, after each fix, rerun the minimal forward-backward test and gradient check. The disappearance of errors, gradients no longer being zero, or the loss starting to decrease can only prove that some symptom has changed; they cannot individually prove that the entire computation graph and all gradients are correct.

This section shows how to use finite differences to check automatic differentiation: backpropagation efficiently computes the gradients of all parameters, while finite differences serve as a relatively independent numerical route for spot-checking a small number of parameters.

Still using the running example:

The model and loss are:

The analytical gradient given by backpropagation is:

At w=2:

Now instead of using this backward formula, we will run only the forward loss and estimate the numerical gradient from both sides of the parameter. The central difference is:

Take:

Perturb in the positive direction:

The corresponding prediction is:

The loss is:

Perturb in the negative direction:

The corresponding prediction is:

The loss is:

Substitute into the central difference:

Therefore:

In this example, the loss with respect to w is quadratic. Under ideal arithmetic, central difference cancels the first-order approximation error caused by the quadratic term, so the result is particularly accurate; in actual programs, finite floating-point errors still exist.

Numerical gradients can serve as an independent reference for automatic gradients because they depend only on:

the loss values from these two forward computations. It does not use the backward formula of the node being checked. If a custom backward rule is written incorrectly while the forward computation is correct, automatic gradients and numerical gradients will usually differ.

When comparing the two, you can use the relative error:

where δ is a small positive number used to prevent division by zero.

Relative error takes into account the scale of the gradient itself. For example, when the gradient magnitude is 10³, 10^-5 an absolute difference can usually be ignored; when the gradient magnitude is 10^-6, the same absolute difference may be serious.

However, when both gradients are close to zero, the relative error may appear large because the denominator is small. Therefore, during actual checking, you should also observe the absolute error:

Only by combining absolute error, relative error, and gradient scale can you judge the result reasonably.

Finite difference checks need to satisfy several conditions; otherwise, even if the backward implementation is correct, inconsistencies may occur.

First, avoid non-differentiable points as much as possible. Take ReLU as an example:

At u=0 the left and right derivatives differ. The central difference observes both sides of the kink, while the automatic differentiation framework returns a subgradient according to a predefined convention. The two may differ, but this does not necessarily mean the backward implementation is wrong.

For piecewise functions, select test points within each differentiable region separately, and handle non-differentiable boundaries according to the operator's convention.

Second, the forward computation must be deterministic. If the two finite-difference forward passes use different dropout masks, then:

contains both parameter perturbation and random noise, so the difference quotient cannot reliably represent the derivative in the parameter direction.

Therefore, you need to:

• Turn off random behavior such as dropout, or strictly restore the same random state;
• Fix the input data;
• Fix the model state;
• Avoid updating running statistics between the two forward passes;
• Avoid operators with uncontrollable side effects.

Third, ε must be in an appropriate range.

When ε is too large, the difference quotient measures the average slope over a relatively wide interval and cannot adequately represent the local derivative at the current point, producing truncation error.

When ε is too small, the two loss values are very close and significant digits are easily lost during subtraction. The resulting rounding error, when divided by the very small 2ε, may be further amplified.

Therefore, you should not try only a single ε, but should test multiple orders of magnitude. The typical phenomenon is:

• a larger ε is affected by truncation error;
• in some middle range, the numerical gradient is relatively stable;
• an extremely small ε starts to be affected by floating-point rounding error.

Checks should preferentially use results from this stable range.

Fourth, gradient checks are best done in double precision. Central difference requires subtracting two very close loss values, and single precision more easily loses significant digits. The goal of a gradient check is to verify the differentiation logic, so you need not stick to the low-precision setting used in actual training.

Fifth, you should reduce the network and parameter scale. Finite differences require an additional forward computation for each parameter element checked. Checking a complete large model is both expensive and not conducive to locating problems.

For tensor parameters, you can randomly select a small number of indices. Each time, perturb only one of the elements while keeping the others unchanged, then compare the numerical gradient with the automatic gradient at the same position.

Sixth, check in an order from local to global:

1. First check the input gradients of a single custom operator;
2. Then check a layer or a combination of a few operators;
3. Finally check a small end-to-end network;
4. Repeat spot checks for different parameter positions and different input regions.

If you directly check the complete network and find an inconsistency, the problem could come from any node, broadcast dimension, or branching path, making it hard to locate. Local testing can more quickly identify which backward rule is wrong.

Test cases should also cover different branches of operators. For example, when checking ReLU, choose positive and negative inputs separately; when checking broadcasting, use small tensors that expose the reduction axes; when checking shared parameters, ensure that the same parameter actually participates in the loss along multiple paths.

If a gradient check fails, further localization can be based on the observed behavior:

• At multiple ε values the numerical gradient is stable but inconsistent with the automatic gradient: first check the backward formula, path accumulation, and broadcast reduction;
• The numerical gradient varies drastically with ε: check floating-point precision, randomness, non-differentiable points, and loss scale;
• The automatic gradient is None: check whether the parameter participates in the loss and whether the computation graph has been cut;
• The automatic gradient is zero while the numerical gradient is nonzero: check detach operations, wrong local derivatives, non-differentiable operations, and low-precision underflow;
• Single-layer checks pass but end-to-end checks fail: check layer combinations, shared variables, branch summation, and state changes.

Even if all spot checks pass, it cannot prove that the entire backward implementation is correct for all inputs and all parameters. It only shows that, at the selected parameter elements, input data, and evaluation points, the automatic gradient and the finite-difference gradient give similar results.

Therefore, gradient checking is a targeted sampling verification. It should cover representative parameters, numerical ranges, and computational branches, and combine small models, double precision, deterministic forward, central difference, and multiple step sizes to make this independent evidence as reliable as possible.

This section connects the complete causal chain in one training iteration and clarifies the boundaries of responsibility among backpropagation, the loss function, the optimizer, and evaluation.

The entire process can be summarized as:

Backpropagation sits between loss computation and parameter updates. It is responsible for answering “how sensitive the current loss is to each parameter,” but it does not define the training objective, nor does it decide how much parameters are actually updated.

The first step is to perform forward propagation using the current parameters.

Given input x, target y and model parameters θ, the model produces a prediction:

The loss function then computes:

In training, the loss is usually reduced to a scalar. Forward propagation also records the dependencies of the computation graph and saves intermediate values needed by the backward formulas, such as input activations, errors, activation masks, and normalization statistics.

What we get at this point is only the prediction and loss of the current parameters on the current data; the parameters themselves have not changed yet.

The second step is to start backpropagation from the scalar loss.

The final node is the loss itself, so the backward seed is:

This 1 represents one unit of sensitivity of the loss to itself. It is not a learning rate, nor is it an artificially set optimization signal, but rather the mathematical starting point of backpropagation.

If the output with respect to which we differentiate is not a scalar, an additional seed gradient with the same shape as the output must be provided. Standard training usually already reduces the output error to a scalar loss, so the initial seed is naturally 1.

The third step is to traverse the computation graph in reverse topological order.

Forward propagation must first obtain node inputs before it can compute node outputs. Backpropagation must first know the gradient of the loss with respect to the node output before it can compute the gradient of the loss with respect to the node input. Therefore, the reverse traversal order is opposite to the forward computation order.

Suppose a node performs:

The node receives the upstream gradient from downstream:

Then, combined with its own local derivative, it computes the gradient passed to the input:

Each operator is only responsible for its own local rule. A multiplication node handles the derivative of multiplication, an activation node handles the derivative of the activation function, and a matrix multiplication node handles the corresponding transposed multiplication. They do not need to know the structure of the complete model.

The fourth step is to directly compute the vector–Jacobian product.

For a tensor node:

The complete local derivative is the Jacobian matrix:

But backpropagation does not need to explicitly construct J. If the downstream gradient is:

The current node only needs to compute:

This is the vector–Jacobian product. The operator uses its own structure to directly complete this computation, avoiding the generation of a huge full Jacobian matrix.

The fifth step is to accumulate gradient contributions at forks and shared variables.

If the same variable affects the loss through multiple paths, each path returns a gradient. The total gradient of that variable must be the sum of all path contributions:

Residual connections, weight sharing, and parameter reuse in recurrent networks all rely on this rule. Keeping only one path, or letting a later-arriving contribution overwrite an earlier one, will miss part of the total derivative.

The sixth step is to obtain the gradients of all trainable parameters.

After the reverse topological traversal is complete, every parameter that participated in computing the current loss and requires gradients θ will receive:

This gradient describes the first-order change in the loss when parameters undergo small changes near the current parameter point:

At this point, the main responsibility of backpropagation has been completed. It delivers gradients, not new parameters.

The seventh step is to verify the gradient implementation using finite differences and local tests.

Finite differences estimate the numerical gradient using the forward loss on both sides of a parameter:

Compare it with the gradient obtained from backpropagation:

By comparison, you can check whether local backward rules, broadcast reductions, path accumulation, or custom operators are correct.

Such checks are not a necessary part of every training iteration, and they cannot prove that all parameters are correct for all possible values. They are a development and diagnostic tool.

The eighth step is when the optimizer updates the parameters.

The simplest gradient descent rule is:

where η is the learning rate. More complex optimizers may also combine momentum, historical squared gradients, weight decay, or other state to process the current gradient.

Therefore, the division of labor between backpropagation and the optimizer is:

• Backpropagation computes gradients;
• The optimizer reads and transforms gradients;
• The optimizer decides the update magnitude and modifies the parameters.

Backpropagation does not automatically modify parameters, nor does it decide the learning rate, momentum, gradient clipping, or weight decay.

Even if the gradient direction is correct, a single actual update does not guarantee a decrease in loss. The negative gradient gives the local descent direction in an infinitesimal neighborhood:

If the learning rate is too large, the actual update may exceed the region where the local linear approximation holds, causing the loss to oscillate or even increase. The gradient provides the local direction and rate of change, while the concrete step size remains the optimizer's responsibility.

Backpropagation is also not responsible for judging whether the training objective is correct. It only faithfully computes the derivatives of the given loss function. If the loss function does not express the true objective, the gradient may still be mathematically completely correct, but the model will optimize the wrong or incomplete objective.

For example, if training loss keeps decreasing while validation metrics do not improve, possible causes include:

• The loss is only a proxy for the actual evaluation objective;
• The model is overfitting;
• The training data distribution differs from actual usage data;
• There are issues with labels or sample weights;
• The evaluation metric includes constraints not reflected by the loss;
• The model exploits incorrect shortcuts in the data.

These problems need to be solved by data, loss design, and independent evaluation; they cannot be discovered by backpropagation alone.

Backpropagation also does not guarantee finding the global optimum. It computes the local derivative at the current point, and long-term training results also depend on model structure, initialization, data order, optimizer, learning rate, and randomness.

A complete training iteration can be organized as:

1. Zero out or retain old gradients according to the training strategy;
2. Perform forward propagation using current parameters;
3. Compute the scalar loss and save necessary intermediate state;
4. From ∂ L/∂ L=1 start the backward pass;
5. Execute the local vector–Jacobian product for each node in reverse order;
6. Accumulate all path contributions at shared variables;
7. Obtain the gradients of all trainable parameters;
8. If necessary, perform unscale, non-finite value checks, or gradient clipping;
9. The optimizer reads the gradients and updates the parameters;
10. Use the loss and independent evaluation to judge whether training actually improves the objective.

The responsibilities of the three parties can be summarized as:

Evaluation is responsible for judging whether the model improves on the criteria we truly care about.

Therefore, the precise responsibility of backpropagation is: given the current forward computation graph, the necessary intermediate state, and the final scalar loss, efficiently compute the gradients of that loss with respect to all variables that require gradients. It connects “the loss has already been computed” with “the optimizer can act,” but it does not replace objective design, parameter update strategy, or effectiveness evaluation.

This section presents the conceptual dependencies and subsequent learning pathway for backpropagation. After understanding backpropagation, you still need to answer five types of questions: from what objective to start differentiation, how gradients become updates, why deep paths weaken gradients, how network structure improves propagation, and how large-scale training balances computation, communication, and memory.

The first dependency is the loss function. Backpropagation must start from a clear output, typically a scalar loss during training:

Backpropagation computes:

Therefore, before discussing gradients, one must first answer: what scalar are we actually differentiating?

The loss function defines the optimization objective of the model. Different loss functions produce different gradients; even if the model, parameters, and data are identical, changing the loss can change the training signal received by parameters.

The relationship between the loss function and backpropagation is:

If the loss does not accurately express the true task objective, backpropagation can still produce mathematically correct gradients, but the model will efficiently optimize a wrong or incomplete objective. Therefore, a correct gradient does not mean a correct objective.

The second dependency is gradient descent. Backpropagation yields:

After this, the parameters have not changed yet. The optimizer needs to convert the gradient into a concrete update. The simplest gradient descent rule is:

where η is the learning rate.

The gradient provides the direction and rate of change near the current parameter point; the learning rate determines the actual step size. If the step is too small, training may progress slowly; if too large, parameters may overshoot the local descending region, causing the loss to oscillate or diverge.

In mini-batch training, the current gradient is usually only a stochastic estimate of the gradient of the overall objective:

Different batches produce different gradients, so updates contain sampling noise. When later studying gradient descent, you need to continue to understand:

• how the learning rate controls stability and convergence speed;
• how batch size affects gradient noise;
• how momentum integrates gradient directions across multiple steps;
• how gradient clipping limits abnormal updates;
• how weight decay changes the actual optimization objective or update rule.

Backpropagation is responsible for computing the direction, and the optimizer is responsible for deciding how to use that direction.

The third extension is vanishing and exploding gradients. Gradients in deep networks are formed by successively multiplying many local Jacobians.

Let the layers of the network be:

Then the gradient of the loss with respect to the early state h₀ contains:

If each local derivative shrinks the gradient in the relevant direction, for example each layer multiplies by about 0.5, after n layers, the gradient scale becomes approximately:

This value decreases exponentially with depth, causing early layers to receive almost no effective training signal. This is one intuitive source of vanishing gradients.

Conversely, if the scaling factor of successive Jacobians in the relevant direction is consistently greater than 1, for example each layer multiplies by about 1.5, the gradient scale may grow to:

thus producing exploding gradients.

In the tensor case, you cannot only look at a single element of the Jacobian; you must also pay attention to how the Jacobian scales different vector directions. Some directions may be compressed, while others may be amplified.

Vanishing gradients do not mean that backpropagation missed a path. It may be exactly the result correctly computed by the chain rule, just that successive local transformations make the signal extremely small. Understanding this problem requires further study of activation functions, initialization, normalization, Jacobian spectra, and numerical stability in deep networks.

The fourth extension is residual connections. A plain deep path requires the gradient to pass through all nonlinear transformations consecutively, while a residual block adds an identity shortcut:

Its derivative with respect to the input is:

Therefore, the backward gradient is:

Here, I comes from the identity path. The gradient can pass not only through the transformation branch F, but it can also propagate directly back to the input along the shortcut.

From the computational graph perspective, a residual connection is exactly the summation of forked paths:

Even if the local Jacobian in the transformation branch significantly shrinks the gradient, the identity branch can still provide a shorter, more direct propagation path.

Residual connections do not bypass the chain rule, nor can they guarantee that gradients never vanish. By changing the structure of the computational graph, they add a path whose derivative is the identity mapping, thereby improving information and gradient propagation in deep networks.

The fifth extension is the memory trade-off in distributed training. Training memory typically includes:

• model parameters;
• forward activations;
• parameter gradients;
• optimizer states;
• temporary computation space.

When these cannot fit on a single device, you can apply sharding, recomputation, or cross-device scheduling to different objects.

Parameters can be fully replicated on every device or stored in shards. Full replication allows each device to directly access all parameters, but every device bears the full model memory; parameter sharding reduces per-device occupancy, but requires additional communication during computation.

Gradients can also be synchronized, summed, averaged, or sharded across devices. In data parallelism, different devices compute local gradients using different data, then aggregate them into gradients consistent with the overall batch. The aggregation rule must be consistent with whether the loss is summed or averaged.

Optimizer state is usually proportional to the parameter size, and a single parameter may correspond to multiple historical statistics. Sharding optimizer state across multiple devices can avoid each device keeping a full copy.

Activations mainly grow with batch size, sequence length, network depth, and feature width. Methods to reduce activation memory include:

• using gradient checkpointing to recompute part of the forward pass during the backward pass;
• adjusting micro-batch size to reduce peak memory per step;
• performing pipeline partitioning so that different devices handle different network layers;
• splitting computation and activations along tensor or sequence dimensions;
• arranging more complex schedules between computation, communication, and storage.

These methods all face a fundamental trade-off:

Distributed training does not change the mathematical rules of backpropagation. Each node still performs local vector-Jacobian products, and shared variables still accumulate contributions from all paths. What changes is where parameters, activations, gradients, and optimizer states are stored, and when they are computed and how they are aggregated across devices.

The pass criterion for this section can be verified with a small computational graph. Let:

Forward propagation computes in order: a, b, s, e and L, and saves the intermediate values needed for the backward pass.

Backward pass starts from:

Start. After the half-square node:

After the subtraction node:

The addition node sends the upstream gradient separately to a and b:

The square path returns to w with a contribution of:

The direct path returns to w with a contribution of:

The two paths add together at the shared parameter w:

This example simultaneously checks three core abilities:

• being able to decompose a composite function into a local computational graph;
• being able to backpropagate node by node starting from the loss seed 1;
• being able to correctly accumulate gradients from forked paths at shared variables.

Finally, you also need to clearly explain the three-way division of labor in training:

The loss function answers "what result is better," backpropagation answers "how current parameters affect this objective," and the optimizer answers "how large a step to actually take based on these gradients."

Truly understanding backpropagation is not just being able to write the chain rule; you also need to know what objective it starts from, how it handles deep paths and forked structures, how the result is handed to the optimizer, and how large-scale training systems arrange computation, communication, and memory for the same mathematical rules.