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.
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.
Here is the code snippet for the forward pass.
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.
Here is the code snippet for the backward pass.
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.