I am trying to implement Direct Feedback Alignment with a custom autograd function, a diagram of the algorithm can be found here in subplot c. Basically what I want to do (for example in a Linear layer) is changing the gradient computation in backward function from
grad_input = grad_output.mm(weight)
grad_weight = grad_output.t().mm(input)
to
grad_input = grad_output # Directly pass down grad_output without modifiying
grad_weight = grad_output.mm(weight_fb).t().mm(input)
But this won’t work as the grad_input does not have the corresponding size to the input.
I was wondering if there is a way to implement this with a custom autograd function. Any hint would be appreciated!