Zygote - parametrize matrix such that gradient is only performed on selected coefficients

Specific Domains Machine Learning
, , , ,
1

I was wondering if it was possible to parametrize matrices such that zygote’s gradient would only be performed on the selected coefficients of the matrix.

Some dummy code would look like:

theta = randn(2)
W = [theta[1] theta[2]; 0.0 theta[1]]
b = randn(100)
x = randn(2, 100)

gs = gradient(w -> sum(w*x+b), W)

However the gradient would only be performed on the values of theta, keeping W[2, 1] set to 0.0 for the entire gradient descent. Also this would imply that the gradients for theta[1] and theta[2] be equal at all times.
Am I completely going at this wrong way? or is it just not possible in julia?

2

You have the right idea, but you need to take the gradient w.r.t. theta instead of W, just wrap the function that accepts W in a function that accepts theta instead.

2 Likes