x-transformers 2.10.2__tar.gz → 2.11.0__tar.gz

{x_transformers-2.10.2 → x_transformers-2.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.10.2
+Version: 2.11.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2598,4 +2598,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Fleuret2025TheFT,
+    title   = {The Free Transformer},
+    author  = {Franccois Fleuret},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:282210283}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.10.2 → x_transformers-2.11.0}/README.md

@@ -2549,4 +2549,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Fleuret2025TheFT,
+    title   = {The Free Transformer},
+    author  = {Franccois Fleuret},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:282210283}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.10.2 → x_transformers-2.11.0}/pyproject.toml

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

{x_transformers-2.10.2 → x_transformers-2.11.0}/train_copy.py

@@ -12,15 +12,16 @@ GENERATE_EVERY  = 100
 NUM_TOKENS = 16 + 2
 ENC_SEQ_LEN = 32
 DEC_SEQ_LEN = 64 + 1
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # helpers
 
 def cycle():
     while True:
-        prefix = torch.ones((BATCH_SIZE, 1)).long()
-        src = torch.randint(2, NUM_TOKENS, (BATCH_SIZE, ENC_SEQ_LEN)).long()
+        prefix = torch.ones((BATCH_SIZE, 1)).long().to(DEVICE)
+        src = torch.randint(2, NUM_TOKENS, (BATCH_SIZE, ENC_SEQ_LEN)).long().to(DEVICE)
         tgt = torch.cat((prefix, src, src), 1)
-        src_mask = torch.ones(BATCH_SIZE, src.shape[1]).bool()
+        src_mask = torch.ones(BATCH_SIZE, src.shape[1]).bool().to(DEVICE)
         yield (src, tgt, src_mask)
 
 # instantiate model
@@ -39,7 +40,7 @@ model = XTransformer(
     dec_heads = 8,
     dec_max_seq_len = DEC_SEQ_LEN,
     dec_attn_cog_signed = True
-)
+).to(DEVICE)
 
 # optimizer
 
@@ -63,7 +64,7 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
         model.eval()
         src, _, src_mask = next(cycle())
         src, src_mask = src[:1], src_mask[:1]
-        start_tokens = (torch.ones((1, 1)) * 1).long()
+        start_tokens = (torch.ones((1, 1)) * 1).long().to(DEVICE)
 
         sample = model.generate(src, start_tokens, ENC_SEQ_LEN, mask = src_mask)
         incorrects = (src != sample).long().abs().sum()

x_transformers-2.11.0/train_free.py

