Masking module parameters

Mikele: <p>Hello,</p>
<p>how would one effectively mask the parameters of a module without losing neither their link to the optimizer?<br>
More in detail:</p>
<p>Approach 1:</p>
<pre><code class="lang-auto">for weights in model.parameters():
    backups.append(weights.clone().detach().data)
    mask = sample_mask(some_arguments...)  # this comes from weights of a wrapper module that needs to be trained jointly
    weights.data *= mask
loss = loss_func(model(data), target)
for backup, weights in zip(backups, model.parameters()):
    weights.data = backup
</code></pre>
<p>This seemed good, but the weights of <code>model</code> lose the grad_fn attribute that the mask accumulated for its generation process and that I need for my training!</p>
<p>Approach 2:</p>
<pre><code class="lang-auto">for weights in model.parameters():
    backup.append(weights.clone().detach().data)
    mask = sample_mask(some_arguments...)
    weights.data *= mask
    weights.grad_fn = mask.grad_fn  # ERROR, grad_fn is not writable!
loss = loss_func(model(data), target)
</code></pre>
<p>Doesn’t work as pytorch forbids assigning to grad_fn (why? It may not be common or best practice, but certainly has its use cases). If it did work, it would be my preferred approach i guess, as it makes clear what I need as well.</p>
<p>Approach 3:</p>
<pre><code class="lang-auto">for weights in model.parameters():
    del weights from optimizer
    mask = sample_mask(some_arguments...)
    weights = weights * mask
    optimizer.add_weights(weights) in the right place somehow
loss = loss_func(model(data), target)
</code></pre>
<p>Basically very hacky and unclear as to how to delete only the exact weights that I want to see deleted, and then afterwards assign them in the correct spot again…</p>
<p>So since approach 3 appears to be pretty unthinkable for a bigger composite model, I am a bit at a loss at how to simply mask parameters and still train the parameters and the mask generation process jointly. Tips would be appreciated</p>

albanD: <p>Hi,</p>
<p>The <code>nn.Parameter</code> are <strong>always</strong> leaf nodes. Meaning that they cannot have <code>.grad_fn</code> field set.<br>
This is because the optimizers use the <code>.grad</code> field to perform updates on the weights, but if it’s not a leaf Tensor, its <code>.grad</code> field won’t be updated and so the optimizer won’t be able to perform update.<br>
So do you actually want to perform updates both on the weights and parameters (naming from your code)? Or just parameters?</p>

Mikele: <p>Hey thanks for the reply,</p>
<p>sorry the naming was confusing. I want gradient updates on:</p>
<ol>
<li>The parameters behind the mask sampling process (to better sample the appropriate mask)</li>
<li>the weights of the underlying model (here named simply <code>model</code> with its weights <code>weights</code>), that are masked by <code>mask</code>.</li>
</ol>
<p>I will edit the confusing double parameters naming part from the code quickly.<br>
Does this clear it up?</p>

albanD: <p>So here you have two parameters in your module:</p>
<ul>
<li>original weights of the module</li>
<li>mask_params that are used to compute the mask</li>
</ul>
<p>I would modify the module to have all the right Parameters and recompute weight for each forward.</p>
<pre><code class="lang-python"># Example for a Linear (handle bias the same way if you want them)
mod = nn.Linear(10, 10, bias=False)
mod.mask_params = nn.Parameter(whatever)
mod.original_weight = mod.weight
del mod.weight # re-populate it for each forward using a forward pre hook
def repopulate_weight(mod, _):
  mask = sample_mask(mod.mask_params)
  mod.weights = mod.original_weight * mask
mod.register_forward_pre_hook(repopulate_weight)

# Use `mod` as any Linear now
</code></pre>

Mikele: <p>Hey alban,<br>
thanks a lot for the help. The hooks were a good idea. Ive now set a forward pre hook to sample and backup the .data and a backward hook to reinstall the backups for gradient updates.</p>