allow for more than one learned head attn sink

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.6.4"
+version = "2.6.5"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_x_transformers.py

@@ -1245,7 +1245,7 @@ def test_learned_head_attn_sink():
             dim = 512,
             depth = 12,
             heads = 8,
-            attn_head_learned_sink = True
+            attn_head_learned_sinks = 4
         )
     )
 

x_transformers/attend.py

@@ -176,7 +176,7 @@ def __init__(
         softclamp_logits = False,
         logit_softclamp_value = 50.,
         add_zero_kv = False,
-        head_learned_sink = False,
+        head_learned_sinks = 0,
         selective = False,
         hard = False,
         cope = None,
@@ -257,10 +257,10 @@ def __init__(
 
         # learned sink concatted pre-softmax, working solution from gpt-oss
 
-        assert not (head_learned_sink and flash), f'not supported for flash attention yet'
+        self.has_head_learned_sinks = head_learned_sinks > 0
+        assert not (self.has_head_learned_sinks and flash), f'not supported for flash attention yet'
 
-        self.head_learned_sink = head_learned_sink
-        self.head_attn_sink = Parameter(torch.zeros(heads)) if head_learned_sink else None
+        self.head_attn_sinks = Parameter(torch.zeros(heads, head_learned_sinks)) if self.has_head_learned_sinks else None
 
         # soft clamp attention logit value
 
@@ -517,9 +517,10 @@ def forward(
         if self.selective:
             sim = selective_attn(sim)
 
-        if self.head_learned_sink:
+        if self.has_head_learned_sinks:
             # add learned attention sink
-            attn_sink = repeat(self.head_attn_sink, 'h -> b h i 1', b = sim.shape[0], i = sim.shape[2])
+            num_sinks = self.head_attn_sinks.shape[-1]
+            attn_sink = repeat(self.head_attn_sinks, 'h sinks -> b h i sinks', b = sim.shape[0], i = sim.shape[2])
             sim = cat((attn_sink, sim), dim = -1)
 
         pre_softmax_attn = sim
@@ -530,9 +531,9 @@ def forward(
 
         post_softmax_attn = attn
 
-        if self.head_learned_sink:
+        if self.has_head_learned_sinks:
             # remove attention sink
-            attn = attn[..., 1:]
+            attn = attn[..., num_sinks:]
 
         attn = self.attn_dropout(attn)
 

x_transformers/x_transformers.py

@@ -1319,7 +1319,7 @@ def __init__(
         value_dim_head = None,
         dim_out = None,
         add_zero_kv = False,         # same as add_zero_attn in pytorch
-        head_learned_sink = False,
+        head_learned_sinks = 0,
         rotate_num_heads = None,
         data_dependent_alibi = False,
         data_dependent_alibi_per_row = False,
@@ -1516,7 +1516,7 @@ def __init__(
             selective = selective,
             custom_attn_fn = custom_attn_fn,
             add_zero_kv = add_zero_kv,
-            head_learned_sink = head_learned_sink,
+            head_learned_sinks = head_learned_sinks,
             flash = flash,
             softclamp_logits = softclamp_logits,
             logit_softclamp_value = logit_softclamp_value,