x-transformers 1.30.23__tar.gz → 1.31.1__tar.gz

{x_transformers-1.30.23/x_transformers.egg-info → x_transformers-1.31.1}/PKG-INFO

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

{x_transformers-1.30.23 → x_transformers-1.31.1}/README.md

@@ -2169,4 +2169,22 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{Guttenberg2023,
+    author  = {Ohad Rubin},
+    url     = {https://medium.com/@ohadrubin/exploring-weight-decay-in-layer-normalization-challenges-and-a-reparameterization-solution-ad4d12c24950}
+}
+```
+
+```bibtex
+@article{Mesnard2024GemmaOM,
+    title   = {Gemma: Open Models Based on Gemini Research and Technology},
+    author  = {Gemma Team Thomas Mesnard and Cassidy Hardin and Robert Dadashi and Surya Bhupatiraju and Shreya Pathak and L. Sifre and Morgane Riviere and Mihir Kale and J Christopher Love and Pouya Dehghani Tafti and L'eonard Hussenot and Aakanksha Chowdhery and Adam Roberts and Aditya Barua and Alex Botev and Alex Castro-Ros and Ambrose Slone and Am'elie H'eliou and Andrea Tacchetti and Anna Bulanova and Antonia Paterson and Beth Tsai and Bobak Shahriari and Charline Le Lan and Christopher A. Choquette-Choo and Cl'ement Crepy and Daniel Cer and Daphne Ippolito and David Reid and Elena Buchatskaya and Eric Ni and Eric Noland and Geng Yan and George Tucker and George-Christian Muraru and Grigory Rozhdestvenskiy and Henryk Michalewski and Ian Tenney and Ivan Grishchenko and Jacob Austin and James Keeling and Jane Labanowski and Jean-Baptiste Lespiau and Jeff Stanway and Jenny Brennan and Jeremy Chen and Johan Ferret and Justin Chiu and Justin Mao-Jones and Katherine Lee and Kathy Yu and Katie Millican and Lars Lowe Sjoesund and Lisa Lee and Lucas Dixon and Machel Reid and Maciej Mikula and Mateo Wirth and Michael Sharman and Nikolai Chinaev and Nithum Thain and Olivier Bachem and Oscar Chang and Oscar Wahltinez and Paige Bailey and Paul Michel and Petko Yotov and Pier Giuseppe Sessa and Rahma Chaabouni and Ramona Comanescu and Reena Jana and Rohan Anil and Ross McIlroy and Ruibo Liu and Ryan Mullins and Samuel L Smith and Sebastian Borgeaud and Sertan Girgin and Sholto Douglas and Shree Pandya and Siamak Shakeri and Soham De and Ted Klimenko and Tom Hennigan and Vladimir Feinberg and Wojciech Stokowiec and Yu-hui Chen and Zafarali Ahmed and Zhitao Gong and Tris Brian Warkentin and Ludovic Peran and Minh Giang and Cl'ement Farabet and Oriol Vinyals and Jeffrey Dean and Koray Kavukcuoglu and Demis Hassabis and Zoubin Ghahramani and Douglas Eck and Joelle Barral and Fernando Pereira and Eli Collins and Armand Joulin and Noah Fiedel and Evan Senter and Alek Andreev and Kathleen Kenealy},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2403.08295},
+    url     = {https://api.semanticscholar.org/CorpusID:268379206}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-1.30.23 → x_transformers-1.31.1}/setup.py

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

{x_transformers-1.30.23 → x_transformers-1.31.1}/x_transformers/x_transformers.py

@@ -93,6 +93,9 @@ def l2norm(t, groups = 1):
     t = F.normalize(t, p = 2, dim = -1)
     return rearrange(t, '... g d -> ... (g d)')
 
+def softclamp(t, value):
+    return (t / value).tanh() * value
+
 def pad_at_dim(t, pad: Tuple[int, int], dim = -1, value = 0.):
     if pad == (0, 0):
         return t
@@ -560,22 +563,34 @@ class Scale(Module):
         return (scale_fn(out[0]), *out[1:])
 
 class LayerNorm(Module):
-    def __init__(self, dim):
+    def __init__(
+        self,
+        dim,
+        unit_offset = 0.
+    ):
         """
         bias-less layernorm has been shown to be more stable. most newer models have moved towards rmsnorm, also bias-less
         """
         super().__init__()
+        self.unit_offset = unit_offset
+
+        self.ln = nn.LayerNorm(dim, elementwise_affine = False)
         self.gamma = nn.Parameter(torch.ones(dim))
-        self.register_buffer("beta", torch.zeros(dim))
+        nn.init.constant_(self.gamma, 1. - unit_offset)
 
-    def forward(self, x):
-        return F.layer_norm(x, x.shape[-1:], self.gamma, self.beta)
+        self.register_buffer('beta', torch.zeros(dim), persistent = False)
 
-if version.parse(torch.__version__) >= version.parse('2.1.0'):
-    LayerNorm = partial(nn.LayerNorm, bias = False)
+    def forward(self, x):
+        normed = self.ln(x)
+        gamma = self.gamma + self.unit_offset
+        return normed * gamma
 
 class AdaptiveLayerNorm(Module):
-    def __init__(self, dim, dim_condition = None):
+    def __init__(
+        self,
+        dim,
+        dim_condition = None
+    ):
         super().__init__()
         dim_condition = default(dim_condition, dim)
 
@@ -590,25 +605,43 @@ class AdaptiveLayerNorm(Module):
         return normed * (gamma + 1.)
 
 class ScaleNorm(Module):
-    def __init__(self, dim):
+    def __init__(
+        self,
+        dim,
+        unit_offset = 0.
+    ):
         super().__init__()
+        self.unit_offset = unit_offset
         self.scale = dim ** 0.5
-        self.g = nn.Parameter(torch.ones(1))
+
+        self.g = nn.Parameter(torch.zeros(1))
+        nn.init.constant_(self.g, 1. - unit_offset)
 
     def forward(self, x):
-        return F.normalize(x, dim = -1) * self.scale * self.g
+        return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
 
 class RMSNorm(Module):
-    def __init__(self, dim):
+    def __init__(
+        self,
+        dim,
+        unit_offset = 0.
+    ):
         super().__init__()
+        self.unit_offset = unit_offset
         self.scale = dim ** 0.5
-        self.g = nn.Parameter(torch.ones(dim))
+
+        self.g = nn.Parameter(torch.zeros(dim))
+        nn.init.constant_(self.g, 1. - unit_offset)
 
     def forward(self, x):
-        return F.normalize(x, dim = -1) * self.scale * self.g
+        return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
 
 class AdaptiveRMSNorm(Module):
-    def __init__(self, dim, dim_condition = None):
+    def __init__(
+        self,
+        dim,
+        dim_condition = None
+    ):
         super().__init__()
         self.scale = dim ** 0.5
         dim_condition = default(dim_condition, dim)
@@ -622,7 +655,11 @@ class AdaptiveRMSNorm(Module):
         return normed * self.scale * (gamma + 1.)
 
 class SimpleRMSNorm(Module):
-    def __init__(self, dim):
+    def __init__(
+        self,
+        dim,
+        **kwargs
+    ):
         super().__init__()
         self.scale = dim ** 0.5
 
@@ -696,10 +733,19 @@ class ShiftTokens(Module):
 # post branch operator
 
 class LayerScale(Module):
-    def __init__(self, fn: Module, dim, init_value = 0.):
+    def __init__(
+        self,
+        fn: Module,
+        dim,
+        init_value = 0.,
+        unit_offset = 0.
+    ):
         super().__init__()
+        self.unit_offset = unit_offset
+
         self.fn = fn
-        self.gamma = nn.Parameter(torch.ones(dim) * init_value)
+        self.gamma = nn.Parameter(torch.zeros(dim))
+        nn.init.constant_(self.gamma, init_value - unit_offset)
 
     def forward(self, x, **kwargs):
         out = self.fn(x, **kwargs)
@@ -711,7 +757,13 @@ class LayerScale(Module):
         return out * self.gamma, *rest
 
 class AdaptiveLayerScale(Module):
-    def __init__(self, fn: Module, dim, dim_condition = None, init_bias_value = -2.):
+    def __init__(
+        self,
+        fn: Module,
+        dim,
+        dim_condition = None,
+        init_bias_value = -2.
+    ):
         super().__init__()
         self.fn = fn
 
@@ -1182,6 +1234,7 @@ class AttentionLayers(Module):
         use_adaptive_layernorm = False,
         use_adaptive_rmsnorm = False,
         use_adaptive_layerscale = False, # paired with use_adaptive_layernorm for ada-ln-zero from DiT paper
+        norm_add_unit_offset = False,
         dim_condition = None,
         adaptive_condition_mlp = False,
         adaptive_condition_mlp_expansion = 4,
@@ -1215,6 +1268,7 @@ class AttentionLayers(Module):
         scale_residual_constant = 1.,
         shift_tokens = 0,
         sandwich_norm = False,
+        softclamp_output_value: float | None = None,
         resi_dual = False,
         resi_dual_scale = 1.,
         zero_init_branch_output = False,
@@ -1315,6 +1369,10 @@ class AttentionLayers(Module):
 
         norm_fn = partial(norm_class, dim)
 
+        if not norm_need_condition and norm_add_unit_offset:
+            # researcher Ohad Rubin shares in a blog post by adding an offset to gammas and betas, they can be subjected to weight decay safely
+            norm_fn = partial(norm_fn, unit_offset = 1.)
+
         self.norm_need_condition = norm_need_condition
         self.dim_condition = dim_condition
 
@@ -1345,6 +1403,9 @@ class AttentionLayers(Module):
 
         self.post_branch_fn_needs_condition = post_branch_fn_needs_condition
 
+        if not post_branch_fn_needs_condition and norm_add_unit_offset:
+            post_branch_fn = partial(post_branch_fn, unit_offset = 1.)
+
         # setup mlp for conditioning
 
         self.need_condition = norm_need_condition or post_branch_fn_needs_condition
@@ -1421,6 +1482,11 @@ class AttentionLayers(Module):
 
         shift_tokens = cast_tuple(shift_tokens, len(layer_types))
 
+        # optional soft clamping just before the final norm
+        # used in gemma 2
+
+        self.softclamp_output_value = softclamp_output_value
+
         # whether it has post norm
 
         self.final_norm = norm_fn() if pre_norm or resi_dual else nn.Identity()
@@ -1652,6 +1718,9 @@ class AttentionLayers(Module):
         if return_hiddens:
             layer_hiddens.append(x)
 
+        if exists(self.softclamp_output_value):
+            x = softclamp(x, self.softclamp_output_value)
+
         final_norm = self.final_norm
 
         if self.need_condition:

{x_transformers-1.30.23 → x_transformers-1.31.1/x_transformers.egg-info}/PKG-INFO

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