x-transformers 2.1.4__tar.gz → 2.1.6__tar.gz

{x_transformers-2.1.4 → x_transformers-2.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.1.4
+Version: 2.1.6
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.1.4 → x_transformers-2.1.6}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.1.4"
+version = "2.1.6"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.1.4 → x_transformers-2.1.6}/tests/test_x_transformers.py

@@ -695,7 +695,7 @@ def test_lime(
     model(x)
 
 def test_belief_state_wrapper():
-    from x_transformers.belief_state import BeliefStateWrapper
+    from x_transformers.belief_state_wrapper import BeliefStateWrapper
 
     forward_model = TransformerWrapper(
         num_tokens = 20000,
@@ -721,9 +721,11 @@ def test_belief_state_wrapper():
 
     model = BeliefStateWrapper(
         forward_decoder = forward_model,
-        backward_decoder = backward_model
+        backward_decoder = backward_model,
+        backward_ar_loss_weight = 0.5
     )
 
     seq = torch.randint(0, 20000, (2, 16))
 
-    loss = model(seq)
+    loss = model(seq, backward = False)
+    loss.backward()

{x_transformers-2.1.4 → x_transformers-2.1.6}/x_transformers/__init__.py

@@ -14,6 +14,7 @@ from x_transformers.x_transformers import (
 
 from x_transformers.autoregressive_wrapper import AutoregressiveWrapper
 from x_transformers.nonautoregressive_wrapper import NonAutoregressiveWrapper
+from x_transformers.belief_state_wrapper import BeliefStateWrapper
 
 from x_transformers.continuous import (
     ContinuousTransformerWrapper,

x_transformers-2.1.4/x_transformers/belief_state.py → x_transformers-2.1.6/x_transformers/belief_state_wrapper.py

@@ -6,7 +6,7 @@
 import torch
 from torch.autograd import Function
 from torch.nn import Module, ModuleList
-from torch import nn, cat, stack, arange, cartesian_prod
+from torch import nn, cat, stack, tensor, arange, cartesian_prod
 import torch.nn.functional as F
 
 from x_transformers.x_transformers import (
@@ -35,11 +35,13 @@ class BeliefStateWrapper(Module):
     def __init__(
         self,
         forward_decoder: TransformerWrapper,
-        backward_decoder: TransformerWrapper
+        backward_decoder: TransformerWrapper,
+        train_frac_forward_backward_pairs: float = 1.,
+        backward_ar_loss_weight: float = 1. # can weigh the training of the backwards decoder differently, perhaps fwd/bwd have a shared backbone etc etc
     ):
         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'
+        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
@@ -47,6 +49,7 @@ class BeliefStateWrapper(Module):
         # 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
 
@@ -61,9 +64,24 @@ class BeliefStateWrapper(Module):
         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.
+
+        # loss weighting
+
+        self.backward_ar_loss_weight = backward_ar_loss_weight
+        self.needs_loss_weight = backward_ar_loss_weight != 1.
+
+        self.register_buffer('loss_weights', tensor([1., self.backward_ar_loss_weight]))
+
     def forward(
         self,
-        seq
+        seq,
+        backward = True
     ):
         batch, seq_len, device = *seq.shape, seq.device
 
@@ -107,6 +125,17 @@ class BeliefStateWrapper(Module):
 
         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)
@@ -129,14 +158,31 @@ class BeliefStateWrapper(Module):
 
         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)')
+            rearrange(labels, 'b n fb -> b (fb n)'),
+            reduction = 'none' if self.needs_loss_weight else 'mean'
         )
 
+        # maybe loss weighting
+
+        if self.needs_loss_weight:
+            fb_loss = rearrange(fb_loss, 'b (fb n) -> b fb n')
+            fb_loss = fb_loss * self.fwd_bwd_loss_weights
+            fb_loss = fb_loss.mean()
+
         # backwards
 
-        fb_loss.backward()
+        orig_backward = getattr(fb_loss, 'backward')
+
+        def patched_backward_fn(*args, **kwargs):
+            orig_backward(*args, **kwargs)
+            orig_forward_embeds.backward(forward_embeds.grad)
+            orig_backward_embeds.backward(backward_embeds.grad)
+
+        # can allow the researcher to call .backward from the outside
 
-        orig_forward_embeds.backward(forward_embeds.grad)
-        orig_backward_embeds.backward(backward_embeds.grad)
+        if backward:
+            patched_backward_fn()
+        else:
+            setattr(fb_loss, 'backward', patched_backward_fn)
 
         return fb_loss