""" NeuroLex v4 — Creative Name Diffusion Engine ============================================= A Uniform Discrete Language Diffusion Model (UDLM) for creative name generation. CORE INSIGHT: Why this works when autoregressive models fail: ───────────────────────────────────────────────────────────── Problem with AR models (GPT, NeuroLex v3): • Left-to-right generation creates probability feedback loops • Once model enters a high-probability mode, it stays there (Holtzman et al. 2019) • Small models can only memorize ~100 "patterns" → repeats them endlessly • Temperature/top-p CANNOT fix this: they only scale the same collapsed distribution • The model assigns ever-higher probability to seen sequences → positive feedback loop UDLM Solution (this model): • Generation starts from RANDOM tokens (uniform noise over vocab) • Iteratively denoises ALL positions SIMULTANEOUSLY in random order • Each noise seed → completely different denoising path → different output • Bidirectional attention sees full context at every step (not just left context) • Classifier-Free Guidance steers toward conditions without mode collapse • ODD (Orthogonal Diversity Diffusion) repels batch samples from each other Architecture: Diffusion Transformer (DiT) on character sequences • Vocab: ~72 chars (a-z, A-Z, 0-9, special + PAD/MASK/BOS/EOS) • Sequence length: 24 (covers most brand/channel names) • Conditioning: language, domain, style, length — via adaptive LayerNorm (adaLN) • Parameters: ~8-12M (fits free Colab T4 16GB easily) Based on: • MDLM (Sahoo et al., NeurIPS 2024, arxiv:2406.07524) • Discrete CFG (Kuleshov Group, arxiv:2412.10193) • ODD diversity sampling (arxiv:2603.04893) • GFlowNet reward-proportional concepts (arxiv:2106.04399) • SimCTG contrastive framework (Su et al., NeurIPS 2022, arxiv:2202.06417) License: Apache 2.0 """ import torch import torch.nn as nn import torch.nn.functional as F import math from dataclasses import dataclass, field from typing import List, Dict, Optional, Tuple import random import json # ═══════════════════════════════════════════════════════════════ # CONFIGURATION # ═══════════════════════════════════════════════════════════════ @dataclass class NeuroLexConfig: """Configuration for NeuroLex v4 Diffusion Model.""" # Character vocabulary vocab_size: int = 72 # a-z(26) + A-Z(26) + 0-9(10) + special(6) + PAD/MASK/BOS/EOS(4) max_seq_len: int = 24 # max name length # Transformer architecture d_model: int = 256 n_heads: int = 8 n_layers: int = 8 d_ff: int = 1024 dropout: float = 0.1 # Conditioning dimensions n_languages: int = 25 # 25 language families n_domains: int = 20 # tech, food, gaming, luxury, etc. n_styles: int = 10 # sharp, warm, elegant, playful, futuristic, etc. n_lengths: int = 12 # target lengths 3-14 # Diffusion parameters n_diffusion_steps: int = 100 # T (number of denoising steps at inference) noise_schedule: str = "cosine" # cosine or linear cfg_dropout: float = 0.15 # classifier-free guidance dropout during training # Training learning_rate: float = 3e-4 weight_decay: float = 0.01 warmup_steps: int = 500 batch_size: int = 256 epochs: int = 30 # Sampling cfg_scale: float = 2.5 # classifier-free guidance strength temperature: float = 0.9 odd_alpha: float = 8.0 # ODD diversity repulsion strength # ═══════════════════════════════════════════════════════════════ # CHARACTER TOKENIZER (No subword tokenization - pure characters) # ═══════════════════════════════════════════════════════════════ class CharTokenizer: """ Pure character-level tokenizer with explicit vocabulary. Why character-level instead of BPE/subword: • Brand names are NOVEL words — subword tokenizers would split them unpredictably or map to UNK • Character-level forces the model to learn PHONOTACTIC patterns (which character combinations sound good) rather than word chunks • Enables generating truly new character sequences, not just recombining known subwords """ # Special tokens PAD = 0 MASK = 1 # The "noise" token for diffusion BOS = 2 EOS = 3 def __init__(self): # Build character vocabulary self.special_tokens = ['', '', '', ''] self.chars = list('abcdefghijklmnopqrstuvwxyz') # 26 self.chars += list('ABCDEFGHIJKLMNOPQRSTUVWXYZ') # 26 self.chars += list('0123456789') # 10 self.chars += list('-_. &+') # 6 special chars useful in names self.vocab = self.special_tokens + self.chars self.char_to_id = {c: i for i, c in enumerate(self.vocab)} self.id_to_char = {i: c for i, c in enumerate(self.vocab)} self.vocab_size = len(self.vocab) def encode(self, text: str, max_len: int = 24) -> List[int]: """Encode text to token IDs with padding.""" ids = [self.BOS] for ch in text[:max_len - 2]: # leave room for BOS/EOS ids.append(self.char_to_id.get(ch, self.PAD)) ids.append(self.EOS) # Pad to max_len ids += [self.PAD] * (max_len - len(ids)) return ids def decode(self, ids: List[int]) -> str: """Decode token IDs to text, stopping at EOS.""" text = [] for i in ids: if i == self.EOS: break if i in (self.PAD, self.MASK, self.BOS): continue text.append(self.id_to_char.get(i, '')) return ''.join(text) def batch_encode(self, texts: List[str], max_len: int = 24) -> torch.Tensor: """Encode a batch of texts.""" return torch.tensor([self.encode(t, max_len) for t in texts]) def batch_decode(self, ids: torch.Tensor) -> List[str]: """Decode a batch of token ID tensors.""" return [self.decode(row.tolist()) for row in ids] # ═══════════════════════════════════════════════════════════════ # CONDITIONING SYSTEM # ═══════════════════════════════════════════════════════════════ # Domain categories (what the name is for) DOMAINS = [ 'tech', 'food', 'gaming', 'luxury', 'health', 'finance', 'education', 'music', 'sports', 'fashion', 'travel', 'crypto', 'eco', 'entertainment', 'social', 'ai', 'automotive', 'beauty', 'fitness', 'general' ] # Style/vibe of the name STYLES = [ 'sharp', 'warm', 'elegant', 'playful', 'futuristic', 'bold', 'minimal', 'organic', 'mystical', 'professional' ] # Language influences (phonotactic patterns to draw from) LANGUAGES = [ 'english', 'spanish', 'french', 'german', 'italian', 'japanese', 'korean', 'mandarin', 'arabic', 'hindi', 'portuguese', 'russian', 'turkish', 'swedish', 'dutch', 'greek', 'latin', 'finnish', 'hawaiian', 'swahili', 'thai', 'vietnamese', 'polish', 'czech', 'esperanto' ] DOMAIN_TO_ID = {d: i for i, d in enumerate(DOMAINS)} STYLE_TO_ID = {s: i for i, s in enumerate(STYLES)} LANG_TO_ID = {l: i for i, l in enumerate(LANGUAGES)} # ═══════════════════════════════════════════════════════════════ # ADAPTIVE LAYER NORM (adaLN) — DiT-style conditioning # ═══════════════════════════════════════════════════════════════ class AdaptiveLayerNorm(nn.Module): """ Adaptive Layer Normalization (adaLN) from DiT paper. Instead of fixed scale/shift, we LEARN them from the conditioning signal. This allows the model to modulate its behavior based on: • What diffusion timestep we're at (how noisy is the input) • What domain/style/language we want This is far more powerful than prepending control tokens (CTRL-style) because it modulates EVERY layer's computation, not just the input. """ def __init__(self, d_model: int, cond_dim: int): super().__init__() self.norm = nn.LayerNorm(d_model, elementwise_affine=False) # Project condition → 2 * d_model (scale and shift) self.projection = nn.Sequential( nn.SiLU(), nn.Linear(cond_dim, 2 * d_model) ) def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor: """ x: [B, L, D] — sequence features cond: [B, D_cond] — conditioning vector """ # Get adaptive scale and shift params = self.projection(cond) # [B, 2*D] scale, shift = params.chunk(2, dim=-1) # each [B, D] # Apply adaptive normalization x = self.norm(x) x = x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) return x # ═══════════════════════════════════════════════════════════════ # DIFFUSION TRANSFORMER BLOCK # ═══════════════════════════════════════════════════════════════ class DiTBlock(nn.Module): """ Diffusion Transformer block with adaLN conditioning. Key differences from standard transformer block: 1. adaLN replaces regular LayerNorm (condition-dependent normalization) 2. No causal mask — BIDIRECTIONAL attention (sees full name at once) 3. This enables the model to consider how ALL characters interact, not just the left context """ def __init__(self, d_model: int, n_heads: int, d_ff: int, cond_dim: int, dropout: float = 0.1): super().__init__() # Adaptive norms self.norm1 = AdaptiveLayerNorm(d_model, cond_dim) self.norm2 = AdaptiveLayerNorm(d_model, cond_dim) # Multi-head self-attention (BIDIRECTIONAL - no causal mask!) self.attn = nn.MultiheadAttention( d_model, n_heads, dropout=dropout, batch_first=True ) # Feed-forward with GELU (smoother than ReLU for character patterns) self.ff = nn.Sequential( nn.Linear(d_model, d_ff), nn.GELU(), nn.Dropout(dropout), nn.Linear(d_ff, d_model), nn.Dropout(dropout) ) self.dropout = nn.Dropout(dropout) def forward(self, x: torch.Tensor, cond: torch.Tensor, pad_mask: Optional[torch.Tensor] = None) -> torch.Tensor: """ x: [B, L, D] cond: [B, cond_dim] pad_mask: [B, L] True = padded position (ignore) """ # Self-attention with adaLN residual = x x = self.norm1(x, cond) x, _ = self.attn(x, x, x, key_padding_mask=pad_mask) x = self.dropout(x) + residual # Feed-forward with adaLN residual = x x = self.norm2(x, cond) x = self.ff(x) + residual return x # ═══════════════════════════════════════════════════════════════ # SINUSOIDAL TIME EMBEDDING # ═══════════════════════════════════════════════════════════════ class TimeEmbedding(nn.Module): """ Sinusoidal embedding for diffusion timestep. Converts continuous t ∈ [0,1] into a rich d-dimensional vector that tells each transformer layer "how noisy is the current input". The model needs to know this because: • At t=1 (fully noisy): it should make bold predictions • At t=0 (nearly clean): it should make conservative refinements """ def __init__(self, d_model: int): super().__init__() self.d_model = d_model self.mlp = nn.Sequential( nn.Linear(d_model, d_model * 4), nn.GELU(), nn.Linear(d_model * 4, d_model) ) def forward(self, t: torch.Tensor) -> torch.Tensor: """t: [B] float in [0, 1]""" half_dim = self.d_model // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, device=t.device) * -emb) emb = t.unsqueeze(-1) * emb.unsqueeze(0) * 1000 # scale up emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) return self.mlp(emb) # ═══════════════════════════════════════════════════════════════ # MAIN MODEL: NeuroLex v4 — UDLM Diffusion Transformer # ═══════════════════════════════════════════════════════════════ class NeuroLexV4(nn.Module): """ NeuroLex v4: Uniform Discrete Language Diffusion Model (UDLM) This is NOT an autoregressive model. It's a denoising model that: 1. Takes a FULLY NOISY name (random characters) 2. Iteratively predicts what the CLEAN name should be 3. Gradually replaces noisy tokens with predicted ones Each generation starts from different random noise → different denoising path → different output name. This STRUCTURALLY prevents the repetition problem because: • There is no "most likely next token" cascade • The model sees ALL positions at once (bidirectional) • Stochastic noise injection at each step adds variation • CFG guidance steers gently rather than collapsing to modes """ def __init__(self, config: NeuroLexConfig): super().__init__() self.config = config # Token embedding self.token_embed = nn.Embedding(config.vocab_size, config.d_model) # Positional embedding (learnable, for short sequences) self.pos_embed = nn.Embedding(config.max_seq_len, config.d_model) # Time embedding (tells model how noisy the input is) self.time_embed = TimeEmbedding(config.d_model) # Condition embeddings self.domain_embed = nn.Embedding(config.n_domains, config.d_model // 4) self.style_embed = nn.Embedding(config.n_styles, config.d_model // 4) self.lang_embed = nn.Embedding(config.n_languages, config.d_model // 4) self.length_embed = nn.Embedding(config.n_lengths, config.d_model // 4) # Null condition embeddings (for CFG — learned "no condition" vectors) self.null_domain = nn.Parameter(torch.randn(config.d_model // 4) * 0.02) self.null_style = nn.Parameter(torch.randn(config.d_model // 4) * 0.02) self.null_lang = nn.Parameter(torch.randn(config.d_model // 4) * 0.02) self.null_length = nn.Parameter(torch.randn(config.d_model // 4) * 0.02) # Condition projection: concat all condition embeds → cond_dim cond_dim = config.d_model self.cond_proj = nn.Sequential( nn.Linear(config.d_model + config.d_model, cond_dim), nn.GELU(), nn.Linear(cond_dim, cond_dim) ) # Transformer blocks with adaLN self.blocks = nn.ModuleList([ DiTBlock(config.d_model, config.n_heads, config.d_ff, cond_dim, config.dropout) for _ in range(config.n_layers) ]) # Final layer norm and output projection self.final_norm = nn.LayerNorm(config.d_model) self.output_proj = nn.Linear(config.d_model, config.vocab_size) # Initialize weights self._init_weights() def _init_weights(self): """Xavier initialization for stable training.""" for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) def get_condition_vector(self, domain_id, style_id, lang_id, length_id, t_embed, cfg_mask=None): """Build conditioning vector from all attributes + time.""" B = domain_id.shape[0] d_emb = self.domain_embed(domain_id) s_emb = self.style_embed(style_id) l_emb = self.lang_embed(lang_id) len_emb = self.length_embed(length_id) # Apply CFG dropout if cfg_mask is not None: mask = cfg_mask.unsqueeze(-1).float() d_emb = d_emb * (1 - mask) + self.null_domain.unsqueeze(0) * mask s_emb = s_emb * (1 - mask) + self.null_style.unsqueeze(0) * mask l_emb = l_emb * (1 - mask) + self.null_lang.unsqueeze(0) * mask len_emb = len_emb * (1 - mask) + self.null_length.unsqueeze(0) * mask cond = torch.cat([d_emb, s_emb, l_emb, len_emb], dim=-1) combined = torch.cat([cond, t_embed], dim=-1) return self.cond_proj(combined) def forward(self, z_t, t, domain_id, style_id, lang_id, length_id, cfg_mask=None): """ Forward pass: predict clean x_0 from noisy z_t at time t. z_t: [B, L] — noisy token indices t: [B] — diffusion time in [0, 1] Returns: [B, L, V] — logits predicting the clean token at each position """ B, L = z_t.shape x = self.token_embed(z_t) + self.pos_embed(torch.arange(L, device=z_t.device)) t_embed = self.time_embed(t) cond = self.get_condition_vector(domain_id, style_id, lang_id, length_id, t_embed, cfg_mask) pad_mask = (z_t == CharTokenizer.PAD) for block in self.blocks: x = block(x, cond, pad_mask) x = self.final_norm(x) logits = self.output_proj(x) return logits def compute_loss(self, x_0, domain_id, style_id, lang_id, length_id): """ UDLM Training Loss. 1. Take clean name x_0 2. Sample random time t ~ Uniform[0, 1] 3. Corrupt x_0 by replacing tokens with uniform random chars 4. Train model to predict x_0 from z_t 5. With probability cfg_dropout, drop all conditions (for CFG) """ B, L = x_0.shape device = x_0.device # Sample diffusion time t = torch.rand(B, device=device) # Cosine noise schedule if self.config.noise_schedule == "cosine": alpha_t = torch.cos(t * math.pi / 2) ** 2 else: alpha_t = 1 - t # Forward diffusion: replace tokens with uniform random noise noise_mask = torch.rand(B, L, device=device) > alpha_t.unsqueeze(-1) # Don't corrupt special tokens special_mask = (x_0 == CharTokenizer.PAD) | (x_0 == CharTokenizer.BOS) | (x_0 == CharTokenizer.EOS) noise_mask = noise_mask & ~special_mask # Replace with uniform random tokens random_tokens = torch.randint(4, self.config.vocab_size, (B, L), device=device) z_t = torch.where(noise_mask, random_tokens, x_0) # CFG dropout cfg_mask = torch.rand(B, device=device) < self.config.cfg_dropout # Forward pass logits = self.forward(z_t, t, domain_id, style_id, lang_id, length_id, cfg_mask) # Loss on noised positions (Rao-Blackwellized ELBO) loss_per_token = F.cross_entropy( logits.reshape(-1, self.config.vocab_size), x_0.reshape(-1), reduction='none' ).reshape(B, L) # Weight: focus on noised positions + small signal on clean ones weight = noise_mask.float() + 0.1 * (~noise_mask & ~special_mask).float() # Time-dependent importance weighting time_weight = 1.0 / (t.unsqueeze(-1) + 0.01) loss = (loss_per_token * weight * time_weight).sum() / (weight * time_weight).sum() return loss @torch.no_grad() def generate(self, domain_id: int, style_id: int, lang_id: int, target_length: int = 8, batch_size: int = 16, cfg_scale: float = 2.5, temperature: float = 0.9, n_steps: int = 80, odd_alpha: float = 8.0, device: str = 'cuda') -> List[str]: """ Generate names using UDLM denoising with CFG and ODD diversity. Algorithm: 1. Start from fully random tokens (uniform noise) 2. At each step, predict clean x_0 with CFG guidance 3. Confidence-based progressive unmasking 4. ODD: repel each sample's predictions from others in batch """ self.eval() tokenizer = CharTokenizer() d_ids = torch.full((batch_size,), domain_id, device=device, dtype=torch.long) s_ids = torch.full((batch_size,), style_id, device=device, dtype=torch.long) l_ids = torch.full((batch_size,), lang_id, device=device, dtype=torch.long) len_ids = torch.full((batch_size,), min(target_length - 3, 11), device=device, dtype=torch.long) # Initialize: BOS + random chars + EOS + PAD seq_len = min(target_length + 4, self.config.max_seq_len) x = torch.randint(4, self.config.vocab_size, (batch_size, seq_len), device=device) x[:, 0] = CharTokenizer.BOS eos_pos = min(target_length + 1, seq_len - 1) x[:, eos_pos] = CharTokenizer.EOS x[:, eos_pos+1:] = CharTokenizer.PAD # Denoising loop for step in range(n_steps): t_val = 1.0 - step / n_steps t = torch.full((batch_size,), t_val, device=device) # === CLASSIFIER-FREE GUIDANCE === logits_cond = self.forward(x, t, d_ids, s_ids, l_ids, len_ids, cfg_mask=torch.zeros(batch_size, device=device, dtype=torch.bool)) logits_uncond = self.forward(x, t, d_ids, s_ids, l_ids, len_ids, cfg_mask=torch.ones(batch_size, device=device, dtype=torch.bool)) logits = logits_uncond + cfg_scale * (logits_cond - logits_uncond) # === ODD: ORTHOGONAL DIVERSITY DIFFUSION === if odd_alpha > 0 and batch_size > 1: logits_flat = logits.reshape(batch_size, -1) logits_norm = F.normalize(logits_flat, dim=-1) # Repel each sample from the mean of others batch_mean = logits_norm.mean(dim=0, keepdim=True) for i in range(batch_size): others_mean = (logits_norm.sum(0) - logits_norm[i]) / (batch_size - 1) others_mean = F.normalize(others_mean.unsqueeze(0), dim=-1) proj = (logits_flat[i:i+1] @ others_mean.T) * others_mean logits_flat[i] = logits_flat[i] - odd_alpha * proj.squeeze(0) * t_val logits = logits_flat.reshape(batch_size, seq_len, -1) # === SAMPLING === logits = logits / temperature logits[:, 1:-1, CharTokenizer.PAD] = -float('inf') logits[:, 1:-1, CharTokenizer.MASK] = -float('inf') logits[:, 1:-1, CharTokenizer.BOS] = -float('inf') probs = F.softmax(logits, dim=-1) predicted = torch.multinomial( probs.reshape(-1, self.config.vocab_size), 1 ).reshape(batch_size, seq_len) # === PROGRESSIVE UNMASKING === confidence = probs.max(dim=-1).values update_prob = (1.0 - t_val) * confidence update_prob[:, 0] = 0 update_prob[:, eos_pos] = 0 update_prob[:, eos_pos+1:] = 0 should_update = torch.bernoulli(update_prob).bool() x = torch.where(should_update, predicted, x) # Keep structure fixed x[:, 0] = CharTokenizer.BOS x[:, eos_pos] = CharTokenizer.EOS x[:, eos_pos+1:] = CharTokenizer.PAD # Final prediction t_final = torch.zeros(batch_size, device=device) logits_final = self.forward(x, t_final, d_ids, s_ids, l_ids, len_ids, cfg_mask=torch.zeros(batch_size, device=device, dtype=torch.bool)) logits_final = logits_final / 0.7 logits_final[:, :, :4] = -float('inf') # block special tokens final_preds = logits_final.argmax(dim=-1) x[:, 1:eos_pos] = final_preds[:, 1:eos_pos] x[:, 0] = CharTokenizer.BOS x[:, eos_pos] = CharTokenizer.EOS x[:, eos_pos+1:] = CharTokenizer.PAD # Decode names = tokenizer.batch_decode(x) results = [] for name in names: name = name.strip() if len(name) >= 3: name = name[0].upper() + name[1:] results.append(name) return results def count_parameters(self) -> int: return sum(p.numel() for p in self.parameters() if p.requires_grad) # ═══════════════════════════════════════════════════════════════ # MODEL FACTORY # ═══════════════════════════════════════════════════════════════ def create_model(size: str = 'base') -> Tuple[NeuroLexV4, NeuroLexConfig]: """ Create model with preset sizes. Sizes (all fit in free Colab T4 16GB): 'tiny': ~2M params — fast experiments, proof of concept 'small': ~5M params — good balance for most use cases 'base': ~12M params — best quality (recommended) 'large': ~25M params — maximum quality (needs more data) """ configs = { 'tiny': NeuroLexConfig(d_model=128, n_heads=4, n_layers=4, d_ff=512), 'small': NeuroLexConfig(d_model=192, n_heads=6, n_layers=6, d_ff=768), 'base': NeuroLexConfig(d_model=256, n_heads=8, n_layers=8, d_ff=1024), 'large': NeuroLexConfig(d_model=384, n_heads=12, n_layers=10, d_ff=1536), } config = configs.get(size, configs['base']) model = NeuroLexV4(config) print(f"NeuroLex v4 ({size}) created:") print(f" Parameters: {model.count_parameters():,}") print(f" Architecture: {config.n_layers}L / {config.n_heads}H / d={config.d_model}") print(f" Vocab size: {config.vocab_size}") return model, config if __name__ == '__main__': model, config = create_model('base') print(f"\nTotal parameters: {model.count_parameters():,}") tokenizer = CharTokenizer() x = tokenizer.batch_encode(['Nexaflow', 'Datavex', 'Sparkify'], max_len=24) t = torch.rand(3) d = torch.tensor([0, 1, 2]) s = torch.tensor([0, 1, 2]) l = torch.tensor([0, 0, 0]) length = torch.tensor([5, 5, 5]) logits = model(x, t, d, s, l, length) print(f"Output shape: {logits.shape}") loss = model.compute_loss(x, d, s, l, length) print(f"Loss: {loss.item():.4f}")