Yaae: Yet another autodiff engine (written in Numpy)

• In this blog post, we will have an in-depth review of what an automatic differentiation is and how to write one in pure Numpy.
• Please refer to Yaae repository for the full code implementation.
• I mostly used Wikipedia to write this blog post which, in my opinion, provides an excellent explanation of what an automatic differentiation is.
• Also, if you want to read other posts, feel free to check them at my blog.

I) Introduction

• Automatic differentiation (AD), also called algorithmic/computational differentiation, is a set of techniques to numerically differentiate a function through the use of exact formulas and floating-point values. The resulting numerical value involves no approximation error.
• AD is different from:
  • numerical differentiation:
    • $f^{\prime }(x)=\underset{h\to 0}{lim}\frac{f(x+h)-f(x)}{h}$
    • can introduce round-off errors.
  • symbolic differentiation:
    • ${\left(\frac{x^2\cos (x-7)}{\sin (x)}\right)}^{\prime }=x\csc (x)(\cos (x-7)(2-x\cot (x))-x\sin (x-7))$
    • can lead to inefficient code.
    • can have difficulty of converting a computer program into a single expression.
    • The goal of AD is not a formula, but a procedure to compute derivatives.
• Both classical methods:
  • have problems with calculating higher derivatives, where complexity and errors increase.
  • are slow at computing partial derivatives of a function with respect to many inputs, as it is needed for gradient-based optimization algorithms.
• Automatic differentiation solves all of these problems.

• For your information, you will often hear "Autograd" instead of automatic differentitation which is the name of a particular autodiff package.
• Backpropagation is a special case of autodiff applied to neural networks. But in machine learning, we often use backpropagation synonymously with autodiff.

II) Automatic differentiation modes

• There are usually 2 automatic differentiation modes:
  • Forward automatic differentiation.
  • Reverse automatic differentiation.
• For each mode, we will have an in-depth explanation and a walkthrough example but let's first define what the chain rule is.
• The chain rule is a formula to compute the derivative of a composite function.
• For a simple composition:

$$\boxed{{\displaystyle y=f(g(h(x))}}$$

1. $w_0=x$
2. $w_1=h(x)$
3. $w_2=g(w_1)$
4. $w_3=f(w_2)=y$

• Thus, using chain rule, we have:

$$\boxed{{\displaystyle \frac{\mathrm{d}y}{\mathrm{d}x}=\frac{\mathrm{d}y}{\mathrm{d}{w}_2}\frac{\mathrm{d}{w}_2}{\mathrm{d}{w}_1}\frac{\mathrm{d}{w}_1}{\mathrm{d}x}=\frac{\mathrm{d}f(w_2)}{\mathrm{d}{w}_2}\frac{\mathrm{d}g(w_1)}{\mathrm{d}{w}_1}\frac{\mathrm{d}h(w_0)}{\mathrm{d}{w}_0}}}$$

1) Forward automatic differentiation

• Forward automatic differentiation divides the expression into a sequence of differentiable elementary operations. The chain rule and well-known differentiation rules are then applied to each elementary operation.
• In forward AD mode, we traverse the chain rule from right to left. That is, according to our simple composition function, we first compute
  $\frac{\mathrm{d}{w}_1}{\mathrm{d}x}$ and then
  $\frac{\mathrm{d}{w}_2}{\mathrm{d}{w}_1}$ and finally
  $\frac{\mathrm{d}y}{\mathrm{d}{w}_2}$.
  
• Here is the recursive relation:

$$\boxed{{\displaystyle \frac{\mathrm{d}{w}_i}{\mathrm{d}x}=\frac{\mathrm{d}{w}_i}{\mathrm{d}{w}_{i-1}}\frac{\mathrm{d}{w}_{i-1}}{\mathrm{d}x}}}$$

• More precisely, we compute the derivative starting from the inside to the outside:

