MLSYS ENGINEERING

10.4. Automatic differentiation

Now that we understand backpropagation, you may have noticed that as long as the forward pass is defined, we can always derive its backward pass. We do not need to know what each op does inside, as long as each op's forward and backward functions are paired together, as sketched in the following base class.

Code 24. An operator with paired forward and backward methods.
class Op:
    def forward(self, x):
        raise NotImplementedError

    def backward(self, grad, x):
        raise NotImplementedError

For simplicity, suppose we have a model structured as a single chain of ops without any weights. During the forward pass, we evaluate the ops sequentially from input to output. Because each op's backward function may need the activation passed into its forward function, we save each intermediate activation along the way, as shown below.

Activation Forward op x0 x1 x2 x3 op1 op2 op3 Saved activations x0 x1 x2
Figure 49. Saving intermediate activations during the forward pass.

Here is the code snippet for the forward pass.

Code 25. Forward pass recording intermediate activations.
saved_activations = []
x = get_model_input()

for op in ops:
    saved_activations.append(x)
    x = op.forward(x)

Once the forward pass completes, the backward pass runs in reverse order using the saved activations, as shown below.

Activation Gradient Backward op Saved activations x0 x1 x2 x0 x1 x2 x3 op1 op2 op3
Figure 50. Backward pass executing in reverse order using saved activations.

Here is the code snippet for the backward pass.

Code 26. Backward pass executing ops in reverse order.
grad = 1.0
for op, saved_x in zip(reversed(ops), reversed(saved_activations)):
    grad = op.backward(grad, saved_x)

Although this example uses a single chain where each op takes only one input, modern deep learning frameworks generalize this exact pattern. An op can accept multiple inputs, such as weights and activations, and its outputs can branch into multiple downstream ops, forming a more complex graph. The user only needs to define the forward pass. The framework computes the gradients automatically. This automated mechanism is what modern frameworks refer to as automatic differentiation, or autodiff for short.

If you are interested in a complete implementation of autodiff, see readable-ml-framework.