Over the last couple of months, I’ve found myself working out the gradient of the cosine loss twice. I seem to have lost the work the first time. This post makes sure I don’t have to work it all out a third time. Plus, if it turns out I’ve screwed this up, hopefully someone will point it out.

## Why derive? Do you even autodiff?

If you’re using Tensorflow, PyTorch, or another library with tape-based auto-differentiation, all you have to implement is the loss, and the gradient will be calculated for you. spaCy uses its own library, Thinc, that works a bit differently. We use function composition to propagate gradients instead of tape-based auto-differentiation. This means instead of the loss, we need to provide the gradient of the loss. Besides, working out the gradient for yourself sometimes is healthy. Eat your vegetables, you know?

## Solution 1: Copy and paste code from StackExchange

Whenever I google for a gradient calculation, I always end up at
`maths.stackexchange.com`

. I’m very grateful for these posts, but the truth is I
seldom follow them very well. I’m just not practiced at manipulating equations
in maths notation, like at all. I stopped taking maths classes at 16. Here’s
what I come up with if I translate the
final equation here
into code.

`def get_cossim_loss(yh, y): # Add a small constant to avoid 0 vectors yh = yh + 1e-8 y = y + 1e-8 # https://math.stackexchange.com/questions/1923613/partial-derivative-of-cosine-similarity xp = get_array_module(yh) norm_yh = xp.linalg.norm(yh, axis=1, keepdims=True) norm_y = xp.linalg.norm(y, axis=1, keepdims=True) mul_norms = norm_yh * norm_y cosine = (yh * y).sum(axis=1, keepdims=True) / mul_norms d_yh = (y / mul_norms) - (cosine * (yh / norm_yh**2)) loss = xp.abs(1-cosine).sum() return loss, -d_yh`

That looks plausible, but is it right? ¯_(ツ)_/¯.

I guess I could run some finite differences tests, but I’d rather be able to work these things out for myself. It’s not like the cosine is complicated — I should definitely be able to do this. So, let’s try.

```
import numpy as np
def sx_cosine_loss(yh, y): # Add a small constant to avoid 0 vectors yh = yh + 1e-8 y = y + 1e-8 # https://math.stackexchange.com/questions/1923613/partial-derivative-of-cosine-similarity norm_yh = np.linalg.norm(yh) norm_y = np.linalg.norm(y) mul_norms = norm_yh * norm_y cosine = (yh * y).sum() / mul_norms d_yh = (y / mul_norms) - (cosine * (yh / norm_yh**2)) loss = np.abs(1-cosine).sum() return loss, -d_yh
```

## Solution 2: Tangent and refactor

Tangent is a neat library for source-to-source autodifferentiation. You implement a function, and it generates code to calculate the gradient with respect to the inputs. This is a fantastic fit for my requirements, as I don’t want a runtime dependency. I just want some tooling assistance to help me make sure I’m doing it right. First, a quick refresher. Here’s the cosine loss:

```
import numpy as np
def cosine_loss(X, Y): xnorm = np.sqrt(np.sum(X*X)) ynorm = np.sqrt(np.sum(Y*Y)) similarity = np.sum(X*Y) / (xnorm * ynorm) return 1 - similarity
```

The similarity calculation ranges from -1 to 1. We want the loss to be 0 when
the similarity is 1, so the loss is `1-similarity`

. Here’s what Tangent
generates for the gradient calculation of this function:

```
import tangent
tangent_cosine_loss = tangent.grad(cosine_loss, verbose=True)
def dcosine_lossdX(X, Y, b_return=1.0): X_times_X = X * X _xnorm = np.sum(X_times_X) xnorm = np.sqrt(_xnorm) Y_times_Y = Y * Y _ynorm = np.sum(Y_times_Y) ynorm = np.sqrt(_ynorm) _similarity2 = xnorm * ynorm X_times_Y = X * Y _similarity = np.sum(X_times_Y) similarity = _similarity / _similarity2 _return = 1 - similarity assert tangent.shapes_match(_return, b_return ), 'Shape mismatch between return value (%s) and seed derivative (%s)' % ( numpy.shape(_return), numpy.shape(b_return))
# Grad of: _similarity = np.sum(X_times_Y) _bsimilarity = -tangent.unbroadcast(b_return, similarity) bsimilarity = _bsimilarity
# Grad of: similarity = np.sum(X * Y) / (xnorm * ynorm) _b_similarity = bsimilarity / _similarity2 _b_similarity2 = -bsimilarity * _similarity / (_similarity2 * _similarity2) b_similarity = _b_similarity b_similarity2 = _b_similarity2 _bX_times_Y = tangent.astype(tangent.unreduce(b_similarity, numpy.shape (X_times_Y), None, False), X_times_Y) bX_times_Y = _bX_times_Y _bX3 = tangent.unbroadcast(bX_times_Y * Y, X) bX = _bX3 _bxnorm = tangent.unbroadcast(b_similarity2 * ynorm, xnorm) bxnorm = _bxnorm
# Grad of: xnorm = np.sqrt(np.sum(X * X)) _xnorm2 = xnorm _b_xnorm = bxnorm / (2.0 * _xnorm2) b_xnorm = _b_xnorm _bX_times_X = tangent.astype(tangent.unreduce(b_xnorm, numpy.shape( X_times_X), None, False), X_times_X) bX_times_X = _bX_times_X _bX = tangent.unbroadcast(bX_times_X * X, X) _bX2 = tangent.unbroadcast(bX_times_X * X, X) bX = tangent.add_grad(bX, _bX) bX = tangent.add_grad(bX, _bX2) return bX
```