$$\boxed{{\displaystyle \frac{\mathrm{d}y}{\mathrm{d}x}=\frac{\mathrm{d}y}{\mathrm{d}{w}_{i-1}}\frac{\mathrm{d}{w}_{i-1}}{\mathrm{d}x}=\frac{\mathrm{d}y}{\mathrm{d}{w}_{i-1}}\left(\frac{\mathrm{d}{w}_{i-1}}{\mathrm{d}{w}_{i-2}}\frac{\mathrm{d}{w}_{i-2}}{\mathrm{d}x}\right)=\frac{\mathrm{d}y}{\mathrm{d}{w}_{i-1}}\left(\frac{\mathrm{d}{w}_{i-1}}{\mathrm{d}{w}_{i-2}}\left(\frac{\mathrm{d}{w}_{i-2}}{\mathrm{d}{w}_{i-3}}\frac{\mathrm{d}{w}_{i-3}}{\mathrm{d}x}\right)\right)=...}}$$

• Let's consider the following example:

$$\boxed{{\displaystyle \begin{array}{rl}{w}_1& =x_1=2\\ w_2& =x_2=3\\ w_3& =w_1{w}_2\\ w_4& =sin(w_1)\\ w_5& =w_3+w_4=y\end{array}}}$$

• It can be represented as a computational graph.

• Since

  $y$ depends on
  $w_1$ and
  $w_2$, its derivative will depend on
  $w_1$ and
  $w_2$ too. Thus, we have to perform 2 chain rules:
  $\frac{\mathrm{d}y}{\mathrm{d}{w}_1}$ and
  $\frac{\mathrm{d}y}{\mathrm{d}{w}_2}$.

• $\boxed{{\displaystyle \text{For }\frac{\mathrm{d}y}{\mathrm{d}{w}_1}:}}$

• Thus,

$$\boxed{{\displaystyle \begin{array}{rl}\frac{\mathrm{d}y}{\mathrm{d}{w}_1}& =\dot{w_4}+\dot{w_3}\\ & =cos(2)\ast 1+(1\ast 3+2\ast 0)\\ & \approx 2.58\end{array}}}$$

---

• $\boxed{{\displaystyle \text{For }\frac{\mathrm{d}y}{\mathrm{d}{w}_2}:}}$

• Thus,

$$\boxed{{\displaystyle \begin{array}{rl}\frac{\mathrm{d}y}{\mathrm{d}{w}_2}& =\dot{w_4}+\dot{w_3}\\ & =cos(2)\ast 0+(0\ast 3+2\ast 1)\\ & \approx 2\end{array}}}$$

---

• Setting
  $\dot{w_1}$ and
  $\dot{w_2}$ to
  $\{0,1\}$ is called seeding.
• Forward AD is more efficient than Reverse AD for functions
  $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$ with
  $m\gg n$ (more outputs than inputs) as only
  $n$ sweeps are necessary compared to
  $m$ sweeps for Reverse AD.

2) Reverse automatic differentiation

• Reverse automatic differentiation is pretty much like Forward AD but this time, we traverse the chain rule from left to right. That is, according to our simple composition function, we first compute
  $\frac{\mathrm{d}y}{\mathrm{d}{w}_2}$ and then
  $\frac{\mathrm{d}{w}_2}{\mathrm{d}{w}_1}$ and finally
  $\frac{\mathrm{d}{w}_1}{\mathrm{d}x}$.
• Here is the recursive relation:

$$\boxed{{\displaystyle \frac{\mathrm{d}y}{\mathrm{d}{w}_i}=\frac{\mathrm{d}y}{\mathrm{d}{w}_{i+1}}\frac{\mathrm{d}{w}_{i+1}}{\mathrm{d}{w}_i}}}$$

• More precisely, we compute the derivative starting from the outside to the inside:

$$\boxed{{\displaystyle \frac{\mathrm{d}y}{\mathrm{d}x}=\frac{\mathrm{d}y}{\mathrm{d}{w}_{i+1}}\frac{\mathrm{d}{w}_{i+1}}{\mathrm{d}x}=\left(\frac{\mathrm{d}y}{\mathrm{d}{w}_{i+2}}\frac{\mathrm{d}{w}_{i+2}}{\mathrm{d}{w}_{i+1}}\right)\frac{\mathrm{d}{w}_{i+1}}{\mathrm{d}x}=\left(\left(\frac{\mathrm{d}y}{\mathrm{d}{w}_{i+3}}\frac{\mathrm{d}{w}_{i+3}}{\mathrm{d}{w}_{i+2}}\right)\frac{\mathrm{d}{w}_{i+2}}{\mathrm{d}{w}_{i+1}}\right)\frac{\mathrm{d}{w}_{i+1}}{\mathrm{d}x}=...}}$$

