Algorithmic differentiation

Although it’s important to understand the backpropagation algorithm, it’s unlikely that you will need to code it in practice.

Modern deep learning frameworks such as PyTorch and TensorFlow calculate the derivatives automatically, given the model specification. This is known as algorithmic differentiation.

Each functional component (linear transform, ReLU activation, loss function) in the framework knows how to compute its own derivative.

For example, the PyTorch ReLU function \(z_{\text{out}} = relu[z_{\text{in}}]\) knows how to compute the derivative of its output \(z_{\text{out}}\) with respect to its input \(z_{\text{in}}\).

Similarly, a linear function \(z_{\text{out}} = \beta+ \Omega z_{\text{in}}\) knows how to compute the derivatives of the output \(z_{\text{out}}\) with respect to the input \(z_{\text{in}}\) and with re- spect to the parameters \(\beta\) and \(\Omega\).

The algorithmic differentiation framework also knows the sequence of operations in the network and thus has all the information required to perform the forward and backward passes.

These frameworks exploit the massive parallelism of modern graphics processing units (GPUs). Computations such as matrix multiplication (which features in both the forward and backward pass) are naturally amenable to parallelization.

Moreover, it’s possible to perform the forward and backward passes for the entire batch in parallel if the modeland intermediate results in the forward pass do not exceed the available memory.

Since the training algorithm now processes the entire batch in parallel, the input becomes a multi-dimensional tensor.

In this context, a tensor can be considered the generalization of a matrix to arbitrary dimensions. Hence, a vector is a 1D tensor, a matrix is a 2D tensor, and a 3D tensor is a 3D grid of numbers.

Until now, the training data have been 1D, so the input for backpropagation would be a 2D tensor where the first dimension indexes the batch element and the second indexes the data dimension.

In subsequent chapters, we will encounter more complex structured input data. For example, in models where the input is an RGB image, the original data examples are 3D (height ×width ×channel). Here, the input to the learning framework would be a 4D tensor, where the extra dimension indexes the batch element.

Extension to arbitrary computational graphs

We have described backpropagation in a deep neural network that is naturally sequential; we calculate the intermediate quantities \(f_0,h_1,f_1,h_2 ...,f_k\) in turn.

However, models need not be restricted to sequential computation.

Later in this book, we will meet models with branching structures.

For example, we might take the values in a hidden layer and process them through two different sub-networks before recombining.

Fortunately, the ideas of backpropagation still hold if the computational graph is acyclic.

Modern algorithmic differentiation frameworks such as PyTorch and TensorFlow can handle arbitrary acyclic computational graphs.

import torch
from matplotlib import pyplot as plt
import numpy as np

X = torch.linspace(-5, 5, 10, requires_grad=True)
phi0, phi1 = torch.randint(-10, 10, (2,)) # actual
Y = phi1*X + phi0

phi0, phi1 = torch.randint(-10, 10, (2,)).float()
phi0.requires_grad=True
phi1.requires_grad=True

Y_hat = phi1*X + phi0

mse = sum((Y_hat - Y)**2)/len(Y)

mse.backward()

with torch.no_grad():

  plt.plot(X, Y, label='Actual')
  plt.plot(X, Y_hat, label='Predicted')
  plt.legend();
  plt.title(mse);
stack = [mse.grad_fn]
seen = set()

while stack:
    curr = stack.pop()

    if curr is None or curr in seen:
        continue

    seen.add(curr)
    print(curr, curr.next_functions)

    stack.extend(node for node, _ in curr.next_functions)
# Uncomment to install libraries 
# !apt-get install graphviz -y
# !pip install torchviz

import torchviz
torchviz.make_dot(mse, params={'phi0':phi0, 'phi1':phi1, 'X':X})