@@ -0,0 +1,134 @@
+# /// script
+# dependencies = [
+#   "tqdm",
+#   "x-transformers>=2.11.0",
+# ]
+# ///
+
+from x_transformers.free_transformer import FreeTransformer
+
+from math import log
+import random
+import tqdm
+import gzip
+import numpy as np
+
+import torch
+import torch.optim as optim
+from torch import tensor
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 250
+GENERATE_LENGTH = 512
+PRIME_LENGTH = 32
+SEQ_LEN = 512
+
+LATENT_BITS = 8
+NAT = log(2)
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate GPT-like decoder model
+
+model = FreeTransformer(
+    num_tokens = 256,
+    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.,
+    kl_loss_threshold = NAT,
+    latent_bits = LATENT_BITS
+).cuda()
+
+rand_index = torch.randint(0, 2 ** LATENT_BITS, ())
+latents = F.one_hot(rand_index, 2 ** LATENT_BITS).float().cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    data = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
+    train_x, valid_x = np.split(data, [int(90e6)])
+    data_train, data_val = torch.from_numpy(train_x), torch.from_numpy(valid_x)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
+        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True))
+val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True))
+
+# optimizer
+
+optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss, (ar_loss, vae_kl_loss) = model(next(train_loader), return_all_losses = True)
+        (loss / GRADIENT_ACCUMULATE_EVERY).backward()
+
+    print(f'training loss: {ar_loss.item():.4f}\t| kl loss: {vae_kl_loss.item():.4f}')
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            loss, (ar_loss, _) = model(next(val_loader), return_all_losses = True)
+            print(f'validation loss: {ar_loss.item():.4f}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:PRIME_LENGTH]
+        prime = decode_tokens(inp)
+        print(f'%s \n\n %s', (prime, '*' * 100))
+
+        sample = model.generate(
+            prompts = inp,
+            seq_len = GENERATE_LENGTH,
+            latents = latents
+        )
+
+        output_str = decode_tokens(sample)
+
+        print(f'\n\nlatent {rand_index.tolist()} - ', output_str)

{x_transformers-2.10.2 → x_transformers-2.11.0}/x_transformers/autoregressive_wrapper.py

@@ -43,6 +43,10 @@ def log(t, eps = 1e-20):
 def gumbel_noise(t):
     return -log(-log(torch.rand_like(t)))
 
+def gumbel_sample(logits, temperature = 1., eps = 1e-6):
+    noise = gumbel_noise(logits)
+    return ((logits / max(temperature, eps)) + noise).argmax(dim = -1)
+
 # function for modifying all the cached key / values
 
 def modify_cached_kv(cache, fn):

x_transformers-2.11.0/x_transformers/free_transformer.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+# https://arxiv.org/abs/2510.17558
+# François Fleuret
+# https://www.youtube.com/watch?v=Nao16-6l6dQ
+
+import math
+
+import torch
+from torch import nn, Tensor, is_tensor, tensor, arange
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from x_transformers.x_transformers import (
+    Encoder,
+    Decoder,
+    TransformerWrapper
+)
+
+from x_transformers.autoregressive_wrapper import (
+    gumbel_sample,
+    top_p,
+    top_k
+)
+
+from einops.layers.torch import Rearrange, Reduce
+from einops import rearrange, reduce, repeat, einsum, pack, unpack
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
+def pack_with_inverse(t, pattern):
+    packed, ps = pack([t], pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        unpacked, = unpack(out, ps, inv_pattern)
+        return unpacked
+
+    return packed, inverse
+
+# binary mapper
+
+NAT = math.log(2)
+
+def binary_entropy(logits):
+    prob = logits.sigmoid()
+    not_prob = 1. - prob
+    return -(prob * F.logsigmoid(logits) + not_prob * F.logsigmoid(-logits)).sum(dim = -1)
+
+class BinaryMapper(Module):
+    def __init__(
+        self,
+        bits = 1,
+        kl_loss_threshold = NAT # 1 bit
+    ):
+        super().__init__()
+
+        self.bits = bits
+        self.num_codes = 2 ** bits
+        self.kl_loss_threshold = kl_loss_threshold
+
+        power_two = 2 ** arange(bits)
+        codes = (arange(self.num_codes)[:, None].bitwise_and(power_two) != 0).byte().bool()
+
+        self.register_buffer('power_two', power_two, persistent = False)
+        self.register_buffer('codes', codes, persistent = False)
+
+    def forward(
+        self,
+        logits,
+        temperature = 1.,
+        straight_through = None
+    ):
+        straight_through = default(straight_through, self.training)
+
+        assert logits.shape[-1] == self.bits, f'logits must have a last dimension of {self.bits}'
+
+        # temperature and prob for sampling
+
+        prob_for_sample = (logits / temperature).sigmoid()
+
+        # sampling
+
+        sampled_bits = (torch.rand_like(logits) <= prob_for_sample).long()
+        indices = (self.power_two * sampled_bits).sum(dim = -1)
+
+        one_hot = F.one_hot(indices, self.num_codes).float()
+
+        # return hard one hot if not training or overridden
+
+        if not straight_through:
+            return one_hot
+
+        # calculate negative entropy
+
+        kl_div = self.bits * NAT - binary_entropy(logits)
+        aux_kl_loss = F.relu(kl_div - self.kl_loss_threshold).mean()
+
+        # get the soft G for the gradients and do a straight through
+
+        soft_G = (
+            einsum(F.logsigmoid(logits), self.codes.float(), '... bits, codes bits -> ... codes') +
+            einsum(F.logsigmoid(-logits), (~self.codes).float(), '... bits, codes bits -> ... codes')
+        ).exp()
+
+        # straight through
+
+        one_hot = one_hot + soft_G - soft_G.detach()
+
+        return one_hot, aux_kl_loss
+
+# classes
+
+class FreeTransformer(Module):
+    def __init__(
+        self,
+        *,
+        num_tokens,
+        dim,
+        dec_head_depth,
+        dec_tail_depth,
+        enc_depth,
+        max_seq_len,
+        dim_latent = None,
+        attn_dim_head = 64,
+        heads = 8,
+        latent_bits = 16,
+        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__()
+        dim_latent = default(dim_latent, dim)
+
+        self.token_emb = nn.Embedding(num_tokens, dim)
+
+        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.encoder = encoder
+
+        self.to_latent_bit_logits = nn.Sequential(
+            Reduce('b n d -> b d', 'mean'),
+            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.decoder_head = Decoder(
+            dim = dim,
+            depth = dec_head_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            pre_norm_has_final_norm = False,
+            **kwargs,
+            **dec_kwargs
+        )
+
+        self.decoder_tail = Decoder(
+            dim = dim,
+            depth = dec_tail_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            pre_norm_has_final_norm = True,
+            **kwargs,
+            **dec_kwargs
+        )
+
+        self.pad_id = pad_id
+
+        self.kl_loss_weight = kl_loss_weight
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def encode_to_latents(
+        self,
+        seq,
+        mask = None,
+        return_kl_loss = False
+    ):
+        pooled = self.encoder(seq, mask = mask)
+
+        bit_logits = self.to_latent_bit_logits(pooled)
+
+        one_hot_latents, kl_loss = self.binary_mapper(bit_logits, straight_through = True)
+
+        if not return_kl_loss:
+            return one_hot_latents
+
+        return one_hot_latents, kl_loss
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompts,
+        seq_len,
+        latents = None,
+        filter_logits_fn = top_p,
+        logit_filter_kwargs: dict = dict(thres = 0.9)
+    ):
+        prompts, inverse_pack = pack_with_inverse(prompts, '* n')
+
+        batch = prompts.shape[0]
+
+        # prepend embeds
+
+        condition = None
+        if exists(latents):
+            if not is_tensor(latents):
+                latents = tensor(latents, device = self.device)
+
+            if latents.ndim == 1: # repeat latents
+                latents = repeat(latents, 'd -> b d', b = batch)
+
+            condition = self.from_latent_to_condition(latents)
+
+        # generated
+
+        prompt_len = prompts.shape[-1]
+
+        generated = prompts
+
+        tokens = self.token_emb(generated)
+
+        for _ in range(max(0, seq_len - prompt_len)):
+
+            head_embed = self.decoder_head(tokens)
+
+            if exists(condition):
+                head_embed = head_embed + condition
+
+            tail_embed = self.decoder_tail(head_embed)
+
+            tail_embed = tail_embed[:, -1]
+
+            logits = self.token_unembed(tail_embed)
+
+            logits = filter_logits_fn(logits, **logit_filter_kwargs)
+
+            sampled = gumbel_sample(logits)
+
+            generated, _ = pack((generated, sampled), 'b *')
+            tokens, _ = pack((tokens, self.token_emb(sampled)), 'b * d')
+
+        return inverse_pack(generated)
+
+    def forward(
+        self,
+        seq,
+        return_all_losses = False
+    ):
+        batch, device = seq.shape[0], seq.device
+
+        seq, labels = seq[:, :-1], seq[:, 1:]
+
+        encoder_mask = seq != self.pad_id
+
+        tokens = self.token_emb(seq)
+
+        # decoder head
+
+        tokens = self.decoder_head(tokens)
+
+        # get latent Z
+
+        latents, kl_loss = self.encode_to_latents(tokens, mask = encoder_mask, return_kl_loss = True)
+
+        condition = self.from_latent_to_condition(latents)
+
+        # decoder tail
+
+        tokens = self.decoder_tail(tokens)
+
+        # cross entropy loss
+
+        logits = self.token_unembed(tokens)
+
+        ar_loss = F.cross_entropy(
+            rearrange(logits, 'b n l -> b l n'),
+            labels,
+            ignore_index = self.pad_id
+        )
+
+        # return losses
+
+        total_loss = (
+            ar_loss +
+            kl_loss * self.kl_loss_weight
+        )
+
+        if not return_all_losses:
+            return total_loss
+
+        losses = (ar_loss, kl_loss)
+
+        return total_loss, losses