Automatic Differentiation

:label:sec_autograd

As we have explained in :numref:sec_calculus, differentiation is a crucial step in nearly all deep learning optimization algorithms. While the calculations for taking these derivatives are straightforward, requiring only some basic calculus, for complex models, working out the updates by hand can be a pain (and often error-prone).

Deep learning frameworks expedite this work by automatically calculating derivatives, i.e., automatic differentiation. In practice, based on our designed model the system builds a computational graph, tracking which data combined through which operations to produce the output. Automatic differentiation enables the system to subsequently backpropagate gradients. Here, backpropagate simply means to trace through the computational graph, filling in the partial derivatives with respect to each parameter.

A Simple Example

As a toy example, say that we are interested in (differentiating the function $y = 2\mathbf{x}^{\top}\mathbf{x}$ with respect to the column vector $\mathbf{x}$.) To start, let us create the variable x and assign it an initial value.

```{.python .input} from mxnet import autograd, np, npx npx.set_np()

x = np.arange(4.0) x

  2. ```{.python .input}
  3. #@tab pytorch
  4. import torch
  6. x = torch.arange(4.0)
  7. x

```{.python .input}

@tab tensorflow

import tensorflow as tf

x = tf.range(4, dtype=tf.float32) x

  2. [**Before we even calculate the gradient
  3. of $y$ with respect to $\mathbf{x}$,
  4. we will need a place to store it.**]
  5. It is important that we do not allocate new memory
  6. every time we take a derivative with respect to a parameter
  7. because we will often update the same parameters
  8. thousands or millions of times
  9. and could quickly run out of memory.
  10. Note that a gradient of a scalar-valued function
  11. with respect to a vector $\mathbf{x}$
  12. is itself vector-valued and has the same shape as $\mathbf{x}$.
  14. ```{.python .input}
  15. # We allocate memory for a tensor's gradient by invoking `attach_grad`
  16. x.attach_grad()
  17. # After we calculate a gradient taken with respect to `x`, we will be able to
  18. # access it via the `grad` attribute, whose values are initialized with 0s
  19. x.grad

```{.python .input}

@tab pytorch

x.requiresgrad(True) # Same as x = torch.arange(4.0, requires_grad=True) x.grad # The default value is None

  2. ```{.python .input}
  3. #@tab tensorflow
  4. x = tf.Variable(x)

(Now let us calculate $y$.)

```{.python .input}

Place our code inside an autograd.record scope to build the computational

graph

with autograd.record(): y = 2 * np.dot(x, x) y

  2. ```{.python .input}
  3. #@tab pytorch
  4. y = 2 * torch.dot(x, x)
  5. y

```{.python .input}

@tab tensorflow

Record all computations onto a tape

with tf.GradientTape() as t: y = 2 * tf.tensordot(x, x, axes=1) y

  2. Since `x` is a vector of length 4,
  3. an inner product of `x` and `x` is performed,
  4. yielding the scalar output that we assign to `y`.
  5. Next, [**we can automatically calculate the gradient of `y`
  6. with respect to each component of `x`**]
  7. by calling the function for backpropagation and printing the gradient.
  9. ```{.python .input}
  10. y.backward()
  11. x.grad

```{.python .input}

@tab pytorch

y.backward() x.grad

  2. ```{.python .input}
  3. #@tab tensorflow
  4. x_grad = t.gradient(y, x)
  5. x_grad

(The gradient of the function $y = 2\mathbf{x}^{\top}\mathbf{x}$ with respect to $\mathbf{x}$ should be $4\mathbf{x}$.) Let us quickly verify that our desired gradient was calculated correctly.

```{.python .input} x.grad == 4 * x

  2. ```{.python .input}
  3. #@tab pytorch
  4. x.grad == 4 * x

```{.python .input}

@tab tensorflow

x_grad == 4 * x

  2. [**Now let us calculate another function of `x`.**]
  4. ```{.python .input}
  5. with autograd.record():
  6.     y = x.sum()
  7. y.backward()
  8. x.grad  # Overwritten by the newly calculated gradient

```{.python .input}

@tab pytorch

PyTorch accumulates the gradient in default, we need to clear the previous

values

x.grad.zero_() y = x.sum() y.backward() x.grad

  2. ```{.python .input}
  3. #@tab tensorflow
  4. with tf.GradientTape() as t:
  5.     y = tf.reduce_sum(x)
  6. t.gradient(y, x)  # Overwritten by the newly calculated gradient

Backward for Non-Scalar Variables

Technically, when y is not a scalar, the most natural interpretation of the differentiation of a vector y with respect to a vector x is a matrix. For higher-order and higher-dimensional y and x, the differentiation result could be a high-order tensor.

However, while these more exotic objects do show up in advanced machine learning (including [in deep learning]), more often (when we are calling backward on a vector,) we are trying to calculate the derivatives of the loss functions for each constituent of a batch of training examples. Here, (our intent is) not to calculate the differentiation matrix but rather (the sum of the partial derivatives computed individually for each example) in the batch.

```{.python .input}

When we invoke backward on a vector-valued variable y (function of x),

a new scalar variable is created by summing the elements in y. Then the

gradient of that scalar variable with respect to x is computed

with autograd.record(): y = x x # y is a vector y.backward() x.grad # Equals to y = sum(x x)

  2. ```{.python .input}
  3. #@tab pytorch
  4. # Invoking `backward` on a non-scalar requires passing in a `gradient` argument
  5. # which specifies the gradient of the differentiated function w.r.t `self`.
  6. # In our case, we simply want to sum the partial derivatives, so passing
  7. # in a gradient of ones is appropriate
  8. x.grad.zero_()
  9. y = x * x
  10. # y.backward(torch.ones(len(x))) equivalent to the below
  11. y.sum().backward()
  12. x.grad

