x-transformers 2.1.1__py3-none-any.whl → 2.1.4__py3-none-any.whl

x_transformers/belief_state.py

@@ -0,0 +1,142 @@
+# Belief State Transformer
+
+# https://arxiv.org/abs/2410.23506
+# https://www.youtube.com/watch?v=aqhbRtB2Fyg
+
+import torch
+from torch.autograd import Function
+from torch.nn import Module, ModuleList
+from torch import nn, cat, stack, arange, cartesian_prod
+import torch.nn.functional as F
+
+from x_transformers.x_transformers import (
+    Decoder,
+    TransformerWrapper
+)
+
+import einx
+from einops import rearrange, repeat
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# wrappers
+
+class BeliefStateWrapper(Module):
+    """
+    Figure 13. in https://arxiv.org/abs/2410.23506
+    """
+
+    def __init__(
+        self,
+        forward_decoder: TransformerWrapper,
+        backward_decoder: TransformerWrapper
+    ):
+        super().__init__()
+        assert forward_decoder.emb_dim == backward_decoder.emb_dim, 'forward and backwards model must have the same embedding dimension'
+        assert forward_decoder.num_tokens == backward_decoder.num_tokens, 'forward and backwards model must have the same embedding dimension'
+
+        dim = forward_decoder.emb_dim
+        num_tokens = forward_decoder.num_tokens
+
+        # the suffix token
+
+        self.suffix_token = nn.Parameter(torch.zeros(dim))
+
+        # the text prediction head, which predicts for the combinations of prefix and suffix the next and previous token for forwards and backward sequences
+
+        self.text_head = nn.Sequential(
+            nn.Linear(dim * 2, dim),
+            nn.LeakyReLU(),
+            nn.Linear(dim, num_tokens * 2),
+        )
+
+        # the two decoders, one which is causal forward, the other causal backwards
+
+        self.forward_decoder = forward_decoder
+        self.backward_decoder = backward_decoder
+
+    def forward(
+        self,
+        seq
+    ):
+        batch, seq_len, device = *seq.shape, seq.device
+
+        # forward autoregressive
+
+        forward_embeds = self.forward_decoder(seq, return_embeddings = True)
+
+        # backward autoregressive
+
+        backward_seq = seq.flip(1)
+
+        suffix_tokens = repeat(self.suffix_token, 'd -> b 1 d', b = batch)
+
+        backward_embeds = self.backward_decoder(
+            backward_seq,
+            prepend_embeds = suffix_tokens,
+            return_embeddings = True
+        )
+
+        backward_embeds = backward_embeds.flip(1)
+
+        # trick to reduce memory on backwards pass
+
+        orig_forward_embeds, forward_embeds = forward_embeds, forward_embeds.detach()
+        orig_backward_embeds, backward_embeds = backward_embeds, backward_embeds.detach()
+
+        forward_embeds.requires_grad_()
+        backward_embeds.requires_grad_()
+
+        # belief state objective
+
+        seq_arange = arange(seq_len, device = device)
+
+        fb_pairs = cartesian_prod(seq_arange, seq_arange)
+
+        # filter down to valid pairs, as in figure 11
+        # f - forward, b - backward, i - indices
+
+        fi, bi = fb_pairs.unbind(dim = -1)
+        valid_mask = (bi - fi) >= 2
+
+        fb_pairs = fb_pairs[valid_mask]
+
+        # get labels for both
+
+        fi, bi = fb_pairs.unbind(dim = -1)
+
+        labels_fi, labels_bi = (fi + 1), bi
+
+        forward_labels, backward_labels = seq[:, fi], seq[:, bi]
+        labels = stack((forward_labels, backward_labels), dim = -1)
+
+        # get the forward and backward embedding pairs and feed them through the text head for both forward and backward predictions
+
+        fb_embeds = cat((
+            forward_embeds[:, fi],
+            backward_embeds[:, bi]
+        ), dim = -1)
+
+        logits = self.text_head(fb_embeds)
+
+        # cross entropy loss
+
+        fb_loss = F.cross_entropy(
+            rearrange(logits, 'b n (fb l) -> b l (fb n)', fb = 2),
+            rearrange(labels, 'b n fb -> b (fb n)')
+        )
+
+        # backwards
+
+        fb_loss.backward()
+
+        orig_forward_embeds.backward(forward_embeds.grad)
+        orig_backward_embeds.backward(backward_embeds.grad)
+
+        return fb_loss

