Zero'ing one input's gradient for Matrix Multiply

username: kgbounce
cooked: <p>Hi! the @ operator, or matrix multiply, is stateless and accepts 2 input tensors. During the backprop, my understanding is that it’ll calculate two gradients w.r.t. the 2 input tensors, which will each update any variables via the chain rule along the paths that produce them, respectively.</p>
<p>In my experiment, I want to zero out the gradient for only one input tensor and keep the other as-is. My guess is that I need to use autograd, but I’m fairly new to this and the toy example (<a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html" class="inline-onebox" rel="noopener nofollow ugc">PyTorch: Defining New autograd Functions — PyTorch Tutorials 1.7.0 documentation</a>) doesn’t seem suffice. Please help!</p>

username: kgbounce
cooked: <p>I figured I could use the tutorial to add a new custom ‘Identity’ layer that passes through activations in forward prop but zeros out the gradient in back prop. But I’m really curious about how to do it properly with bmm (@). Thanks.</p>

username: kgbounce
cooked: <p>It turns out that I also need to modify the gradient of the other operand. Is there any example or tutorial that shows how to use Autograd function with an operation that produces multiple gradients? TIA.</p>

username: kgbounce
cooked: <pre><code class="lang-auto">class FwdOnlyIdentity(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # pass the input feature map unmodified, as Identity
        return input
    
    @staticmethod
    def backward(ctx, grad_output):
        zero_grad_input = torch.zeros_like(grad_output) # set gradient w.r.t. input to 0's
        return zero_grad_input
</code></pre>
<p>Above is that I implemented for what I originally asked about.</p>
<p>Then, instead of directly doing <code>a @ b</code>, I added <code>self.fwd_identity = FwdOnlyIdentity.apply</code>. Then in the forward function, the computation becomes <code>self.fwd_identity(a) @ b</code>. Does this make sense?</p>

username: googlebot
cooked: <p>in the simplest case, you can just do something like x.detach() @ weights</p>

username: kgbounce
cooked: <p>The matrix multiply here doesn’t have any trainable weights. It takes two activation tensors as input.</p>

username: googlebot
cooked: <p>that’s arbitrary variable names, this cuts off the whole computation chain that produced “x”, if you return a gradient of zeros, zeros will [back] propagate further, so it is the same as blocking with detach()</p>

username: kgbounce
cooked: <p>I see. What makes it complicated is that I cannot simply detach the producer of ‘x’, as it’s expected to receive gradient update from another downstream path.</p>

username: googlebot
cooked: <p>another path won’t be affected, you’ll have separate graph nodes for each use of tensor, i.e.</p>
<p>x_ng = x.detach()<br>
y1 = x_ng @ m1<br>
y2 = x @ m2</p>
<p>two paths start from x, original x tensor is not effected by detach, and y1 is trainable only through m1</p>

username: kgbounce
cooked: <p>Hey, I figured out how to do it with <code>autograd.Function</code> after reading through the pages. Let me know if it makes sense.</p>
<pre><code class="lang-auto">class TwoOpGradientController(torch.autograd.Function):
    @staticmethod
    def forward(ctx, *args):
        ctx.args_0_shape = args[0].shape
        return args[0] @ args[1]
    
    @staticmethod
    def backward(ctx, grad_output):
        lhs_zero_grad_input = torch.zeros_like(ctx.args_0_shape)
        rhs_grad_output = # do weird stuff to rhs gradients
        
        return lhs_zero_grad_input, rhs_grad_output
</code></pre>

username: googlebot
cooked: <p>it can be correct but doesn’t make much sense to me, that you’re tying to matmul. you have to at least start from the original formula (IIRC, rhs_grad_output = grad_output @ args[0].t()) if you have a reason to have such a Function.</p>

username: kgbounce
cooked: <p>Yes, that’s a good point. The toy example is just to show how the syntax works, as the existing examples typically show simpler functions, such as ReLU and Exp.</p>
<p>BTW, I think the original formula is <code>args[0].t() @ grad_output</code>, such that the resulting shape is correct ( k by m * m by n → k by n).</p>