```{.python .input}

@tab tensorflow

with tf.GradientTape() as t: y = x x t.gradient(y, x) # Same as `y = tf.reduce_sum(x x)`

  2. ## Detaching Computation
  4. Sometimes, we wish to [**move some calculations
  5. outside of the recorded computational graph.**]
  6. For example, say that `y` was calculated as a function of `x`,
  7. and that subsequently `z` was calculated as a function of both `y` and `x`.
  8. Now, imagine that we wanted to calculate
  9. the gradient of `z` with respect to `x`,
  10. but wanted for some reason to treat `y` as a constant,
  11. and only take into account the role
  12. that `x` played after `y` was calculated.
  14. Here, we can detach `y` to return a new variable `u`
  15. that has the same value as `y` but discards any information
  16. about how `y` was computed in the computational graph.
  17. In other words, the gradient will not flow backwards through `u` to `x`.
  18. Thus, the following backpropagation function computes
  19. the partial derivative of `z = u * x` with respect to `x` while treating `u` as a constant,
  20. instead of the partial derivative of `z = x * x * x` with respect to `x`.
  22. ```{.python .input}
  23. with autograd.record():
  24.     y = x * x
  25.     u = y.detach()
  26.     z = u * x
  27. z.backward()
  28. x.grad == u

```{.python .input}

@tab pytorch

x.grad.zero_() y = x x u = y.detach() z = u x

z.sum().backward() x.grad == u

  2. ```{.python .input}
  3. #@tab tensorflow
  4. # Set `persistent=True` to run `t.gradient` more than once
  5. with tf.GradientTape(persistent=True) as t:
  6.     y = x * x
  7.     u = tf.stop_gradient(y)
  8.     z = u * x
  10. x_grad = t.gradient(z, x)
  11. x_grad == u

Since the computation of y was recorded, we can subsequently invoke backpropagation on y to get the derivative of y = x * x with respect to x, which is 2 * x.

```{.python .input} y.backward() x.grad == 2 * x

  2. ```{.python .input}
  3. #@tab pytorch
  4. x.grad.zero_()
  5. y.sum().backward()
  6. x.grad == 2 * x

```{.python .input}

@tab tensorflow

t.gradient(y, x) == 2 * x

  2. ## Computing the Gradient of Python Control Flow
  4. One benefit of using automatic differentiation
  5. is that [**even if**] building the computational graph of (**a function
  6. required passing through a maze of Python control flow**)
  7. (e.g., conditionals, loops, and arbitrary function calls),
  8. (**we can still calculate the gradient of the resulting variable.**)
  9. In the following snippet, note that
  10. the number of iterations of the `while` loop
  11. and the evaluation of the `if` statement
  12. both depend on the value of the input `a`.
  14. ```{.python .input}
  15. def f(a):
  16.     b = a * 2
  17.     while np.linalg.norm(b) < 1000:
  18.         b = b * 2
  19.     if b.sum() > 0:
  20.         c = b
  21.     else:
  22.         c = 100 * b
  23.     return c

```{.python .input}

@tab pytorch

def f(a): b = a 2 while b.norm() < 1000: b = b 2 if b.sum() > 0: c = b else: c = 100 * b return c

  2. ```{.python .input}
  3. #@tab tensorflow
  4. def f(a):
  5.     b = a * 2
  6.     while tf.norm(b) < 1000:
  7.         b = b * 2
  8.     if tf.reduce_sum(b) > 0:
  9.         c = b
  10.     else:
  11.         c = 100 * b
  12.     return c

Let us compute the gradient.

```{.python .input} a = np.random.normal() a.attach_grad() with autograd.record(): d = f(a) d.backward()

  2. ```{.python .input}
  3. #@tab pytorch
  4. a = torch.randn(size=(), requires_grad=True)
  5. d = f(a)
  6. d.backward()

```{.python .input}

@tab tensorflow

a = tf.Variable(tf.random.normal(shape=())) with tf.GradientTape() as t: d = f(a) d_grad = t.gradient(d, a) d_grad

  2. We can now analyze the `f` function defined above.
  3. Note that it is piecewise linear in its input `a`.
  4. In other words, for any `a` there exists some constant scalar `k`
  5. such that `f(a) = k * a`, where the value of `k` depends on the input `a`.
  6. Consequently `d / a` allows us to verify that the gradient is correct.
  8. ```{.python .input}
  9. a.grad == d / a

```{.python .input}

@tab pytorch

a.grad == d / a

  2. ```{.python .input}
  3. #@tab tensorflow
  4. d_grad == d / a

Summary

  • Deep learning frameworks can automate the calculation of derivatives. To use it, we first attach gradients to those variables with respect to which we desire partial derivatives. We then record the computation of our target value, execute its function for backpropagation, and access the resulting gradient.

Exercises

  1. Why is the second derivative much more expensive to compute than the first derivative?
  2. After running the function for backpropagation, immediately run it again and see what happens.
  3. In the control flow example where we calculate the derivative of d with respect to a, what would happen if we changed the variable a to a random vector or matrix. At this point, the result of the calculation f(a) is no longer a scalar. What happens to the result? How do we analyze this?
  4. Redesign an example of finding the gradient of the control flow. Run and analyze the result.
  5. Let $f(x) = \sin(x)$. Plot $f(x)$ and $\frac{df(x)}{dx}$, where the latter is computed without exploiting that $f’(x) = \cos(x)$.

:begin_tab:mxnet Discussions :end_tab:

:begin_tab:pytorch Discussions :end_tab:

:begin_tab:tensorflow Discussions :end_tab: