Chain rule in autograd / computation graph using .clone()

daaawit: <p>Hi all,</p>
<p>I’m trying to keep this post somewhat short but I’m happy to provide additional details if necessary. I’m trying to implement ProsPr pruning as proposed by <a href="https://arxiv.org/abs/2202.08132" rel="noopener nofollow ugc">Alizadeh et al. (2022)</a>. The basic idea is to add a masking <em>c</em> (i.e. matrix of all ones) to all to-be-pruned layers at initialization and then perform <em>M</em> batches of training (5 at most), as visualized by figure 1 in the paper:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/5/a5204b696a7878ce17aaba391471f9c81bbc76ea.png" data-download-href="https://discuss.pytorch.org/uploads/default/a5204b696a7878ce17aaba391471f9c81bbc76ea" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5204b696a7878ce17aaba391471f9c81bbc76ea_2_690x222.png" alt="image" data-base62-sha1="nyLYlTf3iWsxWlTzYuJORD5vVq2" width="690" height="222" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5204b696a7878ce17aaba391471f9c81bbc76ea_2_690x222.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5204b696a7878ce17aaba391471f9c81bbc76ea_2_1035x333.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/5/a5204b696a7878ce17aaba391471f9c81bbc76ea_2_1380x444.png 2x" data-dominant-color="F7F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1630×526 61.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>After the <em>M</em> iterations, the derivative for each mask element is computed, which requires using the chain rule, as illustrated by equation 6-10 in the original paper:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/2/e29a9d92596af9785ab8786ee63e8199f9f7f9da.png" data-download-href="https://discuss.pytorch.org/uploads/default/e29a9d92596af9785ab8786ee63e8199f9f7f9da" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e29a9d92596af9785ab8786ee63e8199f9f7f9da_2_690x222.png" alt="image" data-base62-sha1="wkDazwQvowvqTSadYYIWU4FJiZA" width="690" height="222" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e29a9d92596af9785ab8786ee63e8199f9f7f9da_2_690x222.png, https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e29a9d92596af9785ab8786ee63e8199f9f7f9da_2_1035x333.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/e/2/e29a9d92596af9785ab8786ee63e8199f9f7f9da_2_1380x444.png 2x" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1682×542 38.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The <a href="https://github.com/mil-ad/prospr/blob/a92177989f4480f1f2b43a48b3e18a6597ebba6d/prospr/prospr.py" rel="noopener nofollow ugc">original implementation</a> achieves this via monkey-patching and re-implementing SGD with a lot of additional storage overhead. I would like to avoid that, if possible.</p>
<p>My idea is to add the weight masks as parameters and set their learning rates to 0, since the masks need to stay constant but also need to have a gradient per step (see equation (10)). I achieves this by initializing the optimizer as follows:</p>
<pre><code>    weights = [
        param for name, param in model.named_parameters() if name.endswith(".weight")
    ]
    weight_masks = [
        param for name, param in model.named_parameters() if name.endswith(".weight_mask")
    ]

    # Initialize optimizer
    inner_sgd: SGD = SGD(
        [{"params": weights}, {"params": weight_masks, "lr": 0.0}],
        lr=learning_rate,
        momentum=learning_momentum,
    )
</code></pre>
<p>During training, I then use <code>autograd</code> to be able to create/retain the computation graph:</p>
<pre><code>    grads = torch.autograd.grad(
        outputs=loss, 
        inputs=self.model.parameters(), 
        create_graph=True, 
        retain_graph=True
    )
</code></pre>
<p>I then clone each gradient to the respective layer using a for loop and <code>grads[i].clone()</code>, since autograd does not apply the gradients (unlike <code>optimizer.step()</code>).</p>
<p>From what I’ve read on the computation graph, this should allow me to implement the chain rules as given by equation (6) by simply applying autograd after the training has finished:</p>
<pre><code>    torch.autograd.grad(loss, weight_masks, create_graph=False, retain_graph=True)
</code></pre>
<p>However, since I do not understand the inner workings of the graph completely (yet), I am not sure whether my assumptions and implementation are correct. As reference: My results diverge slightly from the original implementation, but the differences in computed thresholds are small enough that one might consider them to be related to numerical differences.</p>
<p>I would appreciate any input on whether this actually makes sense or whether I did something wrong. Thanks in advance!</p>