Apply mask softmax

Diego999: <p>Hi everyone,</p>
<p>I try to implement the following function:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/e/edf484646983459fed873d72b26a4e43a21dc421.jpg" data-download-href="https://discuss.pytorch.org/uploads/default/edf484646983459fed873d72b26a4e43a21dc421" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/e/edf484646983459fed873d72b26a4e43a21dc421_2_690x133.jpg" alt="image" data-base62-sha1="xX32eJ9BzMk6ljcNHAQWMV4V37b" width="690" height="133" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/e/edf484646983459fed873d72b26a4e43a21dc421_2_690x133.jpg, https://discuss.pytorch.org/uploads/default/original/2X/e/edf484646983459fed873d72b26a4e43a21dc421.jpg 1.5x, https://discuss.pytorch.org/uploads/default/original/2X/e/edf484646983459fed873d72b26a4e43a21dc421.jpg 2x" data-dominant-color="F4F4F4"><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">1014×196 24.5 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>At this stage, I have e.g. a tensor <code>[[1,0,3], [0, 1, 2], [3, 2, 1]]</code> and the softmax should be done for only values greater than 0 (neighbors), thus <code>[[0.12, 0, 0.88], [0, 0.27, 0.73], [0.66, 0.24, 0.1]]</code>.</p>
<p>Is there an existing method to do this or does anyone have an idea ?</p>
<p>Thank you for your help !</p>

jpeg729: <p>To clarify: you want to calculate the standard softmax BUT you want to ignore any zero values.</p>
<p>In other words for the first row you remove the zero, then you calculate<code>softmax([1,3])</code>, and then you reinsert the zero into the output.</p>
<p>Is that right?</p>
<p>If so, I am stumped. Maybe someone else can help.</p>

Diego999: <p>Yes exactly ! The zeros have been obtained using torch.where.</p>

Diego999: <p>I thought about using -10^10 instead of 0 for the soft max but there should be a better way I guess</p>

jpeg729: <p>I think that is a great idea. Anything else would require funny indexing and would probably end up being super inefficient.</p>

Andy-jqa: <p>Hi Diego. Does this approach work well?</p>

Diego999: <p>Hi Andy-jqa</p>
<p>Yes absolutely !</p>

Andy-jqa: <p>Thanks for the reply.<br>
I did something like this:</p>
<pre><code class="lang-python"># matrix A is the one you want to do mask softmax at dim=1
A_max = torch.max(A,dim=1,keepdim=True)[0]
A_exp = torch.exp(A-A_max)
A_exp = A_exp * (A == 0).type(torch.FloatTensor) # this step masks
A_softmax = A_exp / torch.sum(A_exp,dim=1,keepdim=True)
</code></pre>
<p>which also seems to work.</p>

jpeg729: <p>Instead of <code>.type(torch.FloatTensor)</code> which will move the data to CPU if it was on the GPU, I would do <code>.float()</code> which will convert it while leaving it on the device it was on, as well as being shorter to type and easier to read.</p>

Andy-jqa: <p>Great advice! Thanks.</p>

Ujan_Deb: <p>In case you cannot make any assumptions about your data, you could use  <code>float('-inf')</code></p>

MrTuo: <p>Thanks a lot!<br>
And I have a little suggestion:<br>
<code>A_softmax = A_exp /(torch.sum(A_exp,dim=1,keepdim=True)+epsilon)</code><br>
It can avoid division by zero zero.</p>

krylea: <p>I had to implement something similar. My approach was the following (where mask is a tensor of 1s and 0s indicating the entries to be removed):</p>
<pre><code class="lang-auto">def masked_softmax(vec, mask, dim=1):
    masked_vec = vec * mask.float()
    max_vec = torch.max(masked_vec, dim=dim, keepdim=True)[0]
    exps = torch.exp(masked_vec-max_vec)
    masked_exps = exps * mask.float()
    masked_sums = masked_exps.sum(dim, keepdim=True)
    zeros=(masked_sums == 0)
    masked_sums += zeros.float()
    return masked_exps/masked_sums
</code></pre>

Changan: <aside class="quote no-group" data-username="krylea" data-post="13" data-topic="14212" 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/krylea/48/6812_2.png" class="avatar"> krylea:</div>
<blockquote>
<p>I had to implement something similar. My approach was the following (where mask is a tensor of 1s and 0s indicating the entries to be removed):</p>
<pre><code class="lang-auto">def masked_softmax(vec, mask, dim=1):
    masked_vec = vec * mask.float()
    max_vec = torch.max(masked_vec, dim=dim, keepdim=True)[0]
    exps = torch.exp(masked_vec-max_vec)
    masked_exps = exps * mask.float()
    masked_sums = masked_exps.sum(dim, keepdim=True)
    zeros=(masked_sums == 0)
    masked_sums += zeros.float()
    return masked_exps/masked_sums
</code></pre>
</blockquote>
</aside>
<p>Thanks for your solution. But why do you need to compute zeros and add them to the masked_sums? Without zeros, pytorch can still compute the probability by broadcasting the masked_sums.</p>

nurlano: <p>I wrote even simpler version:</p>
<pre><code class="lang-auto">def masked_softmax(vec, mask, dim=1, epsilon=1e-5):
    exps = torch.exp(vec)
    masked_exps = exps * mask.float()
    masked_sums = masked_exps.sum(dim, keepdim=True) + epsilon
    return (masked_exps/masked_sums)
</code></pre>

WendyShang: <p>ehn… directly feed vectors with <code>float('-inf')</code> to pytorch’s default <code>nn.Softmax</code> in fact returns <code>nan</code> at the entries with negative infty.</p>

Junwu_Weng: <p>There is no need to add the epsilon as exp(x) is always larger than 0.</p>

friskit: <p>Try this</p>
<pre><code class="lang-auto">tsr = torch.Tensor([[1,0,3], [0, 1, 2], [3, 2, 1]]).float()
mask = ((tsr  &gt; 0).float() - 1) * 9999  # for -inf
result = (tsr + mask).softmax(dim=-1)
</code></pre>

deniz: <p>Here is a solution by filling masked placed with float(’-inf’):</p>
<pre><code class="lang-auto">import torch.nn.functional as F
F.softmax(vec.masked_fill((1 - mask).bool(), float('-inf')), dim=1)
</code></pre>

gb_pytorch: <p>What is the logic behind skipping zeroes? Does it improve predictions?</p>