x-transformers 2.6.6__py3-none-any.whl → 2.7.0__py3-none-any.whl

x_transformers/attend.py

@@ -23,7 +23,7 @@ class Intermediates:
     pre_softmax_attn:   Tensor | None = None
     post_softmax_attn:  Tensor | None = None
     values:             Tensor | None = None
-    cached_kv:          Tuple[Tensor, Tensor] | None = None
+    cached_kv:          tuple[Tensor, Tensor] | None = None
     layer_type:         str | None = None
     hybrid_hidden:      Tensor | None = None
 

x_transformers/nonautoregressive_wrapper.py

@@ -1,16 +1,20 @@
+from __future__ import annotations
+
 import math
 from random import random
 from contextlib import nullcontext
 from collections import namedtuple
 
 import torch
+from torch import nn, pi
+from torch.nn import Module
+from torch.func import grad_and_value, vmap
 import torch.nn.functional as F
-from torch import nn
 
+import einx
 from einops import rearrange, repeat, pack, unpack
 
 from x_transformers.x_transformers import TransformerWrapper
-from typing import Optional
 
 # constants
 
@@ -75,12 +79,12 @@ def linear_schedule(t):
 
 def cosine_schedule(t):
     """ https://arxiv.org/abs/2202.04200 """
-    return torch.cos(t * math.pi / 2)
+    return torch.cos(t * pi / 2)
 
 # self token critic
 # inspired by Nijkamp et al. - https://aclanthology.org/2021.naacl-main.409/
 
-class SelfCritic(nn.Module):
+class SelfCritic(Module):
     def __init__(self, net):
         super().__init__()
         self.net = net
@@ -92,7 +96,7 @@ class SelfCritic(nn.Module):
         embed = self.net(x, return_embeddings = True)
         return self.to_logits(embed)
 
