How can I control gradient computation in cascaded operation with shared parameters

prclibo: <p>Hello,</p>
<p>I was designing some structure with cascaded operations sharing same parameters, i.e., something like <code>loss = g(f(x, w), w)</code>. Normally in a backward pass, both <code>f</code> and <code>g</code> will contribute gradients to <code>w</code>.  I was wondering if I am able to control the backward pass so that only one of <code>f</code> or <code>g</code> contributes to the gradient of <code>w</code>.</p>
<p>I know a naive say of doing this is to copy <code>w</code> as <code>w1</code> and operate by <code>loss = g(f(x, w), w1)</code>. However, since <code>f</code> and <code>g</code> might be a complicated network and <code>w</code> be a series of parameters, I was wondering if we can do this without copying parameters but by just setting some flags.</p>

soulitzer: <p>It is possible to do this by detaching before one of the uses i.e. <code>g(f(x, w.detach()), w)</code>. I’m not sure why you would want to do this though because your gradient would be wrong. To consider an example: If f, g are both mul operations, your whole function would be <code>xw^2</code> and <code>dw</code> would be <code>2wx</code>, but if you detach before one of the uses, one of the w’s would be treated as a constant, and your <code>dw</code> would be <code>wx</code> a factor of 2 off from the correct gradient.</p>

prclibo: <p>Thank you for the reply!</p>
<p>Yes, detaching w could be a solution. But as mentioned in the original post, <code>f</code> and <code>g</code> might be a complicated network and contain many layers and <code>w</code> parameters. I was wondering if there is a way to detach all the parameters of a network in just one operation?</p>
<p>My use case comes from invertible networks. The computation is of a form of <code>f_w^{-1} (h(f_w(x)))</code> and I would like to observe how will <code>f^{-1}</code> or <code>f</code> contribute to the training respectively. I would be great if you know any references on similar bidirectional training and could share with me:)</p>
<aside class="quote no-group" data-username="soulitzer" data-post="2" data-topic="156853" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/48.png" class="avatar"> soulitzer:</div>
<blockquote>
<p>It is possible to do this by detaching before one of the uses i.e. <code>g(f(x, w.detach()), w)</code>. I’m not sure why you would want to do this though because your gradient would be wrong. To consider an example: If f, g are both mul operations, your whole function would be <code>xw^2</code> and <code>dw</code> would be <code>2wx</code>, but if you detach before one of the uses, one of the w’s would be treated as a constant, and your <code>dw</code> would be <code>wx</code> a factor of 2 off from the correct gradient.</p>
</blockquote>
</aside>

soulitzer: <p>Perhaps the easiest way to do this is to just have two copies of each <code>f</code> or <code>g</code> module one as-is and one with all the parameters detached, and just swap in the detached one whenever you want to only one to contribute to the backward pass during the next backward. (If you are asking about how to iterate through the parameters of a module, you can use <code>mod.parameters()</code> which will iterate recursively through the current module as well as its submodule’s parameters.)</p>
<p>Invertible networks are pretty cool, are you trying to use them for the purposes of memory reduction or normalizing flows or some other application?</p>

prclibo: <aside class="quote no-group" data-username="soulitzer" data-post="4" data-topic="156853">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/48.png" class="avatar"> soulitzer:</div>
<blockquote>
<p>Invertible networks are pretty cool, are you trying to us</p>
</blockquote>
</aside>
<p>Thanks. Copying <code>f</code> and <code>g</code> is a baseline solution. I will do it this way if no better solution is found.</p>
<p>Yes. I am playing some normalizing-flow-like structures and would like to see if training it in a bidirectional fashion of both GAN and BPD will produce any new discovery (perhaps similar to FlowGAN). Let me know if you have any experience on this:)</p>