x-transformers 2.11.1__py3-none-any.whl → 2.11.4__py3-none-any.whl

x_transformers/free_transformer.py

@@ -128,19 +128,19 @@ class FreeTransformer(Module):
         dim,
         dec_head_depth,
         dec_tail_depth,
-        enc_depth,
         max_seq_len,
+        enc_depth = 1,
         dim_latent = None,
         attn_dim_head = 64,
         heads = 8,
         latent_bits = 16,
+        per_token_latents = True,  # they use a latent per token in the sequence, instead of one for entire sequence, iiuc
         kl_loss_threshold = NAT,
         binary_mapper_kwargs: dict = dict(),
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
         kl_loss_weight = 1.,
         pad_id = -1,
-        encoder: Module | None = None,
         **kwargs
     ):
         super().__init__()
@@ -150,33 +150,30 @@ class FreeTransformer(Module):
 
         self.token_unembed = nn.Linear(dim, num_tokens, bias = False)
 
-        if not exists(encoder):
-            encoder = Encoder(
-                dim = dim,
-                depth = enc_depth,
-                attn_dim_head = attn_dim_head,
-                heads = heads,
-                **kwargs,
-                **enc_kwargs
-            )
+        self.query_token_for_latents = nn.Parameter(torch.randn(dim) * 1e-2)
 
-        self.encoder = encoder
+        self.per_token_latents = per_token_latents
 
-        self.to_latent_bit_logits = nn.Sequential(
-            Reduce('b n d -> b d', 'mean'),
-            nn.Linear(dim, latent_bits, bias = False),
+        self.encoder = Encoder(
+            dim = dim,
+            depth = enc_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            only_cross = True,
+            cross_attend = True,
+            **kwargs,
+            **enc_kwargs
         )
 
+        self.to_latent_bit_logits = nn.Linear(dim, latent_bits, bias = False)
+
         self.binary_mapper = BinaryMapper(
             latent_bits,
             kl_loss_threshold,
             **binary_mapper_kwargs
         )
 
-        self.from_latent_to_condition = nn.Sequential(
-            nn.Linear(2 ** latent_bits, dim, bias = False),
-            Rearrange('b d -> b 1 d')
-        )
+        self.from_latent_to_condition = nn.Linear(self.binary_mapper.num_codes, dim, bias = False)
 
         self.decoder_head = Decoder(
             dim = dim,
@@ -208,11 +205,34 @@ class FreeTransformer(Module):
 
     def encode_to_latents(
         self,
-        seq,
+        decoder_head_embeds,
         mask = None,
         return_kl_loss = False
     ):
-        pooled = self.encoder(seq, mask = mask)
+        batch, seq_len, device = *decoder_head_embeds.shape[:2], decoder_head_embeds.device
+
+        query_tokens = repeat(self.query_token_for_latents, 'd -> b 1 d', b = batch)
+
+        encoder_kwargs = dict()
+
+        # handle the interesting per query token latents, as in the paper
+
+        if self.per_token_latents:
+            query_tokens = repeat(query_tokens, 'b 1 d -> b n d', n = seq_len)
+
+            rotary_pos = torch.arange(seq_len, device = device)
+
+            encoder_kwargs.update(
+                pos = rotary_pos,
+                context_pos = rotary_pos
+            )
+
+        pooled = self.encoder(
+            query_tokens,
+            context = decoder_head_embeds,
+            context_mask = mask,
+            **encoder_kwargs
+        )
 
         bit_logits = self.to_latent_bit_logits(pooled)
 

x_transformers/gpt_vae.py

@@ -44,6 +44,7 @@ class GPTVAE(Module):
         enc_kwargs: dict = dict(),
         dec_kwargs: dict = dict(),
         vae_kl_loss_weight = 1.,
+        vae_kl_div_floor = 0.,      # what was done in free transformer, which in turn came from Kingma 2016
         latents_dropout_prob = 0.5, # what percentage of the time to dropout the latents completely
         pad_id = -1,
         encoder: Module | None = None,
@@ -99,6 +100,7 @@ class GPTVAE(Module):
 
         # loss weights - vae kl loss
 
+        self.vae_kl_div_floor = vae_kl_div_floor
         self.vae_kl_loss_weight = vae_kl_loss_weight
 
         self.latents_dropout = nn.Dropout(latents_dropout_prob)
@@ -190,12 +192,16 @@ class GPTVAE(Module):
 
         # vae kl loss
 
-        vae_kl_loss = (
+        vae_kl_loss = 0.5 * (
             latents_log_var.exp()
             + latents_mean.square()
             - latents_log_var
             - 1.
-        ).sum(dim = -1).mean()
+        )
+
+        vae_kl_loss = F.relu(vae_kl_loss - self.vae_kl_div_floor)
+
+        vae_kl_loss = vae_kl_loss.sum(dim = -1).mean()
 
         # return losses
 

{x_transformers-2.11.1.dist-info → x_transformers-2.11.4.dist-info}/METADATA

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

{x_transformers-2.11.1.dist-info → x_transformers-2.11.4.dist-info}/RECORD

@@ -5,8 +5,8 @@ x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTN
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
-x_transformers/free_transformer.py,sha256=K9dX9xj0pJgiOi9jlOLCD9Nn9eYNNgTWj9YvQLhexHw,8295
-x_transformers/gpt_vae.py,sha256=myYSgcx66V0M4zeEGKyhY1P2HlPDHcezhaZEoo_uMdo,5715
+x_transformers/free_transformer.py,sha256=t-PyeiqNq-B8BZUGICECyWCeU5XLPSGbFsQ3ICxTtsM,9072
+x_transformers/gpt_vae.py,sha256=4QdznXZcU7pmMXUeEocAOKpcTkREYS-zDHktN5ADtNk,5981
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
@@ -14,7 +14,7 @@ x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,
 x_transformers/x_transformers.py,sha256=ADr83Fz2cehj_F7N1bMwxhAg-r48fGhlaZqw3hxoxMQ,125765
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.11.1.dist-info/METADATA,sha256=Rj4l6-wbfFsC7wbSWfUFQyTNWQE5EXu952aSE3B8uas,96011
-x_transformers-2.11.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.11.1.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.11.1.dist-info/RECORD,,
+x_transformers-2.11.4.dist-info/METADATA,sha256=18u2dW5d3aJeH3LaVsRRA_vf1gZzbSjS0Fbd5XM4dJc,96011
+x_transformers-2.11.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.11.4.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.11.4.dist-info/RECORD,,