-class NonAutoregressiveWrapper(nn.Module):
+class NonAutoregressiveWrapper(Module):
     """
     https://arxiv.org/abs/1904.09324
     https://arxiv.org/abs/2202.04200
@@ -110,9 +114,10 @@ class NonAutoregressiveWrapper(nn.Module):
         random_token_prob = 0.1,         # which percentage of tokens to be replaced with random token, done in original MLM paper
         schedule = 'linear',
         can_mask_prev_unmasked = False,  # when unmasking, whether it can remask previously unmasked
-        token_critic: Optional[TransformerWrapper] = None,
+        token_critic: TransformerWrapper | None = None,
         self_token_critic = False,
-        critic_loss_weight = 1.
+        critic_loss_weight = 1.,
+        use_simple_mdlm_loss_weight = True # Sahoo et al. https://arxiv.org/abs/2406.07524
     ):
         super().__init__()
         assert not (self_token_critic and exists(token_critic))
@@ -143,6 +148,23 @@ class NonAutoregressiveWrapper(nn.Module):
         else:
             raise ValueError(f'invalid schedule {schedule}')
 
+        # whether to use the loss weighting proposed in simple diffusion lm paper
+
+        self.loss_weight_fn = None
+
+        if use_simple_mdlm_loss_weight:
+            grad_and_value_schedule_fn = vmap(grad_and_value(self.schedule_fn))
+
+            # eq (10)
+
+            def loss_weight_fn(times):
+                grad, value = grad_and_value_schedule_fn(times)
+                return grad / (1. - value)
+
+            self.loss_weight_fn = loss_weight_fn
+
+        # whether to mask previous - in the simple mdlm paper, they chose not to
+
         self.can_mask_prev_unmasked = can_mask_prev_unmasked
 
         # self conditioning
@@ -311,12 +333,27 @@ class NonAutoregressiveWrapper(nn.Module):
 
         loss_fn = F.cross_entropy if not self.net.output_is_log_prob else F.nll_loss
 
-        # cross entropy loss
+        # loss
 
-        loss = loss_fn(
-            logits[mask],
-            orig_seq[mask]
-        )
+        if exists(self.loss_weight_fn):
+            # using simple mdlm loss weighting
+
+            loss = loss_fn(
+                rearrange(logits, 'b n l -> b l n'),
+                orig_seq,
+                reduction = 'none'
+            )
+
+            loss_weights = self.loss_weight_fn(rand_times)     # calculate loss weight
+            loss = einx.multiply('b n, b', loss, loss_weights) # apply loss weights
+
+            loss = loss[mask].mean()
+
+        else:
+            loss = loss_fn(
+                logits[mask],
+                orig_seq[mask],
+            )
 
         if not exists(self.token_critic) or only_train_generator:
             return Losses(loss, loss, None)

x_transformers/x_transformers.py

@@ -10,7 +10,7 @@ import torch
 from torch.amp import autocast
 import torch.nn.functional as F
 from torch import nn, einsum, tensor, Tensor, cat, stack, arange, is_tensor
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.utils._pytree import tree_flatten, tree_unflatten, tree_map
 from torch.nn import Module, ModuleList, ModuleDict
 
 from functools import partial, wraps
@@ -81,6 +81,9 @@ def cast_tuple(val, depth = 1):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def detach_all(obj):
+    return tree_map(lambda t: t.detach() if is_tensor(t) and t.requires_grad else t, obj)
+
 def maybe(fn = None):
     if not exists(fn):
         fn = identity
@@ -157,6 +160,19 @@ def or_reduce(masks):
         head = head | rest
     return head
 
+# cache helpers
+
+def get_cached_kvs(
+    cache: LayerIntermediates
+) -> list[tuple[Tensor, Tensor]]:
+
+    cached_kvs = []
+
+    for attn_intermediate in cache.attn_intermediates:
+        cached_kvs.append(attn_intermediate.cached_kv)
+
+    return cached_kvs
+
 # entropy
 
 def calc_entropy(
@@ -2441,8 +2457,13 @@ class AttentionLayers(Module):
         context_pos = None,
         attn_bias = None,
         deep_embeds_and_ids: tuple[nn.Parameter, Tensor] | None = None,
-        self_attn_additional_kv: list[tuple[Tensor, Tensor]] | None = None,
+        self_attn_additional_kv: (
+            LayerIntermediates |
+            list[tuple[Tensor, Tensor]]
+            | None
+        ) = None,
         additional_kv_mask = None,
+        detach_additional_kv = False,
         route_additional_kv_to_top = True,
         condition = None,
         in_attn_cond = None, # https://arxiv.org/abs/2105.04090
@@ -2590,6 +2611,13 @@ class AttentionLayers(Module):
         # additional self attn key / values - say coming from vlm
 
         if exists(self_attn_additional_kv) and route_additional_kv_to_top:
+
+            if isinstance(self_attn_additional_kv, LayerIntermediates):
+                self_attn_additional_kv = get_cached_kvs(self_attn_additional_kv)
+
+            if detach_additional_kv:
+                self_attn_additional_kv = detach_all(self_attn_additional_kv)
+
             num_self_attns = sum([layer_type == 'a' for layer_type in first(layer_variables)])
 
             self_attn_additional_kv = self_attn_additional_kv[-num_self_attns:]

{x_transformers-2.6.6.dist-info → x_transformers-2.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.6.6
+Version: 2.7.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2509,11 +2509,22 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 
 ```bibtex
 @misc{openai_gpt_oss,
-  author       = {OpenAI},
-  title        = {Introducing gpt-oss},
-  howpublished = {https://openai.com/index/introducing-gpt-oss},
-  month        = {August},
-  year         = {2025}
+    author       = {OpenAI},
+    title        = {Introducing gpt-oss},
+    howpublished = {https://openai.com/index/introducing-gpt-oss},
+    month        = {August},
+    year         = {2025}
+}
+```
+
+```bibtex
+@article{Sahoo2024SimpleAE,
+    title   = {Simple and Effective Masked Diffusion Language Models},
+    author  = {Subham Sekhar Sahoo and Marianne Arriola and Yair Schiff and Aaron Gokaslan and Edgar Marroquin and Justin T Chiu and Alexander Rush and Volodymyr Kuleshov},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2406.07524},
+    url     = {https://api.semanticscholar.org/CorpusID:270380319}
 }
 ```
 

{x_transformers-2.6.6.dist-info → x_transformers-2.7.0.dist-info}/RECORD

@@ -1,5 +1,5 @@
 x_transformers/__init__.py,sha256=aVuhUU0572TJHW88BVc4yA2tla0Zb8l3NH7W4RZ1AEs,1005
-x_transformers/attend.py,sha256=JJv6ypJbZIFmH1LQ49hFg6hD0Wf9Z7Im1AP2ekm9hVI,18091
+x_transformers/attend.py,sha256=jzOwrtCIdAt1dRQBO68htDsgtjdTx6TAQQVB2xflS1w,18091
 x_transformers/autoregressive_wrapper.py,sha256=BsGO9xfVYkvynqbU1__tu_S_cxl7gss0YwnkhIa2baY,18401
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
 x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
@@ -7,12 +7,12 @@ x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
-x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
+x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=vjRMEMA12Js94YwLVeZksYMEoRgK6CSKT6TJViMPp7U,122186
+x_transformers/x_transformers.py,sha256=txdFN5266Tu-lQVMgyICMWt8azslAkxG5YL4n9tOUIo,122944
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.6.6.dist-info/METADATA,sha256=95CKrJ98X7R0hpb5D8GHSfi372UtxXDSeDaO2qB0Lrs,90445
-x_transformers-2.6.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.6.6.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.6.6.dist-info/RECORD,,
+x_transformers-2.7.0.dist-info/METADATA,sha256=HFH-y2lnS8T-KZkv27z7hBGECYipDSkgtXj9LJbLMHo,90888
+x_transformers-2.7.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.7.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.7.0.dist-info/RECORD,,