""" Cozet: Native SYNAXIM Base Model ================================= Non-transformer architecture using Symbiotic Gate (M-matrix) instead of self-attention. Trained from scratch. Architecture per layer: RMSNorm -> Q/K/V proj -> Symbiotic Gate (M update + retrieval) -> O proj -> Residual RMSNorm -> SwiGLU MLP -> Residual (c) 2026 GRRN Research. All rights reserved. """ import math import torch import torch.nn as nn import torch.nn.functional as F from dataclasses import dataclass from typing import Optional, List, Tuple @dataclass class CozConfig: """Cozet model configuration.""" hidden_size: int = 1024 num_layers: int = 12 num_attention_heads: int = 16 num_kv_heads: int = 4 intermediate_size: int = 4096 vocab_size: int = 32000 max_seq_len: int = 4096 rope_theta: float = 10000.0 rms_norm_eps: float = 1e-6 memory_decay: float = 0.995 tie_word_embeddings: bool = True @property def head_dim(self) -> int: return self.hidden_size // self.num_attention_heads def num_params(self) -> int: """Estimate total parameter count.""" D = self.hidden_size V = self.vocab_size I = self.intermediate_size L = self.num_layers n_kv = self.num_kv_heads hd = self.head_dim embed = V * D per_layer = ( D * D + # q_proj D * (n_kv * hd) + # k_proj D * (n_kv * hd) + # v_proj D * D + # o_proj D * I + # gate_proj D * I + # up_proj I * D + # down_proj 2 * D + # norms 2 # gate params ) lm_head = 0 if self.tie_word_embeddings else V * D total = embed + L * per_layer + lm_head + D return total # ---- Presets ---- COZET_SMALL = CozConfig( hidden_size=1024, num_layers=12, num_attention_heads=16, num_kv_heads=4, intermediate_size=4096, vocab_size=32000, max_seq_len=4096, tie_word_embeddings=True, ) COZET_MEDIUM = CozConfig( hidden_size=2048, num_layers=24, num_attention_heads=16, num_kv_heads=4, intermediate_size=8192, vocab_size=32000, max_seq_len=8192, rope_theta=500000.0, tie_word_embeddings=True, ) COZET_LARGE = CozConfig( hidden_size=4096, num_layers=32, num_attention_heads=32, num_kv_heads=8, intermediate_size=14336, vocab_size=32000, max_seq_len=8192, rope_theta=1000000.0, tie_word_embeddings=False, ) class RMSNorm(nn.Module): """Root Mean Square Layer Normalization.""" def __init__(self, dim: int, eps: float = 1e-6): super().__init__() self.weight = nn.Parameter(torch.ones(dim)) self.eps = eps def forward(self, x: torch.Tensor) -> torch.Tensor: norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() return (x.float() * norm).type_as(x) * self.weight class SymbioticGate(nn.Module): """ The Symbiotic Gate: core SYNAXIM attention replacement. Instead of softmax(QK^T/sqrt(d)) @ V, computes: gate = sigmoid(mean(Q * K) * scale * gate_scale + gate_bias) M = gate * decay * M + (1-gate) * outer(k_norm, v_norm) out = q @ M @ W_o M is persistent state: O(1) memory, infinite context. Fully differentiable for training. """ def __init__(self, config: CozConfig, layer_idx: int): super().__init__() self.D = config.hidden_size self.n_heads = config.num_attention_heads self.n_kv = config.num_kv_heads self.head_dim = config.head_dim self.decay = config.memory_decay self.layer_idx = layer_idx # Projections self.q_proj = nn.Linear(self.D, self.n_heads * self.head_dim, bias=False) self.k_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False) self.v_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False) self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.D, bias=False) # Learnable gate parameters self.gate_scale = nn.Parameter(torch.ones(1)) self.gate_bias = nn.Parameter(torch.zeros(1)) # RoPE tables (precomputed, not learned) self._rope_built = False def _build_rope(self, max_seq: int, device: torch.device): """Precompute RoPE cos/sin tables.""" if self._rope_built and max_seq <= self.rope_cos.shape[0]: return self._rope_built = False # Rebuild if we need more positions half = self.head_dim // 2 freqs = 1.0 / (10000.0 ** (torch.arange(0, half, device=device).float() / half)) t = torch.arange(max_seq, device=device).float() angles = torch.outer(t, freqs) self.register_buffer("rope_cos", angles.cos(), persistent=False) self.register_buffer("rope_sin", angles.sin(), persistent=False) self._rope_built = True def _apply_rope(self, x: torch.Tensor, pos: int) -> torch.Tensor: """Apply RoPE to (n_heads, head_dim) tensor at position pos.""" half = self.head_dim // 2 cos = self.rope_cos[pos, :half] sin = self.rope_sin[pos, :half] x1, x2 = x[..., :half], x[..., half:] return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1) def forward(self, h: torch.Tensor, M: torch.Tensor, position: int) -> Tuple[torch.Tensor, torch.Tensor]: """ Symbiotic Gate forward pass for a single token. Args: h: (D,) hidden state M: (D, D) persistent memory matrix position: token position for RoPE Returns: output: (D,) attention output M_new: (D, D) updated memory """ self._build_rope(position + 1, h.device) # Project q = self.q_proj(h) # (n_heads * hd,) k = self.k_proj(h) # (n_kv * hd,) v = self.v_proj(h) # (n_kv * hd,) # Reshape for per-head ops q_h = q.view(self.n_heads, self.head_dim) k_h = k.view(self.n_kv, self.head_dim) v_h = v.view(self.n_kv, self.head_dim) # RoPE q_h = self._apply_rope(q_h, position) k_h = self._apply_rope(k_h, position) # GQA: repeat KV heads if self.n_kv < self.n_heads: repeat = self.n_heads // self.n_kv k_h = k_h.repeat_interleave(repeat, dim=0) v_h = v_h.repeat_interleave(repeat, dim=0) # Gate score scale = 1.0 / math.sqrt(self.head_dim) gate_score = (q_h * k_h).sum(-1).mean() * scale gate_score = gate_score * self.gate_scale + self.gate_bias gate = torch.sigmoid(gate_score) # Flatten to D-space k_flat = k_h.reshape(-1) v_flat = v_h.reshape(-1) # Normalize k_norm = k_flat / (k_flat.norm() + 1e-8) v_norm = v_flat / (v_flat.norm() + 1e-8) * h.norm() # M-matrix update (differentiable) imprint = torch.outer(k_norm, v_norm) M_new = gate * self.decay * M + (1.0 - gate) * imprint # Retrieve q_flat = q_h.reshape(-1) output = q_flat @ M_new # Output projection result = self.o_proj(output) return result, M_new class SynaxBlock(nn.Module): """One SYNAXIM layer: SymbioticGate + SwiGLU MLP.""" def __init__(self, config: CozConfig, layer_idx: int): super().__init__() self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps) self.attn = SymbioticGate(config, layer_idx) self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps) # SwiGLU MLP self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) def forward(self, h: torch.Tensor, M: torch.Tensor, position: int) -> Tuple[torch.Tensor, torch.Tensor]: # Attention h_normed = self.norm_attn(h) attn_out, M_new = self.attn(h_normed, M, position) h = h + attn_out # MLP h_normed = self.norm_mlp(h) gate_out = self.gate_proj(h_normed) up_out = self.up_proj(h_normed) mlp_out = self.down_proj(F.silu(gate_out) * up_out) h = h + mlp_out return h, M_new class CozModel(nn.Module): """ Cozet: Native SYNAXIM base model. Non-transformer architecture trained from scratch. Uses Symbiotic Gate (M-matrix) instead of self-attention. O(1) memory at inference. Infinite context. No KV cache. """ def __init__(self, config: CozConfig): super().__init__() self.config = config self.D = config.hidden_size self.n_layers = config.num_layers # Embedding self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) # SYNAXIM layers self.layers = nn.ModuleList([ SynaxBlock(config, i) for i in range(config.num_layers) ]) # Final norm self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps) # LM head if not config.tie_word_embeddings: self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) # Initialize weights self._init_weights() def _init_weights(self): """SYNAXIM-specific weight initialization.""" depth_scale = 1.0 / math.sqrt(2 * self.n_layers) for name, p in self.named_parameters(): if p.dim() == 1: continue # Norms already initialized to ones if "embed" in name: nn.init.normal_(p, std=0.02) elif "down_proj" in name or "o_proj" in name: nn.init.normal_(p, std=0.02 * depth_scale) elif "lm_head" in name: nn.init.normal_(p, std=0.02 * depth_scale) elif p.dim() == 2: nn.init.normal_(p, std=0.02) def init_m_states(self, device: torch.device) -> List[torch.Tensor]: """Create fresh M-matrix states for a new sequence.""" return [torch.zeros(self.D, self.D, device=device) for _ in range(self.n_layers)] def forward_token(self, token_id: int, M_states: List[torch.Tensor], position: int) -> Tuple[torch.Tensor, List[torch.Tensor]]: """ Process one token through the full model. Returns: (logits, updated_M_states) """ h = self.embed_tokens.weight[token_id] for i, layer in enumerate(self.layers): h, M_states[i] = layer(h, M_states[i], position) h = self.final_norm(h) if self.config.tie_word_embeddings: logits = h @ self.embed_tokens.weight.T else: logits = self.lm_head(h) return logits, M_states def forward_sequence(self, token_ids: torch.Tensor, chunk_size: int = 256) -> torch.Tensor: """ Process a sequence with truncated BPTT. Args: token_ids: (seq_len,) token IDs chunk_size: gradient truncation window Returns: total_loss: scalar loss for the sequence """ device = token_ids.device seq_len = token_ids.shape[0] M_states = self.init_m_states(device) total_loss = torch.tensor(0.0, device=device) n_tokens = 0 for chunk_start in range(0, seq_len - 1, chunk_size): chunk_end = min(chunk_start + chunk_size, seq_len - 1) # Detach M at chunk boundaries (truncated BPTT) M_states = [m.detach() for m in M_states] chunk_loss = torch.tensor(0.0, device=device, requires_grad=True) for t in range(chunk_start, chunk_end): logits, M_states = self.forward_token( token_ids[t].item(), M_states, t ) loss = F.cross_entropy( logits.unsqueeze(0), token_ids[t + 1].unsqueeze(0) ) chunk_loss = chunk_loss + loss n_tokens += 1 total_loss = total_loss + chunk_loss return total_loss / max(n_tokens, 1) def create_model(size: str = "small") -> CozModel: """Create a Cozet model by size name.""" configs = { "small": COZET_SMALL, "medium": COZET_MEDIUM, "large": COZET_LARGE, } config = configs.get(size, COZET_SMALL) model = CozModel(config) n_params = sum(p.numel() for p in model.parameters()) print(f"Cozet-{size}: {n_params:,} parameters") print(f" D={config.hidden_size}, L={config.num_layers}, " f"H={config.num_attention_heads}/{config.num_kv_heads} GQA, " f"I={config.intermediate_size}, V={config.vocab_size}") return model if __name__ == "__main__": # Quick sanity check model = create_model("small") device = torch.device("cpu") M_states = model.init_m_states(device) # Forward one token logits, M_states = model.forward_token(42, M_states, 0) print(f"Logits shape: {logits.shape}") print(f"M[0] norm after 1 token: {M_states[0].norm():.4f}") # Forward second token logits2, M_states = model.forward_token(100, M_states, 1) print(f"M[0] norm after 2 tokens: {M_states[0].norm():.4f}") # Test gradient flow M_states_fresh = model.init_m_states(device) logits3, _ = model.forward_token(42, M_states_fresh, 0) loss = F.cross_entropy(logits3.unsqueeze(0), torch.tensor([100])) loss.backward() grad_norms = {n: p.grad.norm().item() for n, p in model.named_parameters() if p.grad is not None} print(f"Gradient flow: {len(grad_norms)} parameters received gradients") print(f"Loss: {loss.item():.4f}") print("Sanity check PASSED.")