Create custom gradient descent in pytorch
I am trying to use PyTorch autograd to implement my own batch gradient descent algorithm. I want to create a simple one-layer neural net with a linear activation function and the mean squared error as the loss function. I can't seem to get my head around what exactly is happening in the backward pass and how PyTorch understands my outputs. I have coded one class specifying the linear function in the forward pass, and in the backward pass, I calculated the gradients with respect to each variable. I also coded a class for the MSE function and specified the gradients with respect to ITS variables in the backward pass. When I run a simple gradient descent algorithm, I get no errors, but the MSE only goes down in the first iteration, and after that, it continually goes up. This leads me to believe that I have made a mistake, but I am not sure, where. Does anybody see the error in my code? Also, if somebody could explain to me what exactly the grad_output stands for, that would be amazing.
Here are the functions:
import torch
from torch.autograd import Function
from torch.autograd import gradcheck


class Context:
    def __init__(self):
        self._saved_tensors = ()

    def save_for_backward(self, *args):
        self._saved_tensors = args

    @property
    def saved_tensors(self):
        return self._saved_tensors

class MSE(Function):
    @staticmethod
    def forward(ctx, yhat, y):
        ctx.save_for_backward(yhat, y)
        q = yhat.size()[0]
        mse = torch.sum((yhat-y)**2)/q
        return mse

    @staticmethod
    def backward(ctx, grad_output):
        yhat, y = ctx.saved_tensors
        q = yhat.size()[0]
        return 2*grad_output*(yhat-y)/q, -2*grad_output*(yhat-y)/q  

class Linear(Function):
    @staticmethod
    def forward(ctx, X, W, b):
        rows = X.size()[0]
        yhat = torch.mm(X,W) + b.repeat(rows,1)
        ctx.save_for_backward(yhat, X, W)
        return yhat

    @staticmethod
    def backward(ctx, grad_output):
        yhat, X, W = ctx.saved_tensors
        q = yhat.size()[0]
        p = yhat.size()[1]
        return torch.transpose(X, 0, 1), W, torch.ones(p)

And here is my gradient descent:
import torch
from torch.utils.tensorboard import SummaryWriter       
from tp1moi import MSE, Linear, Context

 
x = torch.randn(50, 13)         
y = torch.randn(50, 3)

w = torch.randn(13, 3, requires_grad=True)
b = torch.randn(3, requires_grad=True)

epsilon = 0.05

writer = SummaryWriter()
for n_iter in range(100):
    linear = Linear.apply
    mse = MSE.apply
    loss = mse(linear(x, w, b), y)

    writer.add_scalar('Loss/train', loss, n_iter)

    print(f&quot;Itérations {n_iter}: loss {loss}&quot;)

    loss.backward()

    with torch.no_grad(): 
        w -= epsilon*w.grad
        b -= epsilon*b.grad
        w.grad.zero_()
        b.grad.zero_()

