Usage of alpha about entmax HOT 3 CLOSED

Hi,

We did define a trainable parameter for each head per layer. Here's a small code snippet that we used:

class AlphaChooser(torch.nn.Module):

    def __init__(self, head_count):
        """head_count (int): number of attention heads"""
        super(AlphaChooser, self).__init__()
        self.pre_alpha = nn.Parameter(torch.randn(head_count))

    def forward(self):
        alpha = 1 + torch.sigmoid(self.pre_alpha)
        return torch.clamp(alpha, min=1.01, max=2)

However, it's possible to have a single alpha per layer, or per transformer block!

from entmax.

Thanks so much! @goncalomcorreia
Nice work though!

from entmax.

Hello!

How can I compute the entmax_bisect when the size of alpha is larger than 1 ?
When I use this:
p_attn = entmax_bisect(x, alpha, n_iter=25)
where alpha=tensor([1.3691, 1.5766, 1.7588, 1.9206],grad_fn=), the shape of x is [batch_size, 4, d,d].
There arises the error that
The expanded size of the tensor (1) must match the existing size (4) at non-singleton dimension 3.

from entmax.