x_transformers/x_transformers.py

@@ -970,6 +970,7 @@ class DynamicLIMe(Module):
         dim,
         num_layers,
         num_views = 1,
+        norm = True,
         use_softmax = True
     ):
         super().__init__()
@@ -977,6 +978,7 @@ class DynamicLIMe(Module):
         self.multiple_views = num_views > 1
 
         self.to_weights = Sequential(
+            RMSNorm(dim) if norm else None,
             nn.Linear(dim, num_views * num_layers),
             Rearrange('... (views layers) -> views ... layers', views = num_views),
             nn.Softmax(dim = -1) if use_softmax else nn.ReLU()
@@ -987,6 +989,7 @@ class DynamicLIMe(Module):
         x,
         hiddens
     ):
+
         if not is_tensor(hiddens):
             hiddens = stack(hiddens)
 

{x_transformers-2.1.1.dist-info → x_transformers-2.1.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.1
+Version: 2.1.4
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -41,7 +41,6 @@ Requires-Dist: packaging>=21.0
 Requires-Dist: torch>=2.0
 Provides-Extra: examples
 Requires-Dist: lion-pytorch; extra == 'examples'
-Requires-Dist: torchvision; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'
 Provides-Extra: test
 Requires-Dist: pytest; extra == 'test'
@@ -2436,4 +2435,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Hu2024TheBS,
+    title   = {The Belief State Transformer},
+    author  = {Edward S. Hu and Kwangjun Ahn and Qinghua Liu and Haoran Xu and Manan Tomar and Ada Langford and Dinesh Jayaraman and Alex Lamb and John Langford},
+    year    = {2024},
+    url     = {https://api.semanticscholar.org/CorpusID:273707334}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.1.1.dist-info → x_transformers-2.1.4.dist-info}/RECORD

@@ -1,15 +1,16 @@
 x_transformers/__init__.py,sha256=l0dom8ZYkRzFvnDdgzDboXqrI1tKav3beVE7TN2nHko,844
 x_transformers/attend.py,sha256=-5BWWhFsp7tvZTdN91Ay5SqOjyj9uOs-122vFvoO6b4,17253
 x_transformers/autoregressive_wrapper.py,sha256=reLCno9Z9pchVU79tBF8OMo21LwSZ67KAeB83jqkyAc,10505
+x_transformers/belief_state.py,sha256=5E_08m6kvXROYgROazTJFxuUsDPNjIVM3AJxg3CJNmU,3966
 x_transformers/continuous.py,sha256=p0sCAiH1na236ygwgL1Yyhu36eZBf9cZvoW1JyP_fFE,7073
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 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/x_transformers.py,sha256=bIlP-NHj0SB2joklpxicoaD1HVpRMGIulMF8WYEsOAQ,110076
+x_transformers/x_transformers.py,sha256=-80N4sqUr3sR51Ms4wCfc4jhxnPwf0ApNR4xfIsasfQ,110142
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=7S00kCuab4tWQa-vf-z-XfzADjVj48MoFIr7VSIvttg,8575
-x_transformers-2.1.1.dist-info/METADATA,sha256=BBGKnocyDvj_ynWM5dtrbyX1iodI4eWEnn9TWrw38kc,87275
-x_transformers-2.1.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.1.dist-info/RECORD,,
+x_transformers-2.1.4.dist-info/METADATA,sha256=-jme9jyXVeVlo1T7nPi2iGFfoxfjSwdj2D33_dr4yxQ,87570
+x_transformers-2.1.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.4.dist-info/RECORD,,