• Let's consider the following example:

$$\boxed{{\displaystyle \begin{array}{rl}{w}_1& =x_1=2\\ w_2& =x_2=3\\ w_3& =w_1{w}_2\\ w_4& =sin(w_1)\\ w_5& =w_3+w_4=y\end{array}}}$$

• It can be represented as a computational graph.

• The reverse automatic differentiation is divided into 2 parts:
  • A forward pass to store intermediate values, needed during the backward pass.
  • A backward pass to compute the derivatives.

---

☰ Forward pass

• We simply compute the value of each node. They will be used during the backward pass to compute the derivatives.

☰ Backward pass

• Using the previous recursive formula, we start from the output to the inputs.

• Thus,

$$\boxed{{\displaystyle \begin{array}{rl}\frac{\mathrm{d}y}{\mathrm{d}{w}_1}& =\overline{w_1}=cos(2)+3\approx 2.58\\ \frac{\mathrm{d}y}{\mathrm{d}{w}_2}& =\overline{w_2}=2\end{array}}}$$

---

• Reverse AD is more efficient than Forward AD for functions
  $f:{\mathbb{R}}^n\to {\mathbb{R}}^m$ with
  $m\ll n$ (less outputs than inputs) as only
  $m$ sweeps are necessary, compared to
  $n$ sweeps for Forward AD.

III) Implementation

• There are 2 ways to implement an autodiff:
  • source-to-source transformation (hard):
    • Uses a compiler to convert source code with mathematical expressions to source code with automatic differentiation expressions.
  • operator overloading (easy):
    • Users-defined data types and the operator overloading features of a language are used to implement automatic differentiation.
• Moreover, we usually have less outputs than inputs in a neural network scenario.
• Therefore, it is more efficient to implement the Reverse AD over the Forward AD mode.
• Our implementation revolves around the class Node which is defined below with its main components:

class Node:

    def __init__(self, data, requires_grad=True, children=[]):
        pass
    
    def zero_grad(self):
        pass
        
    def backward(self, grad=None):
        pass
       
    # Operators.
    def __add__(self, other):
        # built-in python operators.
        pass
        
    def matmul(self, other):
        # custom operators.
        pass
        
    ... # More operators.
    
    # Functions. 
    def sin(self):
        pass
        
    ... # More functions.

• Let's go through each method.

---

1) __init__() method

def __init__(self, data, requires_grad=True, children=[]):
    self.data = data if isinstance(data, np.ndarray) else np.array(data)
    self.requires_grad = requires_grad
    self.children = children
    self.grad = None
    if self.requires_grad: 
        self.zero_grad()
    # Stores function.
    self._compute_derivatives = lambda: None

def zero_grad(self):
    self.grad = Node(np.zeros_like(self.data, dtype=np.float64), requires_grad=False)

• The __init__() method has the following arguments:
  • data: a scalar/tensor value.
  • requires_grad: a boolean that tells us if we want to keep track of the gradient.
  • children: a list to keep track of children.
• It has also other variables such as:
  • grad: equal to 0-scalar/tensor (only if requires_grad is True) depending on data type.
  • _compute_derivatives: lambda function, useful during backward() method (we will dive into it later on).

2) Operators & Functions

# Operators.
def __add__(self, other):
    # Operator overloading.
    op = Add(self, other)
    out = op.forward_pass()
    out._compute_derivatives = op.compute_derivatives
    return out

def matmul(self, other):
    # Custom matrix multiplication function.
    op = Matmul(self, other)
    out = op.forward_pass()
    out._compute_derivatives = op.compute_derivatives
    return out

 # Functions.
