What is forward-mode automatic differentiation?
Hello fellow machine learners,
Two weeks ago, I outlined two computational issues that neural networks can succumb to when being trained. Here is what I wrote:
🤔 If our training set is very large, then we’ll need to use the network to make predictions for every data point and take the MSE of all these values. This means that theexecution time of the forward pass scales up with the size of the training set.
🤔 The more network parameters we have, the more partial derivatives we’ll need to evaluate. Although some of them have similar expressions, the quantity and complexity of the partial derivatives can get out of hand very quickly.
We solved the first problem with techniques like stochastic gradient descent and mini-batch gradient descent. Today, we will begin to discuss how we can resolve the issue in the second point. Namely, how can we efficiently compute all the derivatives needed for a neural network?
We’ll begin this article with a comparison between numerical and algebra-symbolic approaches, before introducing a different way of computing and tracking gradients that might make more sense within the context of neural networks.
Let’s get to unpacking!
Numerical differentiation
Differentiation from first principles involves approximating the tangent to a curve at a point with a secant line, then taking a limit to bring the secant line as close to the true tangent line as possible:
If your neural network has a single output node, with n parameters in total (i.e. there are n weights and biases altogether), then you can consider the neural net as a function that takes a vector in R^n as input and returns a value in R, and we can store all the weights and biases in a vector of parameters relabelled as θ=(θ_1,θ_2,…,θ_n). The 1-d differentiation technique can be extended to cope with an n-dimensional input. Namely, to compute the partial derivative of the function f with respect to the i-th parameter, we can isolate the contribution of the scalar h to the i-th dimension using the i-th basis vector e_i:
The basis vector e_i is a vector whose i-th value is 1, with the rest of the values 0:
However, computers are not able to cope with the limit procedure, and so we remove the limit part and simply take h to be really small:
We refer to this as numerical differentiation. Now, since this technique is designed to approximate the true value of the derivative, the limitation of inaccuracy is evident. But how bad can this really be?
Here’s a game I remember playing at primary school. All the students would sit in a circle, and the first person would whisper a sentence to the individual next to them. They would have to whisper to the person next to them what the first person told them, and this would repeat until the message was propagated through the full circle. We would then compare the original messsage with the one that was received by the final person. In almost every instance, the pair of messages didn’t match. This was usually because a stray word would be added, omitted or modified from the original message by accident. But this could happen multiple times across the circle, and every deviation from the source made for a worse mismatch when comparing sentences at the end.
This toy example hopefully helps explain why the approximation error is especially bad for neural networks; we’re often in the case where partial derivatives rely on each other. This means that these errors can compound pretty quick, and so the partial derivative approximation can end up pretty far away from the true value, similar to how the pair of messages from the primary school whispering game usually ended up pretty different in the end.
Symbolic differentiation
Rather than relying on numerical approximation for the computation of derivatives, we can instead appeal to the sum, product and quotient rules of calculus to find explicit expressions for derivatives.
For instance, suppose we have the scalar-valued function
The partial derivatives are given by
This forumulation is useful, because we can evaluate derivatives with proper accuracy by inserting the numerical values for x 1 and x 2 wherever they appear in the output, without having to worry about approximations.
However, there are two big problems with symbolic differentiation:
To describe the first problem, recall from my Backpropagation explained with examples article, where we tried to symbolically represent the partial derivatives for the weights and bias in the penultimate layer. We ended up with quite a few algebraic horrors, one of which is provided below:
This issue we have here is referred to as expression swell, where the number of terms in each derivative rapidly grows as more variables are involved.
On top of this, we then have to actually evaluate each partial derivative expression, i.e. plug all the numbers in. The sigmoid function in the above example uses the exponential function for many different inputs, but some of these are repeated for different derivatives. It would be nice if we had a way of storing the results in memory to save the effort of repeated computations.
The second issue with symbolic differentiation is that it struggles to handle functions with derivative expressions that rely on programmatic tools. Take ReLU as an example:
We covered all the common activation functions in this banger if you’re interested. The gradient of ReLU is defined as:
Due to the piecewise nature of the function, it does not have an explicit form, which symbolic differentiation requires. Thus, the symbolic differentiation representation method is not robust enough for our neural network needs.
Automatic differentiation (forward mode)
How do we achieve the precision of symbolic differentiation while avoiding the issues of expression swell and lack of robustness? Enter automatic differentiation.
Primals
Let’s start with the simple example of
The main idea behind automatic differentiation (which we’ll abbreviate to AD) is to keep track of both the numerical values and the values of gradients for each constituent of the function. To do this, we break the function down into its building blocks, for which we can compute derivatives with ease.
Get AmeerSaleem’s stories in your inbox
Join Medium for free to get updates from this writer.
The standard convention for AD is to denote the input values with negative indices up to the value 0. We will then introduce intermediate variables, which will store the value of any sum, product of single function. This leaves us with the output variables, which we only have one of for the above example. Here is what this all looks like:
We can represent this information more intuitively in a computational graph. This graph depicts the dependencies of each component via graph edges:
Tangents
Now for the ‘differentiation’ part of AD- we compute the derivative of each primal value! To keep the values general, we use dot notation for the derivatives:
It’s time to put these primals and tangents to good use. Suppose we want to find the gradient of f:
💡 With respect to x_1
💡 Evaluated at the point (x_1,x_2)=(3,4).
For the first point, we set the initial tangent pair to (1,0), and for the second we insert initial values of 3 and 4:
We can propagate these values through the table of primals and tangents. With trig function values in radians (and decimals rounded to 3dp), we get:
I encourage you to try filling the values in the table yourself. Once you’ve gotten your head around the computational graph and corresponding table, it’ll help better illustrate the utility of AD’s decomposition technique.
We have just demonstrated the forward mode of automatic differentiation. We can isolate the partial derivative of interest by toggling the pair of initial tangents, and tweak the pair of initial primals to control the gradient evaluation point.
What about vector-valued functions?
Let’s add another output dimension to the previous function to make it vector-valued:
The good news is that forward mode can handle multiple function outputs without too much hassle. In this case, we just have to update the computational graph:
and the primal/tangent tables accordingly:
Crucially, forward mode AD can handle extra output dimensions with ease. The main change is the addition of more intermediate variables. But since these variables are designed to handle the simple computations, the derivatives won’t ever get out of hand like they did with symbolic differentiation. Awesome 😄
Suppose that we want to compute the partial derivative of f with respect to x_1. The good news is that we’ll get both ∂f_1/∂x_1 and ∂f_2/∂x_1 from a single forward pass. But what if we wanted ∂f_1/∂x_2 and ∂f_2/∂x_2 instead? Well, here we would have to conduct a completely new forward pass: the same graph and tables, but different numbers. I guess this is ok when our input space is small. But what about neural networks? The main thing we’ve learned from the past few weeks is that neural networks have loads of parameters, which corresponds to a very large input space. This means that, with forward mode AD, a neural network of n parameters must conduct n separate forward passes. Oh boy…
Packing it all up
We’re not quite done with our discussion of automatic differentiation. Next week, we’ll discuss a different mode of AD that will help tackle functions of large input spaces, hopefully making AD more practical for neural nets.
Today’s read has discussed a few methods of computing and tracking derivatives. The study of gradient evaluation is imperative within the context of neural networks, because their learning procedures are wholly reliant on them. Time for the usual summary:
📦 Numerical differentiation works by approximating the derivatives of functions. Since computers cannot leverage limit-based techniques, we instead rely on approximation methods, where we take a secant line close to the true tangent line of the function. This technique can be extended to functions of multiple inputs by taking aligned basis vectors into account. However, the approximation causes errors, which compound easily under neural networks’ layer-based structure.
📦 Symbolic differentiation solves the issue of inaccuracy by relying on explicit expressions for derivatives. However, this remains susceptible to expression swell, where the algebra balloons as the number of basic operations increases, i.e. the product and chain rules for differentiation have a tendency to spawn further terms. Another problem is that symbolic differentiation is unable to handle control flow. For example, the gradient of the ReLU activation function requires a conditional statement to implement in code- this is enough to throw off the symbolic process entirely.
📦 The forward mode of automatic differentiation allows us to efficiently compute derivatives by breaking the incident function down into its operational constituents and keeping track of both the coordinate values and the gradient values for each. Recording the gradient values in particular is very useful, because it helps us avoid the expression swell of symbolic differentiation.
Training complete!
I hope you enjoyed reading as much as I enjoyed writing 😁
Do leave a comment if you’re unsure about anything, if you think I’ve made a mistake somewhere, or if you have a suggestion for what we should learn about next 😎
Until next Sunday,
Ameer
Sources
- My GitHub repo where you can find the code for the entire newsletter series: https://github.com/AmeerAliSaleem/machine-learning-algorithms-unpacked
- Here is an excellent blog on the same topic that I found on HuggingFace: https://huggingface.co/blog/andmholm/what-is-automatic-differentiation
- This YT video provides a great explanation of AD from one of OpenAI’s current research scientists
Originally published at https://ameersaleem.substack.com.
