x-transformers 2.1.2__py3-none-any.whl → 2.1.5__py3-none-any.whl

x_transformers/belief_state.py

@@ -0,0 +1,162 @@
+# 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,
+        train_frac_forward_backward_pairs: float = 1.
+    ):
+        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 number of tokens'
+
+        dim = forward_decoder.emb_dim
+        num_tokens = forward_decoder.num_tokens
+
+        # the suffix token
+
+        self.suffix_token = nn.Parameter(torch.zeros(dim))
+        nn.init.normal_(self.suffix_token, std = 0.02)
+
+        # 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
+
+        # what fraction of forward backward pairs to train on
+        # for further memory efficiency
+
+        assert 0 < train_frac_forward_backward_pairs <= 1.
+        self.train_frac_fb_pairs = train_frac_forward_backward_pairs
+        self.needs_subsample_fb_pairs = train_frac_forward_backward_pairs < 1.
+
+    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]
+
+        # maybe subsample fb pairs
+
+        if self.needs_subsample_fb_pairs:
+            num_pairs = fb_pairs.shape[0]
+
+            num_subsampled = max(int(num_pairs * self.train_frac_fb_pairs), 1)
+
+            rand_subsampled_indices = torch.randperm(num_pairs, device = device)[:num_subsampled]
+
+            fb_pairs = fb_pairs[rand_subsampled_indices]
+
+        # 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-2.1.2.dist-info → x_transformers-2.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.2
+Version: 2.1.5
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2435,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.2.dist-info → x_transformers-2.1.5.dist-info}/RECORD

@@ -1,6 +1,7 @@
 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=GfYDeDqmhldozECgFsJ9zhd6O5NMvdYA5OwueVs8SB4,4742
 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
@@ -9,7 +10,7 @@ x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dC
 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.2.dist-info/METADATA,sha256=-LsNGhf7qKzttPNU7VOSVqigs61_Nuw4r0LBlZDT_Qo,87227
-x_transformers-2.1.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.1.2.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.1.2.dist-info/RECORD,,
+x_transformers-2.1.5.dist-info/METADATA,sha256=-nQpm1eBGBXkLEXGiQK06NaIWa13CEnBujBNfTzvnJ8,87570
+x_transformers-2.1.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.1.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.1.5.dist-info/RECORD,,