Code! How nice. One thing we can do right away is test the code we cut and paste from StackExchange:

```
vec1 = np.random.uniform(size=128)vec2 = np.random.uniform(size=128)sx_loss, sx_grad = sx_cosine_loss(vec1, vec2)tangent_grad = tangent_cosine_loss(vec1, vec2)
import numpy.testing
numpy.testing.assert_almost_equal(tangent_grad, sx_grad)
```

Checks out! From a practical perspective, our work is done here. We can now be pretty confident that the code from StackExchange is correct, so if we just want to move forward with our experiments, we can go ahead and paste it in. This is still not very satisfying, though. I don’t feel like I’m much closer to working out the gradient for myself. So let’s dig a little deeper. The first step is to break things up into smaller functions, so that Tangent is easier to work with. We’ll start with a function for the norm calculation:

```
def L2_norm(X): XX = X*X XX_sum = np.sum(XX) XX_sqrt = np.sqrt(XX_sum) return XX_sqrt
dL2_normdX = tangent.grad(L2_norm, verbose=True)
```

```
def dL2_normdX(X, bXX_sqrt=1.0): XX = X * X XX_sum = np.sum(XX) XX_sqrt = np.sqrt(XX_sum) assert tangent.shapes_match(XX_sqrt, bXX_sqrt ), 'Shape mismatch between return value (%s) and seed derivative (%s)' % ( numpy.shape(XX_sqrt), numpy.shape(bXX_sqrt))
# Beginning of backward pass _XX_sqrt = XX_sqrt
# Grad of: XX_sqrt = np.sqrt(XX_sum) _bXX_sum = bXX_sqrt / (2.0 * _XX_sqrt) bXX_sum = _bXX_sum
# Grad of: XX_sum = np.sum(XX) _bXX = tangent.astype(tangent.unreduce(bXX_sum, numpy.shape(XX), None, False), XX) bXX = _bXX
# Grad of: XX = X * X _bX = tangent.unbroadcast(bXX * X, X) _bX2 = tangent.unbroadcast(bXX * X, X) bX = _bX bX = tangent.add_grad(bX, _bX2) return bX
```

The first thing we want to do is make some cosmetic improvements to the generated code, to make it more readable. We’ll strip out the assertions, some of the type casting, and make some renames:

```
def dL2_normdX_v2(X, bXX_sqrt=1.0): XX = X * X XX_sum = np.sum(XX) XX_sqrt = np.sqrt(XX_sum) # Beginning of backward pass # Grad of: XX_sqrt = np.sqrt(XX_sum) bXX_sum = bXX_sqrt / (2.0 * XX_sqrt) # Grad of: XX_sum = np.sum(XX) bXX = tangent.unreduce(bXX_sum, numpy.shape(XX), None, False) # Grad of: XX = X * X bX = tangent.unbroadcast(bXX * X, X) bX2 = tangent.unbroadcast(bXX * X, X) return tangent.add_grad(bX, bX2)
numpy.testing.assert_almost_equal(dL2_normdX(vec1), dL2_normdX_v2(vec1))
```

Next we want to remove the ugly `tangent.unreduce`

and `tangent.broadcast`

calls. These just make the code more general to different input shapes, which we
don’t need for now. If you have a reduction operation `total = vector.sum()`

, as
you backprop you’ll end up with a value `d_total`

, and want to use it to
calculate `d_vector`

. This will just be `np.array(d_total, shape=vector.shape)`

.
This makes sense, right? If you needed the summation to come out lower, you
should decrease the value for everything you’re summing up — and vice versa if
you’re trying to make the summation come out lower. Instead of creating the
`d_vector`

array explicitly, we can just use numpy’s broadcasting rules. We’ll
also make the obvious simplicitation of multiplying the gradient by 2, instead
of calculating the same thing twice and adding them together:

```
def dL2_normdX_v2(X, bXX_sqrt=1.0): XX = X * X XX_sum = np.sum(XX) XX_sqrt = np.sqrt(XX_sum) # Beginning of backward pass bXX_sum = bXX_sqrt / (2.0 * XX_sqrt) return 2 * X * bXX_sum
numpy.testing.assert_almost_equal(dL2_normdX(vec1), dL2_normdX_v2(vec1))
```

One further transformation will make the function easier to work with. Instead of calculating the forward and backward passes in one function, let’s move the backward pass into a callback:

