## Bridging the gap between neural networks and functions

19th Apr 2023

As AI increasingly affects our lives, businesses may look to their software engineers for explanations regarding AI model outputs and guidance implementing and integrating these systems. Understanding the foundations of AI may become essential, even for software engineers that do not aspire to becoming AI researchers. In this series of articles, I shall attempt to produce a layperson software engineer’s guide to understanding neural networks, transformers and mechanistic interpretability. Instead of the guide serving as a fully comprehensive reference, I’m aiming for it to be an approachable yet detailed deep-dive into necessary intuitions for understanding and complementary to existing educational materials (which I will link to where I can).

A fundamental technique in AI is the application of neural networks, which are proven by the “universal approximation theorem” to be “universal function approximators”. They can reproduce the behavior of any continuous function within a specified domain, provided they have the appropriate structure and size. In practice, neural networks can also reproduce functions with discontinuities and capture trends in non-deterministic processes by smoothing over issues and learning expected behaviors.

### Interpretability

As software engineers, we construct functions by arranging syntax to express our logical intent — e.g. if n <= 1 { n } else { fn(n-1) + fn(n-2) }. This naturally leads to functions that are easy for humans to understand and interpret, because (1) they are written in a way that is similar to how we would explain the function to others, (2) we have a full range of high-level programming syntaxes and library calls available to us instead of only a limited set of mathematical operations, and finally (3) related logic is generally refactored so as to be colocated, SOLID, etc.

Unfortunately, the functions modeled by neural networks are not inherently human-interpretable due to the way they are constructed. A neural network is essentially an incomprehensibly large mathematical expression composed of stacked layers. Each layer forwards its outputs to another layer, and is made up from mathematical expressions representing artificial neurons computing the weighted sum of inputs plus a bias (an affine transformation but sometimes referred to as a “linear map”) followed by a non-linear activation function (e.g. activation(sum(weights * inputs) + bias)). There is nothing other than the overall mathematical expression formulated by the neural network, the numeric weights and biases used alongside it that are known to produce accurate results, and the scaffolding used to “train” the expression to perform as expected. This does not lead to natural interpretability and represents a significant issue for businesses reliant on being able to provide authentic explanations to clients on outputs or decisions that are called into question.

### Backpropagation is just the chain rule, what’s the problem?

The reason modern neural networks are constructed as giant mathematical expressions, isn’t merely due to computational efficiency, but is substantially because training them requires a process called “backpropagation” (also known as reverse-mode automatic differentiation).

Backpropagation is, in essence, just the chain rule on (multivariate) functions, and the chain rule is just a way to calculate the derivative (read “rate of change”) of a composite function by multiplying the derivatives of the functions within the composition.

If you haven’t used calculus much since school and have forgotten this, an intuitive understanding of the chain rule is that it allows us to understand the relationships between different rates of change that can be observed in nature.

“If a car travels twice as fast as a bicycle and the bicycle is four times as fast as a walking man, then the car travels $2 * 4 = 8$ times as fast as the man.”

— George F. Simmons

As an example, suppose we have a balloon being inflated such that its volume and radius are changing over time. We can use the chain rule to find out how fast the volume is changing by multiplying the rate of radius change with the rate of radius change per unit volume.

There are two ways this can be expressed. The first highlights the derivative as a composition of functions (Lagrange’s notation), while the latter is the more common form and uses the symbol $∂$ to denote “partial derivative” (Leibniz’s notation).

