x-transformers 2.3.20__py3-none-any.whl → 2.3.22__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -140,7 +140,8 @@ class AutoregressiveWrapper(Module):
         ignore_index = -100,
         pad_value = 0,
         mask_prob = 0.,
-        add_attn_z_loss = False
+        add_attn_z_loss = False,
+        next_embed_loss_weight = 0.1
     ):
         super().__init__()
         self.pad_value = pad_value
@@ -156,6 +157,10 @@ class AutoregressiveWrapper(Module):
         # whether to add router z-loss
         self.add_attn_z_loss = add_attn_z_loss
 
+        # whether to add a continuous loss
+        self.add_continuous_pred_head = net.add_continuous_pred_head
+        self.next_embed_loss_weight = next_embed_loss_weight
+
     @torch.no_grad()
     @eval_decorator
     def generate(
@@ -305,9 +310,9 @@ class AutoregressiveWrapper(Module):
         return out
 
     def forward(self, x, return_outputs = False, **kwargs):
-        seq, ignore_index, add_attn_z_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss
+        seq, ignore_index, add_attn_z_loss, add_next_embed_loss = x.shape[1], self.ignore_index, self.add_attn_z_loss, self.add_continuous_pred_head
 
-        inp, target = x[:, :-1], x[:, 1:]
+        inp, target = x, x[:, 1:]
         inp = torch.where(inp == ignore_index, self.pad_value, inp)
 
         if self.mask_prob > 0.:
@@ -318,15 +323,29 @@ class AutoregressiveWrapper(Module):
             mask = ~torch.zeros_like(inp).scatter(1, indices, 1.).bool()
             kwargs.update(self_attn_kv_mask = mask)
 
-        logits, cache = self.net(
-            inp,
+        out, cache = self.net(
+            x,
             return_intermediates = True,
             return_attn_z_loss = add_attn_z_loss,
+            return_next_embed_pred = add_next_embed_loss,
             **kwargs
         )
 
+        # destruct differently if doing continuous pred
+
+        if add_next_embed_loss:
+            logits, (next_embed_pred, init_embeds) = out
+        else:
+            logits = out
+
+        logits = logits[:, :-1]
+
+        # loss function
+
         loss_fn = F.cross_entropy if not self.net.output_is_log_prob else F.nll_loss
 
+        # cross entropy loss
+
         loss = loss_fn(
             rearrange(logits, 'b n c -> b c n'),
             target,
@@ -336,6 +355,16 @@ class AutoregressiveWrapper(Module):
         if add_attn_z_loss:
             loss = loss + cache.attn_z_loss
 
+        if add_next_embed_loss:
+            mask = inp[:, :-1] != ignore_index
+            embed_pred = next_embed_pred[:, :-1]
+            cont_targets = init_embeds[:, 1:].detach()
+
+            cont_loss = F.l1_loss(embed_pred, cont_targets, reduction = 'none')
+            cont_loss = cont_loss[mask].mean()
+
+            loss = loss + cont_loss * self.next_embed_loss_weight
+
         if not return_outputs:
             return loss
 

x_transformers/continuous.py

@@ -141,6 +141,8 @@ class ContinuousTransformerWrapper(Module):
         sum_embeds = None,
         prepend_embeds = None,
         prepend_mask = None,
+        cache: LayerIntermediates | None = None,
+        input_not_include_cache = False,
         seq_start_pos = None,
         **kwargs
     ):
@@ -154,10 +156,17 @@ class ContinuousTransformerWrapper(Module):
 
             mask = einx.less('j, i -> i j', seq_arange, lens)
 
+        # take care of position embedding offsets in the presence of cache and sequence is less than cache length (not full sequence)
+
+        seq_pos_offset = 0
+
+        if exists(cache) and input_not_include_cache:
+            seq_pos_offset = cache.cache_length
+
         # project in + positional embedding
 
         x = self.project_in(x)
-        x = x + self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos)
+        x = x + self.pos_emb(x, pos = pos, seq_start_pos = seq_start_pos, offset = seq_pos_offset)
 
         if exists(sum_embeds):
             x = x + sum_embeds
@@ -193,7 +202,7 @@ class ContinuousTransformerWrapper(Module):
 
         # attention layers
 
-        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, return_hiddens = True, **kwargs)
+        x, intermediates = self.attn_layers(x, mask = mask, mems = mems, mem_masks = mem_masks, cache = cache, input_not_include_cache = input_not_include_cache, seq_pos_offset = seq_pos_offset, return_hiddens = True, **kwargs)
 
         # splice out memory tokens
 

