x-transformers 1.18.1__py3-none-any.whl → 1.18.2__py3-none-any.whl

x_transformers/attend.py

@@ -1,4 +1,5 @@
 from functools import partial
+from typing import Optional
 
 import torch
 from torch import nn, einsum, Tensor
@@ -17,9 +18,9 @@ EfficientAttentionConfig = namedtuple('EfficientAttentionConfig', ['enable_flash
 
 @dataclass
 class Intermediates:
-    qk_similarities: Tensor = None
-    pre_softmax_attn: Tensor = None
-    post_softmax_attn: Tensor = None
+    qk_similarities: Optional[Tensor] = None
+    pre_softmax_attn: Optional[Tensor] = None
+    post_softmax_attn: Optional[Tensor] = None
 
     def to_tuple(self):
         return (self.qk_similarities, self.pre_softmax_attn, self.post_softmax_attn)
@@ -256,7 +257,6 @@ class Attend(nn.Module):
             dots = dots + attn_bias
 
         i, j, dtype = *dots.shape[-2:], dots.dtype
-        pre_softmax_attn = dots.clone()
 
         mask_value = -torch.finfo(dots.dtype).max
 
@@ -272,6 +272,8 @@ class Attend(nn.Module):
             causal_mask = self.create_causal_mask(i, j, device = device)
             dots = dots.masked_fill(causal_mask, mask_value)
 
+        pre_softmax_attn = dots.clone()
+
         attn = self.attn_fn(dots, dim = -1)
         attn = attn.type(dtype)
 

x_transformers/x_transformers.py

@@ -9,7 +9,7 @@ from functools import partial, wraps
 from inspect import isfunction
 from collections import namedtuple
 from dataclasses import dataclass
-from typing import List, Callable
+from typing import List, Callable, Optional
 
 from einops import rearrange, repeat, reduce
 from einops.layers.torch import Rearrange
@@ -23,9 +23,10 @@ DEFAULT_DIM_HEAD = 64
 
 @dataclass
 class LayerIntermediates:
-    hiddens: List[Tensor] = None
-    attn_intermediates: List[Intermediates] = None
-    layer_hiddens: List[Tensor] = None
+    hiddens: Optional[List[Tensor]] = None
+    attn_intermediates: Optional[List[Intermediates]] = None
+    layer_hiddens: Optional[List[Tensor]] = None
+    attn_z_loss: Optional[Tensor] = None
 
 # helpers
 
@@ -90,6 +91,31 @@ def or_reduce(masks):
         head = head | rest
     return head
 
+# auxiliary loss helpers
+
+def calc_z_loss(
+    pre_softmax_attns: List[Tensor],
+    mask = None,
+    weight = 1.
+):
+    # the same loss applied to the mixture of experts router logits in https://arxiv.org/abs/2202.08906
+    # in the paper, in a tiny footnote, they mention using it on attention logits with stabilizing effects
+    # also used in PaLM as one of the measures
+
+    lse = 0.
+
+    for attn in pre_softmax_attns:
+        lse = lse + attn.logsumexp(dim = -1)
+
+    loss = torch.square(lse)
+    loss = reduce(loss, 'b h n -> b n', 'sum')
+
+    if not exists(mask):
+        return loss.mean() * weight
+
+    loss = loss[mask].sum() / mask.sum().clamp(min = 1e-5)
+    return loss * weight
+
 # init helpers
 
 def init_zero_(layer):
@@ -1288,7 +1314,8 @@ class TransformerWrapper(nn.Module):
         use_abs_pos_emb = True,
         scaled_sinu_pos_emb = False,
         l2norm_embed = False,
-        emb_frac_gradient = 1. # GLM-130B and Cogview successfully used this, set at 0.1
+        emb_frac_gradient = 1., # GLM-130B and Cogview successfully used this, set at 0.1
+        attn_z_loss_weight = 1e-4
     ):
         super().__init__()
         assert isinstance(attn_layers, AttentionLayers), 'attention layers must be one of Encoder or Decoder'
@@ -1353,10 +1380,12 @@ class TransformerWrapper(nn.Module):
         pos = None,
         prepend_embeds = None,
         sum_embeds = None,
+        return_attn_z_loss = False,
+        attn_z_loss_weight = 1e-4,
         **kwargs
     ):
         b, n, device, num_mem, emb_frac_gradient = *x.shape, x.device, self.num_memory_tokens, self.emb_frac_gradient
-        return_hiddens = return_mems | return_attn | return_intermediates
+        return_hiddens = return_mems | return_attn | return_intermediates | return_attn_z_loss
 
         # absolute positional embedding
 
@@ -1419,6 +1448,11 @@ class TransformerWrapper(nn.Module):
         else:
             out = self.to_logits(x)
 
+        if return_attn_z_loss:
+            pre_softmax_attns = list(map(lambda t: t.pre_softmax_attn, intermediates.attn_intermediates))
+            intermediates.attn_z_loss = calc_z_loss(pre_softmax_attns, weight = attn_z_loss_weight)
+            return_intermediates = True
+
         if return_intermediates:
             return out, intermediates
 

{x_transformers-1.18.1.dist-info → x_transformers-1.18.2.dist-info}/METADATA

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

{x_transformers-1.18.1.dist-info → x_transformers-1.18.2.dist-info}/RECORD

@@ -1,12 +1,12 @@
 x_transformers/__init__.py,sha256=FDb654rUx8FpXRd76B8q0diH8I7q-ZjTWEtEJ4UM21Y,701
-x_transformers/attend.py,sha256=eCgWuKrJJfOPcrQLBJ8_N0FAjJBOWcdGpuGMdY1UfTM,11854
+x_transformers/attend.py,sha256=DRrcDq3ZwmGfuZ8hjLvKB-IJwPjnnVTGX76ONrlFSA4,11913
 x_transformers/autoregressive_wrapper.py,sha256=LuxuSSmZNXivv1XMijOTpNfPd6_uMemptsufL8l1Y7I,4514
 x_transformers/continuous_autoregressive_wrapper.py,sha256=pTiDqu6JRUlnQJQp_xHATYHy0lgSd6ERLqyiFO3pC-4,1575
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
-x_transformers/x_transformers.py,sha256=8b1fp4ezZ07wLSpR3q7izr-0n_7sKm85WuKxKHdU3bw,55263
+x_transformers/x_transformers.py,sha256=BQtmUlLNGvaB6Xn5zK03dpO2OQT1jwpEODUTHZJz_ps,56443
 x_transformers/xl_autoregressive_wrapper.py,sha256=-CAYjTtqrks8ZTxjYm2stOelZpU4MbZIvLjUxWO0P9Y,3988
-x_transformers-1.18.1.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.18.1.dist-info/METADATA,sha256=g1wD6O2VfVdjqP-OKptUt_7O_J4Pt6vcBs7rbrW-TSg,661
-x_transformers-1.18.1.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
-x_transformers-1.18.1.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.18.1.dist-info/RECORD,,
+x_transformers-1.18.2.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.18.2.dist-info/METADATA,sha256=3F1KCGDFqziINFhFhKLciVog2fnuOhigwAnyG49fb3A,661
+x_transformers-1.18.2.dist-info/WHEEL,sha256=yQN5g4mg4AybRjkgi-9yy4iQEFibGQmlz78Pik5Or-A,92
+x_transformers-1.18.2.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.18.2.dist-info/RECORD,,