def sin(self):
    op = Sin(self)
    out = op.forward_pass()
    out._compute_derivatives = op.compute_derivatives
    return out
    
# ------
# To make things clearer, I removed some line of codes.
# Check Yaae repository for full implementation.

class Add():
    
    def __init__(self, node1, node2):
        pass
    
    def forward_pass(self):
        # Compute value and store children.
        pass

    def compute_derivatives(self):
        pass
            
class Matmul():

    def __init__(self, node1, node2):
        pass

    def forward_pass(self):
        # Compute value and store children.
        pass
        
    def compute_derivatives(self):
        pass

class Sin():

    def __init__(self, node1):
        pass
    
    def forward_pass(self):
        # Compute value and store children.
        pass

    def compute_derivatives(self):
        pass

• As you can see, the main strength of my implementation is that both Operators and Functions follow the same structure.
• Indeed, the code structure was designed in such a way it explicitly reflects the Reverse AD mode logic that is, first a forward pass and then a backward pass (to compute derivatives).
• For each operator/function, the forward_pass() will create a new Node using the operator/function output with its children stored in children.
• And each grad will be a Node whose value is a 0-scalar/tensor depending on data type.
• We then compute the gradient using compute_derivatives() method of each operator/function classes.
• This method is then stored in the lambda function _compute_derivatives() of each resulting nodes.
• Notice that this method is not executed yet. We will see the reason later.

3) backward() method

def backward(self, grad=None):
    L, visited = [], set()
    topo = topo_sort(self, L, visited)
    
    if grad is None:
        if self.shape == ():
            self.grad = Node(1, requires_grad=False)
        else:
            raise RuntimeError('backward: grad must be specified for non-0 tensor')
    else:
        self.grad = Node(grad, requires_grad=False) if isinstance(grad, (np.ndarray, list)) else grad

    for v in reversed(topo):
        v._compute_derivatives()
        
def topo_sort(v, L, visited):
    """
        Returns list of all of the children in the graph in topological order.
        Parameters:
        - v: last node in the graph.
        - L: list of all of the children in the graph in topological order (empty at the beginning).
        - visited: set of visited children.
    """
    if v not in visited:
        visited.add(v)
        for child in v.children:
            topo_sort(child, L, visited)
        L.append(v)
    return L

• In the backward() method, we need to propagate the gradient from the output to the input as done in the previous example.
• To do so, we have to:
  • Set grad of the output to 1.
  • Perform a topological sort.
  • Use the _compute_derivatives() lambda function which store the gradients of each node.
• A topological sort is performed for problems where some objects are ordered in a particular way.
• A typical example would be a set of tasks where certain tasks must be executed after other ones. Such example can be represented by a Directed Acyclic Graph (DAG) which is exactly the graph we designed.

Figure: A computational graph (DAG)

• On top of that, order matters here (since we are working with chain rules). Hence, we want the node
  $w_5$ to propagate the gradient first to node
  $w_3$ and
  $w_4$.
• Now that we know a topological sort is needed, let's perform one. For the above DAG, the topological sort will output us topo = [node_w1, node_w2, node_w3, node_w4, node_w5].
• We then reverse topo and call for each node in topo the _compute_derivatives() method which is a lambda function containing the node derivative computations, in charge of propagating the gradient through the DAG.

IV) Application

• Let’s write down our previous example using Yaae and analyze itstep by step.

# Yaae.
w1 = Node(2, requires_grad=True)
w2 = Node(3, requires_grad=True)
w3 = w2 * w1
w4 = w1.sin()
w5 = w3 + w4
z = w5
z.backward()
print(w1.grad.data) # 2.5838531634528574 ≈ 2.58
print(w2.grad.data) # 2.0

• Let's first define the following:

• In the following, the lambda function l will contain a check symbol to express that l is no more empty.

• It seems to work properly ! Since Yaae also supports tensors operations, we can easily build a small neural network library on top of it and workout some simple regression/classification problems !
• Here are the result from demo_regression.ipynb and demo_classification.ipynb files which can be found in the repository example folder.

• Still skeptical ? We can even go further by writing our own GAN.