Replicating the behaviour from the paper about entmax HOT 1 CLOSED

Hi! Thank you for your interest!

1. Yes, setting it as a Parameter is sufficient since entmax_bisect already takes care of the gradient computation for you.

2. I'm not really sure about your code snippet so maybe this can help:

class EntmaxAlpha(torch.nn.Module):

    def __init__(self, head_count, n_iter=25, dim=0):
        super(EntmaxAlpha, self).__init__()
        self.dim = dim
        self.n_iter = n_iter
        self.alpha_chooser = AlphaChooser(head_count)

    def forward(self, X):
        batch_size, head_count, query_len, key_len = X.size()

        X = X.view(-1, key_len)
        self.alpha = self.alpha_chooser()
        expanded_alpha = self.alpha.unsqueeze(0).unsqueeze(-1)
        expanded_alpha = expanded_alpha.expand((batch_size, -1, query_len))
        expanded_alpha = expanded_alpha.contiguous().view(-1)
        p_star = entmax_bisect(X, expanded_alpha, self.n_iter)

        return p_star.view(batch_size, head_count, query_len, -1)

We used this EntmaxAlpha just as you would use nn.Softmax in a regular multi-head attention implementation (particularly, we used OpenNMT).

3. The code snippet above should answer your question. You can compute the attention probabilities for the several heads in parallel (even if they have different alphas).

4. Note that p.mean() is a constant since p always sums to 1. Since it's a constant, calling backward() on this will get you zero gradients.

5. This is learned automatically via backpropagation! :)

from entmax.