```
def L2_norm(X): XX = X * X XX_sum = np.sum(XX) XX_sqrt = np.sqrt(XX_sum) def get_dX(bXX_sqrt): bXX_sum = bXX_sqrt / (2.0 * XX_sqrt) return 2 * X * bXX_sum return XX_sqrt, get_dX
L2_vec1, get_d_vec1 = L2_norm(vec1)d_vec1 = get_d_vec1(1.0)numpy.testing.assert_almost_equal(dL2_normdX(vec1), d_vec1)
```

Using this callback-based approach makes it easy to write higher-order functions that remain differentiable, without having to repeat calculations. For instance, here’s a function that implements a feed-forward relationship:

`def chain(func1, func2): """Compose two functions f and g, such that g(f(x))""" def forward(*inputs): result1, get_d_inputs = func1(*inputs) outputs, get_d_result1 = func2(*result1) def backprop(*d_outputs): d_result1 = get_d_result1(*d_outputs) d_inputs = get_d_inputs(*d_result1) return d_inputs return outputs, backprop return forward`

With that in mind, let’s do the rest of the work to finish the gradient of the
cosine loss. First we’ll make a little function that isolates the gradient of
the dot product. This means we only have to work out the gradient of the step
`1-(XY_dot / (X_norm * Y_norm)`

. We can use Tangent to make sure we can’t screw
up even this small step:

```
def dot(X, Y): XY = X*Y XY_sum = XY.sum()
def get_dX(dXY_sum): return Y * dXY_sum
def get_dY(dXY_sum): return X * dXY_sum
return XY_sum, (get_dX, get_dY)
def cosine_division(XY_dot, X_norm, Y_norm): similarity = XY_dot / (X_norm * Y_norm) loss = 1-similarity return loss
tangent_backprop_cosine_division = tangent.grad(cosine_division, verbose=True, wrt=(0, 1, 2))
```

```
def dcosine_divisiondXY_dotX_normY_norm(XY_dot, X_norm, Y_norm, bloss=1.0): _similarity = X_norm * Y_norm similarity = XY_dot / _similarity loss = 1 - similarity assert tangent.shapes_match(loss, bloss ), 'Shape mismatch between return value (%s) and seed derivative (%s)' % ( numpy.shape(loss), numpy.shape(bloss))
# Grad of: loss = 1 - similarity _bsimilarity = -tangent.unbroadcast(bloss, similarity) bsimilarity = _bsimilarity
# Grad of: similarity = XY_dot / (X_norm * Y_norm) _bXY_dot = bsimilarity / _similarity _b_similarity = -bsimilarity * XY_dot / (_similarity * _similarity) bXY_dot = _bXY_dot b_similarity = _b_similarity _bX_norm = tangent.unbroadcast(b_similarity * Y_norm, X_norm) _bY_norm = tangent.unbroadcast(b_similarity * X_norm, Y_norm) bX_norm = _bX_norm bY_norm = _bY_norm return bXY_dot, bX_norm, bY_norm
```

Cleaning this up and moving it into our callback-based style, we get:

```
def tangent_backprop_cosine_division(XY_dot, X_norm, Y_norm, d_loss=1.0): similarity = XY_dot / _similarity loss = 1 - similarity
# Grad of: loss = 1 - similarity d_similarity = -d_loss # Grad of: similarity = XY_dot / (X_norm * Y_norm) d_XY_dot = d_similarity / similarity d_similarity = -d_similarity * XY_dot / (similarity * similarity) d_X_norm = d_similarity * Y_norm d_Y_norm = d_similarity * X_norm return dXY_dot, dX_norm, dY_norm
def cosine_divsion_v2(XY_dot, X_norm, Y_norm): similarity = XY_dot / (X_norm * Y_norm) loss = 1 - similarity
def get_d_XY_dot(d_loss): return -d_loss / similarity
def get_dX_norm(d_loss): d_similarity = d_loss * XY_dot / (similarity * similarity) return d_similarity * Y_norm
def get_dY_norm(d_loss): d_similarity = d_loss * XY_dot / (similarity * similarity) return d_similarity * X_norm
return loss, (get_d_XY_dot, get_dX_norm, get_dY_norm)
```

Finally, let’s put it all together. I’ll ignore the calculations for dY, and just provide the gradient with respect to dX.

```
def cosine_loss_v3(X, Y): XY_dot = (X*Y).sum() X_norm = np.sqrt((X*X).sum()) Y_norm = np.sqrt((Y*Y).sum()) similarity = XY_dot / (X_norm * Y_norm) loss = 1 - similarity
def get_dX(d_loss): d_XY_dot = -d_loss / similarity d_X_norm = d_loss * XY_dot / (similarity ** 2) dX = d_XY_dot * Y + d_X_norm * (2.0 * X_norm) # dX = (X / (X_norm*Y_norm)) - (similarity * (X / X_norm**2)) return dX return loss, get_dX
loss, get_d_vec1 = cosine_loss_v3(vec1, vec2)my_grad = get_d_vec1(1.0)numpy.testing.assert_almost_equal(tangent_grad, my_grad)
```