Long Range Arena benchmark code

[vpuri3]

Hi LION team,

Thanks for the really interesting paper.

I see in your repo you have shared the code for MLM and image classification. Could you also share the code for LRA as well? I'd like to take a look at your implementation as I don't quite understand how you got the performance bump over vanilla transformer or vanilla linear attention. Both saturate at ~40% on listops whereas LION-S works extremely well for some cases and not so much on others (ref. Table 6 in your manuscript).

Thanks,
Vedant

[vpuri3]

@arshiaafzal, could you confirm that this is a faithful reproduction of Model [1] in Table 6, page 31 of your manuscript?

#======================================================================#
# Linear Attention with data-dependent gating
# https://arxiv.org/pdf/2502.16249v1
# Page 31, Table 6, Row 1
#
# LION-S for λi = σ(Wxi) being input dependent, and bi-directional Linear Transformer
# inspired by selectivity of Mamba2 (Dao & Gu, 2024).
# Adds content-aware selectivity (irrelevant tokens are suppressed).
#======================================================================#
class LinearAttentionSelective(nn.Module):
    def __init__(
        self,
        channel_dim: int,
        num_heads: int,
    ):

        assert channel_dim % num_heads == 0

        super().__init__()
        self.channel_dim = channel_dim
        self.num_heads = num_heads
        self.head_dim = channel_dim // num_heads

        self.qkv_proj = nn.Linear(channel_dim, 3 * channel_dim, bias=True)
        self.out_proj = nn.Linear(channel_dim, channel_dim)
        self.gate_mlp = nn.Linear(channel_dim, num_heads)

    def create_mask(self, beta):
        beta = beta.mT.sigmoid().neg().add(1).to(torch.float32)[:, :, 1:]
        beta = torch.cat((torch.ones_like(beta[..., :1]), beta), dim=-1)
        beta = torch.log(beta)

        vec_shape = beta.shape

        A_for = self.create_matrix_from_tensor(beta.unsqueeze(-1).transpose(-1,-2)).squeeze()
        A_back = self.create_matrix_from_tensor(torch.cat((beta,torch.ones((vec_shape[0],vec_shape[1],1),device=beta.device)),dim=-1)[:,:,1:].unsqueeze(-1).transpose(-1,-2)).transpose(-1,-2).squeeze()    
        M = A_for + A_back - torch.eye(A_for.shape[-1]).to(A_for.device)

        return M

    def create_matrix_from_tensor(self, x):

        cumsum = torch.cumsum(x, dim=-1)
        prepend_zeros = torch.zeros(*x.shape[:-1], 1, dtype=x.dtype, device=x.device)
        cumsum = torch.cat((prepend_zeros, cumsum), dim=-1)
        A = cumsum[..., 1:].unsqueeze(-2) - cumsum[..., :-1].unsqueeze(-1)
        A = torch.tril(A.transpose(-1, -2))
        zero_row = torch.zeros(*A.shape[:-2], 1, A.shape[-1], dtype=A.dtype, device=A.device)
        A = torch.cat((zero_row, A[..., :-1, :]), dim=-2)

        return torch.tril(torch.exp(A))

    def forward(self, x):
        B, N, C = x.shape

        q, k, v = rearrange(self.qkv_proj(x), 'b n (h d) -> b h n d', h=self.num_heads).chunk(3, dim=-1)
        beta = self.gate_mlp(x)

        q = F.sigmoid(q)
        k = F.sigmoid(k)

        attn = (q @ k.mT)
        attn = attn * self.create_mask(beta)

        # Do not normalize attention!
        # # normalize attention
        # attn = attn / (attn.sum(dim=-1,keepdim=True) + 1e-6)

        # get output
        y = (attn @ v)

        y = rearrange(y, 'b h n d -> b n (h d)')
        y = self.out_proj(y)

        return y

class LinearAttentionSelectiveBlock(nn.Module):
    def __init__(
        self,
        channel_dim: int,
        num_heads: int,
        mlp_ratio: float = 4.0,
        act: str = None,
    ):
        super().__init__()
        self.ln1 = nn.LayerNorm(channel_dim)
        self.ln2 = nn.LayerNorm(channel_dim)
        self.att = LinearAttentionSelective(channel_dim, num_heads)
        self.mlp = MLPBlock(channel_dim, int(channel_dim * mlp_ratio), channel_dim, act=act, drop=proj_drop)

    def forward(self, x):
        x = self.ln1(x + self.att(x))
        x = self.ln2(x + self.mlp(x))
        return x

#======================================================================#

[arshiaafzal]

Hello, and thanks for raising the issue!

Apologies for the delayed reply, we were preparing the manuscript, and the new version is now available here: https://arxiv.org/abs/2502.16249
.
Regarding LRA: as noted in our paper, in the experiments we used diagonal decays rather than a scalar. The initialization follows HIPPO (or LRU), which allows the model to effectively learn the task. In addition, the LRA tasks were highly sensitive to the choice of activation function. As shown in Table 6, the activation applied to the query/key was either sigmoid or tanh. The reason for this choice is to reduce and bound the selectivity such that the model can "over-fit" to each of these tasks.

So for above code actually LION needs to be applied with digonal decays not selective. Which they will then be injected into query/key pairs as shown in Appendix section B.5.