So I mailed the authors, and they were kind enough to reply.
The reasoning behind using gradients instead of matrix multiplications is the convenience we get in calculating gradients. The convenience is not much if you are just dealing with linear layers. For example: consider the following code
import torch
torch.manual_seed(123)
fc1 = torch.nn.Linear(4,2,bias=False) # disabling bias is important
i = torch.rand(4, requires_grad=True) # activations from previous layer aLminus1
o = fc1.forward(i) # notice this is w*aLminus1 without activation
out = torch.relu(o) # this is activation of aL layer, which is also our output layer, hence they become Relevance values.
print('activation j',i)
print('forward k',o)
print('activation k',out)
print('w_jk',fc1.weight)
print('bias_k',fc1.bias)
print()
w = fc1.weight.data
Rk = aL = out.data
aLminus1 = i.data
z = aLminus1.unsqueeze(0).numpy().dot(w.T.numpy())
s = Rk/z
c = s.numpy().dot(w.numpy())
print("With matrix multiplication")
print(c)
z = 1e-9 + o
s = Rk/(z+1e-9)
(z*s.data).sum().backward()
c = i.grad
print('With gradients method')
print(c)
You will see that both the method gives the same answer. But the real convenience is seen when you are dealing with convolutional layers. Without the gradient method, you will need to write an iterative code to pass the relevance values from kth layer to jth layer because c = s.numpy().dot(w.numpy())
will no longer hold for convolutional layers since shape of w will be [in_maps, filter_size, filter_size]
.
Okay but then why this gradient method works for linear layers? Well the derivation is quite simple.
target: c_i = w @ s
; where shape(w) = [ j , k ]
, shape(s) = [ k , 1 ]
, and @ is matrix multiplication
let us do this for simplicity, w = rho(w)
- we know that
z = w.forward(a)
which is equivalent to w.T @ a
where shape(a) = [ j , 1 ]
, hence shape(z) = [ k , 1 ]
s = Rk/z
where shape(Rk) = [ k , 1 ]
and hence shape(s) = [ k , 1 ]
- Now if we use gradients, then
temp = z*s.data
is equivalent to ( w.T @ a ) * s.data
where * is element-wise multiplication and shape(temp) = [ k , 1 ]
- when this expression is summed and gradients are back-propagated using
.backward()
The gradients will be attached to w and a both
a.grad
will be the differentiation of temp with respect to a and treating others as constant
- hence
a.grad
will be w @ s
why? For example: if f = x*y*z
then df/dx = y*z
i.e. d(temp)/da = w @ s
it is easy to observe this on pen-paper, just put 1 in all the values of a in temp, and we will get w @ s
- Hence
c_i = a.grad = w @ s
- The reason why
s.data
is used instead of simply s
is because, if we use simply s
then gradients will find a new path to flow into a because s
is dependent on a
.
However, be mindful of one small thing. If you are trying to see for yourself whether gradient based method and matrix multiplication method gives the same values of $c_i$, then make sure your linear layers have bias set to False
. Otherwise both the methods will give different values to $c_i$. Why? Bias changes the way how you should calculate $z_{jk}$, hence the derivation changes.