x-transformers 1.37.8__tar.gz → 1.37.9__tar.gz

{x_transformers-1.37.8/x_transformers.egg-info → x_transformers-1.37.9}/PKG-INFO

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

{x_transformers-1.37.8 → x_transformers-1.37.9}/README.md

@@ -2252,4 +2252,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Ramapuram2024TheoryAA,
+    title   = {Theory, Analysis, and Best Practices for Sigmoid Self-Attention},
+    author  = {Jason Ramapuram and Federico Danieli and Eeshan Gunesh Dhekane and Floris Weers and Dan Busbridge and Pierre Ablin and Tatiana Likhomanenko and Jagrit Digani and Zijin Gu and Amitis Shidani and Russ Webb},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:272463580}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-1.37.8 → x_transformers-1.37.9}/setup.py

@@ -3,7 +3,7 @@ from setuptools import setup, find_packages
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.37.8',
+  version = '1.37.9',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

{x_transformers-1.37.8 → x_transformers-1.37.9}/x_transformers/attend.py

@@ -36,6 +36,9 @@ def exists(val):
 def default(val, d):
     return val if exists(val) else d
 
+def at_most_one_of(*bools):
+    return sum([*map(int, bools)]) <= 1
+
 def compact(arr):
     return [*filter(exists, arr)]
 
@@ -100,6 +103,7 @@ class Attend(Module):
         scale = None,
         qk_norm = False,
         l2_distance = False,
+        sigmoid = False,
         flash = False,
         softclamp_logits = False,
         logit_softclamp_value = 50.,
@@ -116,10 +120,25 @@ class Attend(Module):
         super().__init__()
         self.scale = scale
 
+        # causal related
+
         self.causal = causal
         self.create_causal_mask = onnx_create_causal_mask if onnxable else create_causal_mask
 
-        self.attn_fn = partial(F.softmax, dtype = torch.float32) if not qk_norm else F.softmax
+        # attention type
+
+        assert not (flash and sigmoid), 'sigmoid attention not available for flash'
+        assert at_most_one_of(sigmoid, l2_distance)
+
+        self.sigmoid = sigmoid
+
+        if not sigmoid:
+            softmax_fn = partial(F.softmax, dim = -1)
+            self.attn_fn = partial(softmax_fn, dtype = torch.float32) if not qk_norm else softmax_fn
+        else:
+            self.attn_fn = F.sigmoid
+
+        # dropouts
 
         self.dropout = dropout
         self.attn_dropout = nn.Dropout(dropout)
@@ -410,7 +429,7 @@ class Attend(Module):
         if self.sigsoftmax:
             sim = sim + sim.sigmoid().log()
 
-        attn = self.attn_fn(sim, dim = -1)
+        attn = self.attn_fn(sim)
         attn = attn.type(dtype)
 
         post_softmax_attn = attn

{x_transformers-1.37.8 → x_transformers-1.37.9}/x_transformers/x_transformers.py

@@ -924,6 +924,7 @@ class Attention(Module):
         qk_norm_scale = 10,
         qk_norm_dim_scale = False,
         l2_distance = False,
+        sigmoid = False,
         one_kv_head = False,
         kv_heads = None,
         shared_kv = False,
@@ -1039,6 +1040,7 @@ class Attention(Module):
             qk_norm = qk_norm,
             scale = qk_norm_scale if qk_norm else self.scale,
             l2_distance = l2_distance,
+            sigmoid = sigmoid,
             add_zero_kv = add_zero_kv,
             flash = flash,
             softclamp_logits = softclamp_logits,

{x_transformers-1.37.8 → x_transformers-1.37.9/x_transformers.egg-info}/PKG-INFO

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