Vectorized implementation of back-propagation
By: Assaad Moawad
Tags: Vectorization, Machine Learning, TensorFlow, Backpropagation, Neural Networks

What is vectorization and why it matters?

“Vectorization” (simplified) is the process of rewriting a loop so that instead of processing a single element of an array N times, it processes several or all elements of the array simultaneously.

Let’s start by an example of a dataset with 1000 houses sold in a specific city. For each house we have 5 information: its area, the number of rooms, the construction year, the price paid and the agency fees. The goal is to train a model that predicts the price and the agency fees from the first 3 features.

Dataset example

Let’s consider a simple linear feed-forward model with 6 weights (W11,W12,W13,W21,W22,W23) where:

Price = W11.Area + W12.NbRooms + W13.Year

Fees = W21.Area + W22.NbRooms + W23.Year

As explained in more details previously, the goal of the machine learning, is to find which values for these 6 weights fit the model’s output the closest to the dataset’s real output. We start by initialising the weights randomly. Then, we forward-propagate to calculate the predicted price and agency fees. By comparing the results with the real price and fees from the dataset, we can get a gradient of the error to back-propagate later and update the weights accordingly.

A simple implementation for this, would look something like this:

This sequential for loop on the dataset, is however too slow, and does not take advantage of modern parallelism in CPU and GPU.

Vectorizing forward-propagation

In order to achieve high performance, we need to transform the dataset into a matrix representation. If we take the column-based representation, every input from our dataset is copied to a column in the matrix.

Our weight matrix will be a matrix of 2 rows x 3 columns.

Our input matrix will be a matrix of 3 rows x 1000 columns.

Our output matrix will be a matrix of 2 rows x 1000 columns.

Our linear model that we are searching to solve, can then be presented in the following matrix-based form:

Matrix or a vectorized-form of: Weights x Inputs = Outputs

The reason why this representation works, is because this is exactly how matrix multiplication operates:

Matrix multiplication

Matrix-multiplication: a row i of the first matrix is multiplied by a column j of the second matrix to calculate the value of the cell (i , j) of the output

With the vectorized implementation, the previous for loop with 1000 iterations can now be done with very few, high-performance vectorized operations as following:

Predictions = Matrix.Multiply(Weights, Inputs)

Error = Matrix.Substract(Predictions, Outputs)

On big datasets, and using GPUs (some have 1000 cores), we can expect in thousands of times of speedup! On CPUs, there are many advanced math libraries that implement high-performance matrix operations such as openBLAS.

Vectorizing back-propagation

Vectorizing forward-propagation is easy and straightforward, it follows the model definition. The challenge is in vectorizing the back-propagation of errors.

With numbers, if we pass a number x through a function f to get y=f(x), the derivative f’ of f gives us the rate of change on y , when x changes .

through a function to get the derivative of gives us the , when . With Matrices, we need to use the Jacobian, which is a Matrix made of the partial-derivatives in respect to the different elements of the input matrix.

The rational behind, is to fix all the elements in the input matrix except one, where we add a small delta 𝛿 to it and see what elements in the output matrix are affected, to which rate, and add them together. We do this for all elements of the input matrix, we get its gradient matrix at the input. Thus it has the same shape (number of rows and columns).

Consider the following matrix operation. RxP=S, (And the first 3 equations that calculates the first 3 elements of the output S)

Equation 1: s11 = r11.p11 + r12.p21 + r13.p31 (red output)

Equation 2: s12 = r11.p12 + r12.p22 + r13.p32 (green output)

Equation 3: s21 = r21.p11 + r22.p21 + r23.p31 (yellow output)

Let’s say we already have the gradient matrix ΔS at the output S, and we want to back-propagate it to the input R (respectively P) to calculate ΔR (resp. ΔP). Since r11 is only involved in the calculation of s11 and s12 (red and green but not yellow), we can expect that only 𝛿s11 and 𝛿s12 back-propagate to 𝛿r11.

In order to find the rate of back-propagation of 𝛿s11 , we partially derivate equation 1 in respect to r11 (and consider everything else as constant), we get the rate of p11 (Another way to explain this: a small change in r11 , will be amplified by a p11 factor in s11 ).

, we partially derivate equation 1 in respect to (and consider everything else as constant), we get the rate of (Another way to explain this: a small change in , will be amplified by a factor in ). By doing the same for 𝛿s12 and equation 2, we get the rate of p12 .

and equation 2, we get the rate of . If we try to back-propagate 𝛿s13 to r11 , and derivate equation 3 in respect to r11 , we get 0 since equation 3 does not depend at all on r11. Another way to see this: if we have an error or s21, there is nothing that can be done on r11 to reduce this error. Since r11 is not involved in the calculation of s21! The same applies for all the other elements of the S matrix (s21,s22,s31,s32).

to , and derivate equation 3 in respect to , we get 0 since equation 3 does not depend at all on Another way to see this: if we have an error or s21, there is nothing that can be done on r11 to reduce this error. Since r11 is not involved in the calculation of s21! The same applies for all the other elements of the (s21,s22,s31,s32). Finally, by adding up, we get 𝛿r11=𝛿s11.p11 + 𝛿s12.p12

By doing the same for all the elements of matrix R, we get the following:

𝛿r11=𝛿s11.p11 + 𝛿s12.p12

𝛿r12=𝛿s11.p21 + 𝛿s12.p22

𝛿r13=𝛿s11.p31 + 𝛿s12.p32

𝛿r21=𝛿s21.p11 + 𝛿s22.p12

𝛿r22=𝛿s21.p21 + 𝛿s22.p22

𝛿r23=𝛿s21.p31 + 𝛿s22.p32

𝛿r31=𝛿s31.p11 + 𝛿s32.p12

𝛿r32=𝛿s31.p21 + 𝛿s32.p22

𝛿r33=𝛿s31.p31 + 𝛿s32.p32

If we look closely to the pattern, we see we can put it in a vectorized matrix multiplication way as following:

ΔR = ΔS x Transpose(P)

Similarly, if we follow the same procedure to back-propagate to P, we get the following equation: