x-transformers 1.41.1__py3-none-any.whl → 1.41.2__py3-none-any.whl

x_transformers/x_transformers.py

@@ -514,11 +514,48 @@ class DataDependentAlibi(Module):
         nn.init.constant_(linear.bias, 5.)
 
     def forward(self, x):
-        seq = x.shape[-2]
-
         forget_gates = self.to_forget_gates(x)
         forget_gates = forget_gates.cumsum(dim = -1)
         forget_gates = einx.subtract('b h i, b h j -> b h i j', forget_gates, forget_gates)
+        return forget_gates
+
+class PerRowDataDependentAlibi(Module):
+    """ same as data dependent alibi from forgetting transformer, but the forgetting gates are also derived by a queris and keys with a small head dimension """
+
+    def __init__(
+        self,
+        dim,
+        heads,
+        dim_head = 8
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+
+        linear = nn.Linear(dim, heads * dim_head * 2, bias = False)
+
+        self.to_forget_gates = nn.Sequential(
+            linear,
+            Rearrange('b n (kv h d) -> kv b h n d', kv = 2, d = dim_head)
+        )
+
+    def forward(self, x):
+        q, k = self.to_forget_gates(x)
+        forget_gates = einsum('... i d, ... j d -> ... i j', q, k) * self.scale
+
+        forget_gates = F.logsigmoid(forget_gates)
+
+        # mask out upper triangle + diagonal
+
+        n = x.shape[-2]
+        causal_mask = torch.ones((n, n), dtype = torch.bool, device = x.device).triu()
+
+        forget_gates = forget_gates.masked_fill(causal_mask, 0.)
+
+        # reverse cumsum
+
+        forget_gates = forget_gates.flip(dims = (-1,))
+        forget_gates = forget_gates.cumsum(dim = -1)
+        forget_gates = forget_gates.flip(dims = (-1,))
 
         return forget_gates
 
@@ -971,6 +1008,8 @@ class Attention(Module):
         add_zero_kv = False,         # same as add_zero_attn in pytorch
         rotary_embed_values = False,
         data_dependent_alibi = False,
+        data_dependent_alibi_per_row = False,
+        data_dependent_alibi_per_row_dim_head = 8,
         use_cope = False,
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
@@ -1082,10 +1121,13 @@ class Attention(Module):
         if data_dependent_alibi:
             assert causal, 'data dependent alibi only works for autoregressive for now until further research'
 
-            self.data_dependent_alibi = DataDependentAlibi(
-                dim,
-                heads = heads
-            )
+            dda_klass = DataDependentAlibi if not data_dependent_alibi_per_row else PerRowDataDependentAlibi
+            dda_kwargs = dict(dim = dim, heads = heads)
+
+            if data_dependent_alibi_per_row:
+                dda_kwargs.update(dim_head = data_dependent_alibi_per_row_dim_head)
+
+            self.data_dependent_alibi = dda_klass(**dda_kwargs)
 
         # attend class - includes core attention algorithm + talking heads
 

{x_transformers-1.41.1.dist-info → x_transformers-1.41.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.41.1
+Version: 1.41.2
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.41.1.dist-info → x_transformers-1.41.2.dist-info}/RECORD

@@ -5,11 +5,11 @@ x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/x_transformers.py,sha256=n8W19Pnhbz-JxbC7QATApWrhI_yC4oqTHGQ1NLuindY,89814
+x_transformers/x_transformers.py,sha256=5yoN96mTpw3qygZg2plVed5lTOs9p0REtt50lOlllHU,91334
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-1.41.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.41.1.dist-info/METADATA,sha256=UIPYbEVBLrWDGuezlnyh2tFKPlM_Mdj-pYTGxse_NMI,689
-x_transformers-1.41.1.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
-x_transformers-1.41.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.41.1.dist-info/RECORD,,
+x_transformers-1.41.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.41.2.dist-info/METADATA,sha256=84WHjWdVHIxkHdHxMEKK5OEssLU80cdj4aGTJQ7SpO4,689
+x_transformers-1.41.2.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.41.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.41.2.dist-info/RECORD,,