x_transformers/x_transformers.py

@@ -2855,6 +2855,7 @@ class TransformerWrapper(Module):
         sigsoftmax_logits = False,
         ff_deep_embed = False,
         to_logits: Module | None = None,
+        add_continuous_pred_head = False
     ):
         super().__init__()
 
@@ -2975,6 +2976,18 @@ class TransformerWrapper(Module):
         else:
             self.to_logits = LinearNoBias(dim, logits_dim) if not exists(to_logits) else to_logits
 
+        # add a head that predicts the embedding of the next step
+
+        self.add_continuous_pred_head = add_continuous_pred_head
+
+        if add_continuous_pred_head:
+
+            self.to_next_embed_pred = nn.Sequential(
+                LinearNoBias(dim, dim),
+                nn.SiLU(),
+                LinearNoBias(dim, dim)
+            )
+
         # memory tokens (like [cls]) from Memory Transformers paper
 
         num_memory_tokens = default(num_memory_tokens, 0)
@@ -3009,6 +3022,7 @@ class TransformerWrapper(Module):
         return_intermediates = False,
         return_embeddings_and_intermediates = False,
         return_logit_entropies = False,
+        return_next_embed_pred = False,
         mask = None,
         return_mems = False,
         return_attn = False,
@@ -3100,6 +3114,10 @@ class TransformerWrapper(Module):
             assert emb_frac_gradient > 0
             x = x * emb_frac_gradient + x.detach() * (1 - emb_frac_gradient)
 
+        # init embed
+
+        init_embed = x
+
         # embedding dropout
 
         x = self.emb_dropout(x)
@@ -3261,6 +3279,14 @@ class TransformerWrapper(Module):
         else:
             out = logits
 
+        # maybe next embed pred
+
+        if return_next_embed_pred:
+            assert self.add_continuous_pred_head
+            next_embed_out = self.to_next_embed_pred(x)
+
+            out = (out, (next_embed_out, init_embed))
+
         # logit entropies
 
         if return_logit_entropies:

{x_transformers-2.3.20.dist-info → x_transformers-2.3.22.dist-info}/METADATA

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

{x_transformers-2.3.20.dist-info → x_transformers-2.3.22.dist-info}/RECORD

@@ -1,17 +1,17 @@
 x_transformers/__init__.py,sha256=h3I2ejobgEdy8H7NgV-rP8UaBCnd16-MysvDXH9GMEA,985
 x_transformers/attend.py,sha256=fXMuwHuBAFB4f4_U6j5_uVeK7N4cV0PDd6UTqtkjKKM,17333
-x_transformers/autoregressive_wrapper.py,sha256=LW1gr3cFONDEPA_HHhaTE7mk-JWbaINuB1fc_DfbCqw,10791
+x_transformers/autoregressive_wrapper.py,sha256=BWFaO-3YWzCcEfp-EC1ZkdckqDpPIOQG6_uyyP6AmhM,11753
 x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTNY__GGZEKU,13374
-x_transformers/continuous.py,sha256=CHta8vizKl85n220fv5278fwjSU-vrN_FBy-m831_go,12551
+x_transformers/continuous.py,sha256=hpb1sSbt3k2LNzzjrjSd8F5xOIbKj7IluV9MBEAFLkw,13031
 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/x_transformers.py,sha256=l2p-r0iJNlYHUB3vM4lb6ptzNCx9HgA7UfgieEcQT6w,115521
+x_transformers/x_transformers.py,sha256=7phSZvP1_SDRIkVMwVR4cz1dFU2UlR2Wf1HJHEQlcQg,116222
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.3.20.dist-info/METADATA,sha256=ygWyfnlIh2Mw6bd12gJjjZJyM9vfnXmvvOLyrd2El2k,89897
-x_transformers-2.3.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.3.20.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.3.20.dist-info/RECORD,,
+x_transformers-2.3.22.dist-info/METADATA,sha256=_8m8ftpHRKjbEUDuoeYPcVh4yan1FxNRj3seJwiZzl8,89897
+x_transformers-2.3.22.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.3.22.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.3.22.dist-info/RECORD,,