Optimize masking parameter

the-dharma-bum: <p>Hi,<br>
I’d like to optimize a threshold value, based on which I’ll mask a Tensor, which itself will be used to compute the final loss.<br>
Here is a dummy example, where I generate 1000 3D points, and minimize the mean norm of points whose norm is below a threshold:</p>
<pre data-code-wrap="python"><code class="lang-python">X = torch.randn(1000, 3, requires_grad=True)
threshold = torch.tensor(2., requires_grad=True)

loss = X[X.norm(dim=1) &lt; threshold].norm(dim=1).mean()
loss.backward()

assert threshold.grad is None
</code></pre>
<p>Here the threshold can’t be updated by an optimizer as its gradient is None. I’m sure this can easily be rewritten so as to compute the gradient w.r.t the mask but I can’t figure how ?</p>

KFrank: <p>Hi Luc!</p>
<aside class="quote no-group quote-modified" data-username="the-dharma-bum" data-post="1" data-topic="204381" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/the-dharma-bum/48/32575_2.png" class="avatar"> the-dharma-bum:</div>
<blockquote>
<p>I’d like to optimize a threshold value, based on which I’ll mask a Tensor,<br>
…<br>
Here the threshold can’t be updated by an optimizer as its gradient is None.</p>
</blockquote>
</aside>
<p>The short story is use <code>sigmoid()</code> to create a “soft” mask.</p>
<p>The problem is that the threshold-masking operation is not (usefully)<br>
differentiable.</p>
<p>As you vary <code>threshold</code> over some range, the same masked elements<br>
of your tensor are kept, so <code>loss</code> is constant over this range of<br>
<code>threshold</code>.  Mathematically, over this range, the gradient of <code>loss</code><br>
is zero, which isn’t useful for gradient descent (and pytorch is smart<br>
enough not to compute this not-useful gradient).</p>
<p>At the specific (discrete) values of <code>threshold</code> where the set of masked<br>
elements changes, the gradient is mathematically undefined (or <code>inf</code>, if<br>
you prefer).  This is also not useful (and in any event, is a set of measure<br>
zero).</p>
<p>Instead, you want to smoothly turn the masked elements on and off.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.3.1
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2024)
&gt;&gt;&gt;
&gt;&gt;&gt; X = torch.randn (10, 3, requires_grad = True)
&gt;&gt;&gt; threshold = torch.tensor (1.5, requires_grad = True)
&gt;&gt;&gt;
&gt;&gt;&gt; loss = X[X.norm (dim = 1) &lt; threshold].norm (dim = 1).mean()
&gt;&gt;&gt; loss
tensor(1.0206, grad_fn=&lt;MeanBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; hard = 10.0                                            # larger values make the mask "harder"
&gt;&gt;&gt; X_norm = X.norm (dim = 1)
&gt;&gt;&gt; soft_mask = torch.sigmoid (hard * (threshold - X_norm))
&gt;&gt;&gt; soft_mask
tensor([9.0025e-03, 9.9987e-01, 8.5999e-01, 9.0971e-07, 9.6722e-01, 9.2340e-08,
        9.6193e-03, 9.9644e-01, 9.9148e-01, 9.8524e-01],
       grad_fn=&lt;SigmoidBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; lossB = (soft_mask * X_norm).sum() / soft_mask.sum()   # weighted mean
&gt;&gt;&gt; lossB
tensor(1.0156, grad_fn=&lt;DivBackward0&gt;)
&gt;&gt;&gt;
&gt;&gt;&gt; lossB.backward()
&gt;&gt;&gt; threshold.grad
tensor(0.1020)
</code></pre>
<p>Here we smoothly turn the masked elements on and off by multiplying<br>
them with <code>soft_mask</code> that contains values that are (usually) close to<br>
zero or one.</p>
<p>As the parameter <code>hard</code> is increased, in principle to infinity, the elements<br>
of <code>soft_mask</code> become zero and one, <code>soft_mask</code> becomes effectively<br>
the same as your “hard” boolean mask, and <code>lossB</code> becomes equal to<br>
<code>loss</code>.</p>
<p>However, the larger you make <code>hard</code>, the smaller the range becomes<br>
over which the gradient of loss with respect to <code>threshold</code> differs from<br>
zero enough to be of practical use.  (And even is when the gradient is<br>
mathematically non-zero, it can underflow to zero.)</p>
<p>Think of <code>soft_mask</code> as a differentiable proxy for (approximation to)<br>
your hard boolean mask.</p>
<p>It’s up to your use case how hard or soft you want <code>soft_mask</code> to be and<br>
what value you should use for the parameter <code>hard</code>.</p>
<p>Best.</p>
<p>K. Frank</p>