\begin{equation*} \begin{aligned} \textcolor{red}{V(r(t))'} &= \textcolor{green}{r'(t)} \cdot \textcolor{blue}{V'(r)} \\\\ \textcolor{red}{\frac{∂V}{∂t}} &= \textcolor{green}{\frac{∂r}{∂t}} \cdot \textcolor{blue}{\frac{∂V}{∂r}} \end{aligned} \end{equation*}

The equations above represent the chain rule applied to the balloon example, where $r$ is the radius, $V$ is the volume, and $t$ is time. The left-hand side represents “the rate of change of the volume with respect to time”, while the right-hand side represents the product of “the rate of change of the radius with respect to time” and “the rate of change of the volume with respect to the radius”.

If you still feel unclear on the chain rule check out the article “You Already Know Calculus: Derivatives” (2011) which skillfully uses everyday examples to explain calculus rules including the chain rule.

Returning to neural networks, a common pattern is to implement the mathematical expressions within it so that instead of only computing outputs they build a computational graph on application of each operation/function. This computational graph can later be traversed outwards-in to compute the derivatives of each value with respect to the final output value.

To vastly over-simplify, for scalar values it might look like this:

javascript// This is a very cut-down example to help give you the gist of the// high-level architecture of a neural network and what we mean when// we say "computational graph".function neuron(inputs) {  // ...weights and bias available for each neuron...  return {    weights,    bias,    forwardPass(inputs) {      // ...activation, add, sum, and multiply functions      // would be available and would allow us to create mathematical      // expressions that would create a "computational graph".      //      // Note: In languages with operator overloading (e.g. Python)      //       you can avoid creating functions for every mathematical      //       operation through the use of operator overloading.      //       This leads to easier to read mathematical expressions,      //       however, the con is that the creation of the computational      //       graph might not be clear to whoever reads the code.      return activation(add(sum(multiply(weights, inputs)), bias));    },  };}function layer() {  // ...neurons available for each layer...  return {    neurons,    forwardPass(inputs) {      return neurons.map((neuron) => neuron.forwardPass(inputs));    },  };}function model() {  // ...layers would be available...  return {    *parameters() {      for (let layer of layers) {        for (let neuron of layer.neurons) {          yield* neuron.weights;          yield neuron.bias;        }      }    },    forwardPass(inputs) {      let outputs = inputs;      for (let layer of layers) {        outputs = layer.forwardPass(outputs);      }      return outputs;    },  };}// The computational graph looks like a "parse tree" or "directed acyclic// graph" (DAG), in that, it is a a hierarchical data structure representing// the computations of the function.//// e.g.//// [//   {//     operation: 'relu-activation',//     output: 5,//     gradient: 0.0,//     inputs: [//       {//         operation: 'add',//         output: 5,//         gradient: 0.0,//         inputs: [//           {//             operation: 'add',//             output: 3,//             gradient: 0.0,//             inputs: [//               {//                 operation: 'source',//                 output: 1,//                 gradient: 0.0,//                 inputs: []//               },//               {//                 operation: 'source',//                 output: 2,//                 gradient: 0.0,//                 inputs: []//               }//             ]//           },//           {//             operation: 'source',//             output: 2,//             gradient: 0.0,//             inputs: []//           }//         ]//       }//     ]//   },//   ...// ]//const computationalGraph = model().forwardPass(inputs);

Being able to decompose large functions into compositions of many smaller functions is helpful when implementing a neural network, as when coupled with the chain rule and the ability to calculate the local derivative of an input with respect to its output, this allows us to decompose the relative “impact” of each input parameter on the final output. This is incredibly useful as it means we can determine the impact of each weight and bias on the overall model outputs.

$\textcolor{red}{\frac{∂L}{∂input\_value}} \mathrel{+}= \textcolor{green}{\frac{∂current\_value}{∂input\_value}} \cdot \textcolor{blue}{\frac{∂L}{∂current\_value}}$

The equation above represents the chain rule applied to a node within a neural network’s computational graph and shows how we can compute the partial derivative of “the loss function with respect to an input weight or bias” by multiplying the local derivative of the “current weight or bias with respect to its input weight or bias” by the partial derivative of “the loss function with respect to the current weight or bias”. (Note: we’ll discuss loss functions later on — for now substitute the final output wherever you see the loss function $L$ mentioned.)

• The $input\_value$ and $current\_value$ will alternate between being weights and biases, transitory computed values, hardcoded values that are part of computations, and at the edges of the computational graph, its input values and predicted/actual output values. However, from the perspective of training our network we ultimately care about the updates made to the “the loss function with respect to a weight or bias” (the gradient).

• In the example above $\textcolor{blue}{\frac{∂L}{∂current\_value}}$ (the gradient) would have been computed as $\textcolor{red}{\frac{∂L}{∂input\_value}}$ by a prior iteration of the backpropagation algorithm and therefore can be substituted with the gradient of the current value.

• On the other hand, $\textcolor{green}{\frac{∂current\_value}{∂input\_value}}$ is the local derivative and must be computed based on the type of operation/function and its input values.

• A function is differentiable if it is continuous and has a derivative at every point in its domain.

• Basic mathematical operators are trivially differentiable. For example:

Addition: When $current\_value$ was produced by $weighted\_sum + bias$ where $weighted\_sum = sum(weights \times inputs)$, calculating the derivative of $\textcolor{green}{\frac{∂current\_value}{∂input\_value}}$ for each input while holding the other constant:

\begin{equation*} \begin{aligned} \textcolor{green}{\frac{∂(weighted\_sum + bias)}{∂weighted\_sum}} &= \textcolor{green}{\frac{∂weighted\_sum} {∂weighted\_sum} + \frac{∂bias}{∂weighted\_sum}} \\ &= \textcolor{green}{1 + 0} \\ &= \textcolor{green}{1} \\\\ \textcolor{green}{\frac{∂(weighted\_sum + bias)}{∂bias}} &= \textcolor{green}{\frac{∂weighted\_sum}{∂bias} + \frac{∂bias}{∂bias}} \\ &= \textcolor{green}{0 + 1} \\ &= \textcolor{green}{1} \end{aligned} \end{equation*}

Multiplication: When $current\_value$ was produced by $weight \times input$, calculating the derivative of $\textcolor{green}{\frac{∂current\_value}{∂input\_value}}$ for each input while holding the other constant:

\begin{equation*} \begin{aligned} \textcolor{green}{\frac{∂(weight \times input)}{∂weight}} &= \textcolor{green}{input \times \frac{∂weight}{∂weight} + weight \times \frac{∂input}{∂weight}} \\ &= \textcolor{green}{input \times 1 + weight \times 0} \\ &= \textcolor{green}{input} \\\\ \textcolor{green}{\frac{∂(weight \times input)}{∂input}} &= \textcolor{green}{input \times \frac{∂weight}{∂input} + weight \times \frac{∂input}{∂input}} \\ &= \textcolor{green}{input \times 0 + weight \times 1} \\ &= \textcolor{green}{weight} \end{aligned} \end{equation*}
• Discontinuities in a function can make it non-differentiable at those specific points. For example, the non-linear activation function $ReLU(x) = max(0, x)$ is discontinuous at $x = 0$ and therefore is not differentiable at that point, however, it is still differentiable otherwise (e.g. $ReLU'(x) = 0, x < 0$ or $ReLU'(x) = 1, x > 0$). In practice, $x = 0$ is very rare and we can safely set the subderivative to 0 at that point.

• We accumulate (e.g. $\mathrel{+}=$) the result of multiplying these two partial derivatives into $\textcolor{red}{\frac{∂L}{∂input\_value}}$ which means that multiple output values of the network could contribute to the gradient of a single input weight or bias. Only after all functions/operations that an input weight or bias is involved in have been processed will the $\textcolor{red}{\frac{∂L}{∂input\_value}}$ have been computed and be ready for use as a $\textcolor{blue}{\frac{∂L}{∂current\_value}}$ in a future iteration of the backpropagation algorithm. A topological sort may be used to ensure that this is the case.

As long as there is a way to compute or approximate the local derivative of every function/operation, we can use this to help compute the derivative of the loss function with respect to every input weight and bias in the neural network.

A tiny implementation of backpropagation showing how gradients can be computed is given below:

### Neural networks: magical function approximators that sacrifice interpretability

In our pursuit of computers that can autonomously figure out how to achieve desired outcomes, we discovered that neural networks are “universal function approximators” capable of approximating arbitrary functions.

To achieve this, we stack layers (parameterized affine transformations followed by non-linear activation functions) and train our model by repeatedly adjusting its parameters until it behaves like the function we’d like to approximate.

As layers are stacked, the overall function becomes increasingly non-linear, allowing the model to represent more complex functions. However, this comes with a bit of a devil’s bargain, as the mathematics behind neural networks constrain our ability to represent functions in ways that are naturally interpretable:

1. All operations/functions used by a neural network to produce its output must be differentiable and composable.

2. Any logic or understanding learnt will be generically represented by the network’s weights and biases (its parameters) and because these participate in the calculations of derivatives these must be real-valued numeric values.

3. Mathematical expressions and parameter initializations must be carefully designed to avoid issues such as “symmetry” (neurons that produce the same outputs), “dead neurons” (neurons that always output zero), “exploding gradients” (gradients that grow exponentially in magnitude), and “vanishing gradients” (gradients that shrink exponentially in magnitude), as well as other numerical stability issues.

The price of magic was interpretability — whether we can get this back is a question for another day.

If you want to learn further about neural networks, I highly recommend following along with “The spelled-out intro to neural networks and backpropagation: building micrograd” and implementing your own version of micrograd in a language of your choosing (see my heavily-commented Rust version here).