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

x_transformers/x_transformers.py

@@ -499,7 +499,9 @@ class DataDependentAlibi(Module):
     def __init__(
         self,
         dim,
-        heads
+        heads,
+        bias_init = 5.,
+        post_log_scale = 1.
     ):
         super().__init__()
 
@@ -511,14 +513,55 @@ class DataDependentAlibi(Module):
             nn.LogSigmoid()
         )
 
-        nn.init.constant_(linear.bias, 5.)
+        nn.init.constant_(linear.bias, bias_init)
+        self.post_log_scale = post_log_scale
 
     def forward(self, x):
-        seq = x.shape[-2]
-
-        forget_gates = self.to_forget_gates(x)
+        forget_gates = self.to_forget_gates(x) * self.post_log_scale
         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 queries and keys with a small head dimension """
+
+    def __init__(
+        self,
+        dim,
+        heads,
+        dim_head = 8,
+        post_log_scale = 1.
+    ):
+        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 (qk h d) -> qk b h n d', qk = 2, d = dim_head)
+        )
+
+        self.post_log_scale = post_log_scale
+
+    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) * self.post_log_scale
+
+        # 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 +1014,9 @@ 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,
+        data_dependent_alibi_kwargs: dict = dict(),
         use_cope = False,
         cope_max_pos = 16,
         cope_soft_onehot_pos = False,
@@ -1082,10 +1128,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, **data_dependent_alibi_kwargs)
 
         # attend class - includes core attention algorithm + talking heads
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.41.1
+Version: 1.41.3
 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.3.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=aaxMw4iJkLAh8D4W0g8EwWDfIJVPYdhpmj5T9ZFa0js,91642
 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.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.41.3.dist-info/METADATA,sha256=GyMdKP9ErEnGHbBqY3MwpSGSE41lK_KQbeE5n_O2Py4,689
+x_transformers-1.41.3.dist-info/WHEEL,sha256=P9jw-gEje8ByB7_hXoICnHtVCrEwMQh-630tKvQWehc,91
+x_transformers-1.41.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.41.3.dist-info/RECORD,,