The Chain Rule in Computational Graphs

For a function of one variable, asking how fast the output changes is straightforward: you look at one input and one output. Machine learning rarely works that way. A model usually depends on many parameters at once, and its quality is summarized by a loss value that depends on all of them together. To decide how to adjust those parameters, we need a way to describe how the output responds to each input, and then a way to package all of those responses into one object.

This article builds that object step by step. We start with a single slice of a function, move to the full collection of slices, and then see how the chain rule and PyTorch's autograd engine compute these quantities for a loss function. Along the way you will be asked to predict a gradient before the code reveals it, so keep a pencil nearby.

Take a function f that accepts several inputs, such as x and y. The partial derivative with respect to x answers this question: if I nudge x a little and hold every other input fixed, how much does the output move per unit of nudge? The same logic applies to y. Each partial derivative is therefore a one-dimensional view of a higher-dimensional surface, taken along the direction of one chosen variable.

In the nudge above, the output rose by about 0.12 for a change of 0.01 in x, so the slope along x is roughly 12 at this point. Notice that y did not move at all during the nudge. If we had instead nudged y and frozen x, the output would have changed at a different rate. That is the whole point of a partial derivative: it reports change along one axis and ignores the others.

Individual partial derivatives are useful, but a function with many inputs produces many of them, and keeping them in a loose list is awkward. The gradient packages them into one object. A gradient is a vector containing all of a multivariate function's partial derivatives, giving one entry per input variable.

The position of each entry tells you which input it belongs to. The first entry is always the partial with respect to the first input, and so on. The gradient is therefore a compact record of how the function responds to every input at once, describing the rate of change with respect to all of its variables in a single object.

Notice that the two entries can behave differently. The entry for x depends on where you are, because the slope along x changes as x changes. For $f(x, y) = x^{2} y$ the y entry is $x^{2}$, which is also location-dependent here. Both entries are still just partial derivatives, each measured along its own axis. A function with ten inputs has a gradient with ten entries, and the same reading applies to each one.

In training a model, the quantity we usually care about is a loss, a single number that serves as a proxy for how well the model performs. The loss depends on the model's parameters. The chain rule is the way to compute gradients of a loss function with respect to those parameters in a computational graph.

Those gradients provide the information needed to update each parameter so that the loss goes down. A method such as gradient descent uses them to decide the direction and rough size of each adjustment. Because the loss is a function of many parameters, the gradient with respect to all of them is exactly the multivariate object introduced earlier.

Real losses are not single formulas written in one step. They are built from a sequence of smaller operations, each feeding its result into the next. A parameter enters near the start, is combined with data, passes through further operations, and finally produces the loss. Following this path is what the chain rule does: it links how a parameter affects an intermediate result, and how that intermediate result affects the loss.

In this illustrative graph, changing w1 changes y, which changes the difference, which changes the squared loss. The chain rule multiplies these local sensitivities together along the path. This is how the gradient of the loss with respect to w1 is obtained even though w1 appears several steps before the loss.

You do not have to write the chain rule by hand in PyTorch. The source explains that PyTorch's autograd engine constructs a computational graph in the background by tracking every operation performed on tensors. When you compute the loss, the operations that produced it are already recorded, so the engine can walk back through them and produce gradients for the tracked parameters.

The printed value matches the hand calculation exactly. The engine did not need a special formula for this graph. It applied the same local-sensitivity multiplication you performed above, one recorded operation at a time, and deposited the result on the parameter. Calling backward on the loss is what produces the gradients for tracked parameters.