Here is one output I got (they all look similar to this one):
Itérations 0: loss 72.99712371826172
Itérations 1: loss 7.509067535400391
Itérations 2: loss 7.309497833251953
Itérations 3: loss 7.124927997589111
Itérations 4: loss 6.955358982086182
Itérations 5: loss 6.800788402557373
Itérations 6: loss 6.661219596862793
Itérations 7: loss 6.536648750305176
Itérations 8: loss 6.427078723907471
Itérations 9: loss 6.3325090408325195
Itérations 10: loss 6.252938747406006
Itérations 11: loss 6.188369274139404
Itérations 12: loss 6.138798713684082
Itérations 13: loss 6.104228973388672
Itérations 14: loss 6.084658145904541
Itérations 15: loss 6.0800886154174805
Itérations 16: loss 6.090517520904541
Itérations 17: loss 6.115947723388672
Itérations 18: loss 6.156377792358398
Itérations 19: loss 6.2118072509765625
Itérations 20: loss 6.2822370529174805
Itérations 21: loss 6.367666721343994
Itérations 22: loss 6.468096733093262
Itérations 23: loss 6.583526611328125
Itérations 24: loss 6.713956356048584
Itérations 25: loss 6.859385967254639
Itérations 26: loss 7.019815444946289
Itérations 27: loss 7.195245742797852
Itérations 28: loss 7.385674953460693
Itérations 29: loss 7.591104507446289
Itérations 30: loss 7.811534881591797
Itérations 31: loss 8.046965599060059
Itérations 32: loss 8.297393798828125
Itérations 33: loss 8.562823295593262
Itérations 34: loss 8.843254089355469
Itérations 35: loss 9.138683319091797
Itérations 36: loss 9.449112892150879
Itérations 37: loss 9.774543762207031
Itérations 38: loss 10.114972114562988
Itérations 39: loss 10.470401763916016
Itérations 40: loss 10.840831756591797
Itérations 41: loss 11.226261138916016
Itérations 42: loss 11.626690864562988
Itérations 43: loss 12.042119979858398
Itérations 44: loss 12.472548484802246
Itérations 45: loss 12.917980194091797
Itérations 46: loss 13.378408432006836
Itérations 47: loss 13.853838920593262
Itérations 48: loss 14.344267845153809
Itérations 49: loss 14.849695205688477
Itérations 50: loss 15.370124816894531
Itérations 51: loss 15.905555725097656
Itérations 52: loss 16.455984115600586
Itérations 53: loss 17.02141571044922
Itérations 54: loss 17.601844787597656
Itérations 55: loss 18.19727325439453
Itérations 56: loss 18.807701110839844
Itérations 57: loss 19.43313217163086
Itérations 58: loss 20.07356071472168
Itérations 59: loss 20.728988647460938
Itérations 60: loss 21.3994197845459
Itérations 61: loss 22.084848403930664
Itérations 62: loss 22.7852783203125
Itérations 63: loss 23.50070571899414
Itérations 64: loss 24.23113441467285
Itérations 65: loss 24.9765625
Itérations 66: loss 25.73699188232422
Itérations 67: loss 26.512422561645508
Itérations 68: loss 27.302854537963867
Itérations 69: loss 28.108285903930664
Itérations 70: loss 28.9287166595459
Itérations 71: loss 29.764144897460938
Itérations 72: loss 30.614578247070312
Itérations 73: loss 31.48000717163086
Itérations 74: loss 32.36043930053711
Itérations 75: loss 33.2558708190918
Itérations 76: loss 34.16630172729492
Itérations 77: loss 35.091732025146484
Itérations 78: loss 36.032161712646484
Itérations 79: loss 36.98759460449219
Itérations 80: loss 37.95802307128906
Itérations 81: loss 38.943458557128906
Itérations 82: loss 39.943885803222656
Itérations 83: loss 40.959320068359375
Itérations 84: loss 41.98974609375
Itérations 85: loss 43.03517532348633
Itérations 86: loss 44.09561538696289
Itérations 87: loss 45.171043395996094
Itérations 88: loss 46.261474609375
Itérations 89: loss 47.366905212402344
Itérations 90: loss 48.487335205078125
Itérations 91: loss 49.62276840209961
Itérations 92: loss 50.773197174072266
Itérations 93: loss 51.93863296508789
Itérations 94: loss 53.11906433105469
Itérations 95: loss 54.31448745727539
Itérations 96: loss 55.524925231933594
Itérations 97: loss 56.75035095214844
Itérations 98: loss 57.990787506103516
Itérations 99: loss 59.2462158203125```



Let's take a look at the implementation of MSE, the forward pass will be MSE(y, y_hat) = (y_hat-y)² which is straightforward. For the backward pass, we are looking to compute the derivative of the output with regards to the input, as well as the derivative with regards to each of the parameters. Here MSE does not have any learned parameters, so we just want to compute dMSE/dy*dz/dMSE using the chain rule, which is d(y_hat-y)²/dy*dz/dMSE, i.e. -2(y_hat-y)*dz/dMSE. Not to confuse you here: I wrote dz/dMSEas the incoming gradient. It corresponds to the gradient following backward towards the MSE layer. From your notation grad_output is dz/dMSE. Therefore the backward pass is simply -2*(y_hat-y)*grad_output. Then normalized by the batch size q, retrieved from y_hat.size(0).
The same thing goes with the Linear layer. It will involve some more computation since, this time, the layer is parametrized by w and b. The forward pass is essentially x@w + b. While the backward pass, consists in calculating dz/dx, dz/dw, and dz/db. Writing f as x@w + b. After some work you can find that that:

dz/dx = d(x@w + b)/dx * dz/df = dz/df*W.T,
dz/dw = d(x@w + b)/dw * dz/df = X.T*dz/df,
dz/db = d(x@w + b)/db * dz/df = 1.

In terms of implementation this would look like:

output_grad@w.T for the gradient w.r.t x,
x.T@output_grad for the gradient w.r.t w,
torch.ones_like(b) for the gradient w.r.t b.

