x-transformers 2.11.2__tar.gz → 2.11.5__tar.gz

{x_transformers-2.11.2 → x_transformers-2.11.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.11.2
+Version: 2.11.5
 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.2 → x_transformers-2.11.5}/pyproject.toml

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

{x_transformers-2.11.2 → x_transformers-2.11.5}/tests/test_x_transformers.py

@@ -586,12 +586,12 @@ def test_cross_attn_rotary(
     context_pos = torch.arange(128) if cross_attn_rotary else None
 
     embed = model(
-      x = x,
-      mask = mask,
-      context = context,
-      pos = pos,
-      context_pos = context_pos,
-      context_mask = context_mask
+        x = x,
+        mask = mask,
+        context = context,
+        pos = pos,
+        context_pos = context_pos,
+        context_mask = context_mask
     )
 
 @param('tanh', (True, False))
@@ -1408,3 +1408,29 @@ def test_attn_negative_weights(
     x = torch.randint(0, 256, (1, 10))
 
     logits = model(x)
+
+@param('per_token_latents', (False, True))
+def test_free(
+    per_token_latents
+):
+    from x_transformers.free_transformer import FreeTransformer
+
+    model = FreeTransformer(
+        num_tokens = 256,
+        max_seq_len = 1024,
+        dim = 512,
+        heads = 8,
+        dec_head_depth = 4,
+        dec_tail_depth = 4,
+        enc_depth = 3,
+        kl_loss_weight = 1.,
+        per_token_latents = per_token_latents,
+        latent_bits = 8
+    )
+
+    seq = torch.randint(0, 256, (1, 1024))
+
+    loss, (ar_loss, aux_loss) = model(seq, return_all_losses = True)
+    loss.backward()
+
+    assert aux_loss.numel() == 1

{x_transformers-2.11.2 → x_transformers-2.11.5}/train_free.py

@@ -54,11 +54,11 @@ model = FreeTransformer(
     max_seq_len = SEQ_LEN,
     dim = 512,
     heads = 8,
-    rotary_pos_emb = True,
     dec_head_depth = 4,
     dec_tail_depth = 4,
     enc_depth = 3,
     kl_loss_weight = 1.,
+    per_token_latents = True,
     kl_loss_threshold = NAT,
     latent_bits = LATENT_BITS
 ).cuda()

{x_transformers-2.11.2 → x_transformers-2.11.5}/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,39 +150,40 @@ 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,
+            use_rmsnorm = True,
+            rotary_pos_emb = 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,
             depth = dec_head_depth,
             attn_dim_head = attn_dim_head,
             heads = heads,
+            rotary_pos_emb = True,
+            use_rmsnorm = True,
             pre_norm_has_final_norm = False,
             **kwargs,
             **dec_kwargs
@@ -193,6 +194,8 @@ class FreeTransformer(Module):
             depth = dec_tail_depth,
             attn_dim_head = attn_dim_head,
             heads = heads,
+            rotary_pos_emb = True,
+            use_rmsnorm = True,
             pre_norm_has_final_norm = True,
             **kwargs,
             **dec_kwargs
@@ -208,11 +211,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)