Backpropagation

Motivation

Background

Tensor Notation

General Backpropagation

Requirements

1. A computation graph where nodes represent variables and edges their dependencies. The function is assumed to exist.

2. A set of target variables whose gradients must be computed.

3. The variable to be differentiated .

4. A backpropagation operation for each variable : .

Algorithm

def general_back_propagation(T, z):
    grad_table = {z: 1}  // mutable
    for V in T:
        build_grad(V, grad_table)
    return {v: g for v, g in grad_table if v in V}

def build_grad(V, grad_table):
    if V in grad_table:
        return grad_table[V]
    grad_v = 0
    for W in get_children(V):
        D = build_grad(W, grad_table)
        grad_v += bprop(V, W, D)
    grad_table[V] = grad_v
    # + function to insert grad_table[V] into graph, along with ops creating it
    return grad_v

Minimal Backpropagation Example

Setup

import math
import torch
import torchvision
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
from torch.nn.functional import one_hot
from functools import reduce
import altair as alt
import pandas as pd

batch_sz = 64
train = DataLoader(MNIST('data/', train=True, download=True,
                         transform=torchvision.transforms.Compose([
                             torchvision.transforms.ToTensor(),
                             lambda x: torch.flatten(x),
                         ])
                        ), batch_size=batch_sz, shuffle=True)

Layer Overview

• __call__(self, x): implements the forward pass. __call__ allows us to feed an input through the layer by treating the initialised layer object as a function. For example: relu_layer = ReLU(); output = relu_layer(input).

• bp_input(self, grad): implements the backward pass, allowing us to backpropagate the gradient vector through the layer. The grad parameter is a matrix of partial derivatives of the loss, with respect to the data sent from the given layer to the next. As such, it is a (batch_sz, out) matrix. The job of bp_input is to return a (batch_sz, in) matrix to be sent to the next layer by multiplying grad by the derivative of the forward pass with respect to the input (or an equivalent operation).

• __init__(self, ...): initialises the layer, e.g. weights.

• bp_param(self, grad): the "final stop" of the backwards pass. Only applicable for layers with trainable weights. Similar to bp_input, but calculates the derivative with respect to the weights of the layer. Should return a matrix with the same shape as the weights (self.W) to be updated.

Linear Layer

1. We start by initialising the weights (in this case using the Xavier initialisation).

2. We then implement the call method. Rather than adding an explicit bias, we append a vector of ones to the layer's input (this is equivalent, and makes backpropagation simpler).

3. Backpropagation with respect to the input is just right multiplication by the transpose of the weight matrix (adjusted to remove the added 1s column)

4. Backpropagation with respect to the output is left multiplication by the transpose of the input matrix.

class LinearLayer:
    def __init__(self, in_sz, out_sz):
	self.W = self._xavier_init(in_sz + 1, out_sz)  # (in+1, out)

    def _xavier_init(self, i, o):
	return torch.Tensor(i, o).uniform_(-1, 1) * math.sqrt(6./(i + o))
    
    def __call__(self, X):  # (batch_sz, in)
	# (batch_sz, in+1)
        self.X = torch.cat([X, torch.ones(X.shape[0], 1)], dim=1)
	# (batch_sz, in+1) @ (in+1, out) = (batch_sz, out)
        return self.X @ self.W
    
    def bp_input(self, grad):
	# (batch_sz, out) @ (out, in) =  (batch_sz, in) 
	return (grad @ self.W.T)[:,:-1]

    def bp_param(self, grad):
	# (in+1, batch_sz) @ (batch_sz, out) = (in+1, out)
	return self.X.T @ grad

ReLU Layer

class ReLU:
    def __call__(self, X):
        self.X = X
        return X.clamp(min=0)  # (batch_sz, in)
    
    def bp_input(self, grad):
	return grad * (self.X > 0).float()  # (batch_sz, in)

Softmax & Cross Entropy Loss

1. This layer-loss combination often go together, so why not put them all in one layer? This saves us from having to do two separate forward and backward propagation steps.

2. I won't prove it here, but it turns out that the derivative of the loss with respect to the input to the softmax, is much simpler than the two intermediate derivative operations (see
   🔟
   Logistic Regression for full details), and bypasses the numerical stability issues that arise when we do the exponential and the logarithm.

# (batch_sz, in=out) for all dims in this layer
class SoftmaxCrossEntropyLoss:
    def __call__(self, X, Y):
        self.Y = Y
        self.Y_prob = self._softmax(X)
        self.loss = self._cross_entropy_loss(Y, self.Y_prob)
        return self.Y_prob, self.loss
    
    def _softmax(self, X):
        self.X = X
        X_adj = X - X.amax(dim=1, keepdim=True)
        exps = torch.exp(X_adj)
        return exps / exps.sum(axis=1, keepdim=True)
    
    def _cross_entropy_loss(self, Y, Y_prob):
	return (-Y * torch.log(Y_prob)).sum(axis=1).mean()
    
    def bp_input(self, grad):
	return (self.Y_prob - self.Y) * grad

Putting it all together

class NeuralNet:
    def __init__(self, input_size=28*28, hidden_size=32, output_size=10, alpha=0.001):
        self.alpha = alpha
        self.z1 = LinearLayer(input_size, hidden_size)
        self.a1 = ReLU()
        self.z2 = LinearLayer(hidden_size, output_size)
        self.loss = SoftmaxCrossEntropyLoss()
    
    def evaluate(self, X, Y):
        out = self.z2(self.a1(self.z1(X)))
        correct = torch.eq(out.argmax(axis=1), Y).double().mean()
        Y_prob, loss = self.loss(out, one_hot(Y, 10))
        return Y_prob, correct, loss
    
    def gradient_descent(self):
        delta_W1, delta_W2 = self.backprop()
        self.z1.W -= self.alpha * delta_W1
        self.z2.W -= self.alpha * delta_W2
    
    def backprop(self):
        d_out = torch.ones(*self.loss.Y.shape)
        d_z2 = self.loss.bp_input(d_out)
        d_a1 = self.z2.bp_input(d_z2)
        d_z1 = self.a1.bp_input(d_a1)

        d_w2 = self.z2.bp_param(d_z2)
        d_w1 = self.z1.bp_param(d_z1)
        return d_w1, d_w2

Training the model

model = NeuralNet()
stats = {'correct': [], 'loss': [], 'epoch': []}
for epoch in range(2):
    correct, loss = 0, 0
    for i, (X, y) in enumerate(train):
        y_prob, batch_correct, batch_loss = model.evaluate(X, y)
        model.gradient_descent()
        
        correct += batch_correct / len(train)
        loss += batch_loss / len(train)
    stats['correct'].append(correct.item())
    stats['loss'].append(loss.item())
    stats['epoch'].append(epoch)
    print(f'epoch: {epoch} | correct: {correct:.2f}, loss: {loss:.2f}')

base = alt.Chart(pd.DataFrame.from_dict(stats)).mark_line() \
          .encode(alt.X('epoch', axis=alt.Axis(title='epoch')))
line1 = base.mark_line(stroke='#5276A7', interpolate='monotone') \
            .encode(alt.Y('loss'   , axis=alt.Axis(title='Loss'   , titleColor='#5276A7'), scale=alt.Scale(domain=[0.0, max(stats['loss'   ])])), tooltip='loss'   )
line2 = base.mark_line(stroke='#57A44C', interpolate='monotone') \
            .encode(alt.Y('correct', axis=alt.Axis(title='Correct', titleColor='#57A44C'), scale=alt.Scale(domain=[min(stats['correct']), 1.0])), tooltip='correct')
alt.layer(line1, line2).resolve_scale(y = 'independent')