K-winner-take-all advanced indexing

username: dizcza
cooked: <p>I’m trying to implement efficient k-winner-take-all activation function in PyTorch.</p>
<pre><code class="lang-auto">class KWinnersTakeAll(torch.autograd.Function):

    @staticmethod
    def forward(ctx, tensor, sparsity: float):
        batch_size, embedding_size = tensor.shape
        _, argsort = tensor.sort(dim=1, descending=True)
        k_active = math.ceil(sparsity * embedding_size)
        active_indices = argsort[:, :k_active]
        mask_active = torch.ByteTensor(tensor.shape).zero_()
        for sample_id in range(batch_size):
            mask_active[sample_id, active_indices[sample_id]] = 1
        tensor[~mask_active] = 0
        tensor[mask_active] = 1
        ctx.save_for_backward(mask_active)
        return tensor
</code></pre>
<p>What I want is to get rid of for-loop by vectorizing the computations it does. Any idea how to do that?<br>
Let <code>mask_active.shape == (256, 128)</code> and <code>active_indices.shape == (256, 7)</code> then<br>
<code>mask_active[active_indices].shape == (256, 7, 128)</code> but I’d like to have <code>mask_active[active_indices].shape == (256, 7)</code> so then I could assign <code>mask_active[active_indices] = 1</code></p>

username: ptrblck
cooked: <p>I think you could use advanced indexing for your use case.<br>
Here is a small example with smaller sizes for better understanding:</p>
<pre><code class="lang-python">batch_size = 5
features = 4

idx = torch.empty(batch_size, 2, dtype=torch.long).random_(features)
mask_active = torch.zeros(x.shape, dtype=torch.uint8)
mask_active_fast = torch.zeros(x.shape, dtype=torch.uint8)
# Your approach
for i in range(batch_size):
    mask_active[i, idx[i]] = 1

# Advanced indexing
mask_active_fast[torch.arange(batch_size).unsqueeze(1), idx] = 1

# Check
(mask_active==mask_active_fast).all()
</code></pre>

username: adamk
cooked: <p>Sorry for bringing this up again, but has anybody an idea how to implement the backward path?</p>
<p>or is it simply passing the gradient through, where the mask is one? like this:</p>
<pre><code class="lang-auto">
    @staticmethod
    def backward(ctx, grad_output):
        mask_active, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[~mask_active] = 0
        return grad_input
</code></pre>

username: dizcza
cooked: <p><a class="mention" href="/u/adamk">@adamk</a> zeroing-out masked values like you do won’t work for complex data (only MNIST works out such a cruel cut-off in my experiments). However, the paper linked below did use this approach and showed that it works.</p>
<p>For further inquiries on how to broadcast the gradients of the k-winners-take-all function backward, refer to my kwta.py script below. The solution is, as always, to approximate the non-differentiable function with a soft approximation (see <code>KWinnersTakeAllSoft</code> class).</p>
<p>For convenience, I made a list of how people implement kWTA in PyTorch and deal with its gradients.</p>
<ul>
<li>my <a href="https://github.com/dizcza/EmbedderSDR/blob/master/models/kwta.py" rel="noopener nofollow ugc">kwta.py</a> script: straightforward hard kWTA (<code>KWinnersTakeAll</code>) and its soft approximation (<code>KWinnersTakeAllSoft</code>).</li>
<li>two implementations of “Haw can we be so dense?” (Subutai Ahmad, Luiz Scheinkman, 2019) paper. Both employ the straightforward hard kWTA.
<ul>
<li><a href="https://github.com/numenta/htmpapers/blob/cc50752afc010f1c67fac672651baba48ab47e81/arxiv/how_can_we_be_so_dense/src/pytorch/functions/k_winners.py#L116-L129" rel="noopener nofollow ugc">original</a></li>
<li>
<a href="https://github.com/marty1885/sparsenet-pytorch/blob/master/sparsenet.py" class="inline-onebox" rel="noopener nofollow ugc">sparsenet-pytorch/sparsenet.py at master · marty1885/sparsenet-pytorch · GitHub</a>.</li>
</ul>
</li>
</ul>