Gradient descent for linear regression using Zygote AutoDiff

ad_linreg.jl

# Gradient descent with autodiff for linear regression
using Zygote

# Data
X = randn(1000, 10)
b = (1:10)
y = X * b + randn(1000)

# MSE for linear model
function mse(bhat)
    res = y - X * bhat
    res'res / length(res)
end

# Gradient w.r.t. parameters
grad(bhat) = gradient(mse, bhat)[1]

# Initialize params & control
bhat  = zeros(10)
step  = 1.0e-3
maxit = 1_000_000
tol   = eps(1.0)

# gradient descent
for i = 1:maxit
    global bhat
    Δ = step .* grad(bhat)
    if (Δ'Δ < tol) 
        break
    end
    bhat -= Δ
    print("loss ", i, " ", mse(bhat), "\n")
end