x-transformers 1.38.3__py3-none-any.whl → 1.39.1__py3-none-any.whl

x_transformers/attend.py

@@ -128,7 +128,9 @@ class Attend(Module):
         dropout = 0.,
         causal = False,
         heads = None,
-        talking_heads = False,
+        pre_talking_heads = False,
+        post_talking_heads = False,
+        pre_scale_post_talking_heads = False,
         sparse_topk = None,
         scale = None,
         qk_norm = False,
@@ -179,16 +181,22 @@ class Attend(Module):
 
         # talking heads
 
-        assert not (flash and talking_heads), 'talking heads not compatible with flash attention'
+        assert not (flash and (pre_talking_heads or post_talking_heads or pre_scale_post_talking_heads)), 'talking heads not compatible with flash attention'
 
-        self.talking_heads = talking_heads
-        if talking_heads:
-            self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
-            self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False)
+        self.pre_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if pre_talking_heads else None
+        self.post_softmax_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if post_talking_heads else None
+        self.pre_scale_post_talking_heads = nn.Conv2d(heads, heads, 1, bias = False) if pre_scale_post_talking_heads else None
 
+        if exists(self.pre_softmax_talking_heads):
             nn.init.dirac_(self.pre_softmax_talking_heads.weight)
+
+        if exists(self.post_softmax_talking_heads):
             nn.init.dirac_(self.post_softmax_talking_heads.weight)
 
+        if exists(self.pre_scale_post_talking_heads):
+            # an improvisation where heads are combined pre-softmax attention, then used to scale post-softmax attention
+            nn.init.dirac_(self.pre_scale_post_talking_heads.weight)
+
         # selective attention
 
         assert not (flash and selective), 'selective attention cannot work on flash attention'
@@ -434,8 +442,11 @@ class Attend(Module):
 
         qk_similarities = sim.clone()
 
-        if self.talking_heads:
-            sim = self.pre_softmax_talking_heads(sim)
+        if exists(self.pre_scale_post_talking_heads):
+            pre_to_post_scale = self.pre_scale_post_talking_heads(sim)
+
+        if exists(self.pre_softmax_talking_heads):
+            sim = sim + self.pre_softmax_talking_heads(sim)
 
         if exists(attn_bias):
             sim = sim + attn_bias
@@ -482,9 +493,12 @@ class Attend(Module):
 
         attn = self.attn_dropout(attn)
 
-        if self.talking_heads:
+        if exists(self.post_softmax_talking_heads):
             attn = self.post_softmax_talking_heads(attn)
 
+        if exists(self.pre_scale_post_talking_heads):
+            attn = attn * pre_to_post_scale
+
         out = einsum(f'b h i j, {kv_einsum_eq} -> b h i d', attn, v)
 
         intermediates = Intermediates(

x_transformers/x_transformers.py

@@ -907,7 +907,9 @@ class Attention(Module):
         heads = 8,
         causal = False,
         flash = False,
-        talking_heads = False,
+        pre_talking_heads = False,
+        post_talking_heads = False,
+        pre_scale_post_talking_heads = False,
         head_scale = False,
         sparse_topk = None,
         num_mem_kv = 0,
@@ -1036,7 +1038,9 @@ class Attention(Module):
         self.attend = Attend(
             heads = heads,
             causal = causal,
-            talking_heads = talking_heads,
+            pre_talking_heads = pre_talking_heads,
+            post_talking_heads = post_talking_heads,
+            pre_scale_post_talking_heads = pre_scale_post_talking_heads,
             dropout = dropout,
             sparse_topk = sparse_topk,
             qk_norm = qk_norm,

{x_transformers-1.38.3.dist-info → x_transformers-1.39.1.dist-info}/METADATA

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

{x_transformers-1.38.3.dist-info → x_transformers-1.39.1.dist-info}/RECORD

@@ -1,15 +1,15 @@
 x_transformers/__init__.py,sha256=-MkQrSc37cTVDX7AOykxunYnqVtFlQ7lb0Cse5dsGWU,793
-x_transformers/attend.py,sha256=Z9cnbY3f9zCl1yUOMWVQT5_ee4keC0lo_NX4cd0rbKk,15393
+x_transformers/attend.py,sha256=JlWwEzigY5sl9ktDwDWWQD9np9uPRoj2eRo9XU6tJc0,16273
 x_transformers/autoregressive_wrapper.py,sha256=DOJJCMMDOqDYKWy_IaG5IyKsXD3AW6amzfUgdAADOLY,10500
 x_transformers/continuous.py,sha256=cIVEdhfei258__ziV7kQBrJMxCel54bExBTDrO9rfCI,6450
 x_transformers/dpo.py,sha256=LjvWgCkqTl-UuehrzQ8nkX5guLr4whYwsmm7SKSwdls,3450
 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=P3o1R2DY2ic71FlJJ4ie4w_z-g3jIIrkBXcbllRoXHA,84240
+x_transformers/x_transformers.py,sha256=8ZQR6OLT4vusIjJXzrdSp12Fydmmpcc2t5cDE6SxPNc,84460
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=QE1ltYZTR_eGgIHPP2BrMWVWVLqMW-OpDZh87BSmQEg,8563
-x_transformers-1.38.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.38.3.dist-info/METADATA,sha256=XbmJ91NWqmOjO0xvyXh9uh82TIwVbU54L3eFmaFVvYs,661
-x_transformers-1.38.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-x_transformers-1.38.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.38.3.dist-info/RECORD,,
+x_transformers-1.39.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.39.1.dist-info/METADATA,sha256=KJfw4hIDzozyRlsnBaqbbfcFSXygbv1y2B_6cl9pu-4,661
+x_transformers-1.39.1.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
+x_transformers-1.39.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.39.1.dist-info/RECORD,,