Add TRM-Nano model implementation (h=96, 4 heads, 2 layers, unrolled recursion)"

trm_solver/trm_nano/model.py

@@ -0,0 +1,211 @@
+"""
+TRM-Nano: Tiny Recursive Model for ARC-AGI NeuroGolf 2026.
+
+Based on Samsung TRM (arxiv:2510.04871).
+Adapted to fit 1.44MB ONNX constraint with unrolled recursion.
+
+Architecture: h=96, 4 heads, 2 layers SwiGLU, T=3 outer, n=4 inner = 15 net passes.
+~200K params → ~0.8MB FP32 ONNX.
+
+No banned ops (Loop, Scan, NonZero, Unique, Script, Function).
+All recursion unrolled statically at export time.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+
+
+@dataclass
+class TRMNanoConfig:
+    vocab_size: int = 12       # PAD=0, EOS=1, colors 2-11
+    hidden_size: int = 96
+    num_heads: int = 4         # head_dim = 24
+    expansion: int = 4
+    num_layers: int = 2        # single 2-layer block (recursed)
+    H_cycles: int = 3          # outer recursions (T)
+    L_cycles: int = 4          # inner recursions per cycle (n)
+    seq_len: int = 900         # 30x30 flattened grid
+    puzzle_emb_len: int = 8    # learned puzzle embedding tokens
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000.0
+
+
+def rms_norm(x: torch.Tensor, eps: float = 1e-5) -> torch.Tensor:
+    return x * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps).to(x.dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    q_embed = (q * cos.unsqueeze(-2)) + (rotate_half(q) * sin.unsqueeze(-2))
+    k_embed = (k * cos.unsqueeze(-2)) + (rotate_half(k) * sin.unsqueeze(-2))
+    return q_embed, k_embed
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, hidden_size: int, expansion: int):
+        super().__init__()
+        inter = ((expansion * hidden_size * 2 // 3 + 63) // 64) * 64
+        self.gate_up = nn.Linear(hidden_size, inter * 2, bias=False)
+        self.down = nn.Linear(inter, hidden_size, bias=False)
+
+    def forward(self, x):
+        gate, up = self.gate_up(x).chunk(2, dim=-1)
+        return self.down(F.silu(gate) * up)
+
+
+class Attention(nn.Module):
+    def __init__(self, hidden_size: int, num_heads: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.qkv_proj = nn.Linear(hidden_size, 3 * hidden_size, bias=False)
+        self.o_proj = nn.Linear(hidden_size, hidden_size, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, S, _ = x.shape
+        qkv = self.qkv_proj(x).view(B, S, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+        out = out.transpose(1, 2).contiguous().view(B, S, self.hidden_size)
+        return self.o_proj(out)
+
+
+class TRMBlock(nn.Module):
+    """Post-norm residual: Attention + RMSNorm + SwiGLU + RMSNorm."""
+    def __init__(self, config: TRMNanoConfig):
+        super().__init__()
+        self.attn = Attention(config.hidden_size, config.num_heads)
+        self.mlp = SwiGLU(config.hidden_size, config.expansion)
+        self.eps = config.rms_norm_eps
+
+    def forward(self, x, cos, sin):
+        x = rms_norm(x + self.attn(x, cos, sin), self.eps)
+        x = rms_norm(x + self.mlp(x), self.eps)
+        return x
+
+
+class ReasoningModule(nn.Module):
+    """Stack of TRM blocks with input injection."""
+    def __init__(self, config: TRMNanoConfig):
+        super().__init__()
+        self.layers = nn.ModuleList([TRMBlock(config) for _ in range(config.num_layers)])
+
+    def forward(self, hidden, injection, cos, sin):
+        x = hidden + injection
+        for layer in self.layers:
+            x = layer(x, cos, sin)
+        return x
+
+
+class TRMNano(nn.Module):
+    """
+    Tiny Recursive Model - Nano variant for NeuroGolf.
+
+    Unrolled recursion (no Loop/Scan):
+        for T in range(H_cycles):
+            for n in range(L_cycles):
+                z_L = net(z_L, z_H + x_embed)
+            z_H = net(z_H, z_L)
+        output = lm_head(z_H)
+    """
+
+    def __init__(self, config: TRMNanoConfig):
+        super().__init__()
+        self.config = config
+        self.embed_scale = math.sqrt(config.hidden_size)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.puzzle_emb = nn.Parameter(torch.zeros(1, config.puzzle_emb_len, config.hidden_size))
+
+        # RoPE precomputed
+        total_len = config.seq_len + config.puzzle_emb_len
+        head_dim = config.hidden_size // config.num_heads
+        inv_freq = 1.0 / (config.rope_theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        t = torch.arange(total_len).float()
+        freqs = torch.outer(t, inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos())
+        self.register_buffer("sin_cached", emb.sin())
+
+        # Single reasoning module (weight-shared across all recursion steps)
+        self.net = ReasoningModule(config)
+
+        # Initial latent states
+        self.H_init = nn.Parameter(torch.randn(config.hidden_size) * 0.02)
+        self.L_init = nn.Parameter(torch.randn(config.hidden_size) * 0.02)
+
+        # Output head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
+        """
+        Args: tokens [B, 900] int (vocab 0-11)
+        Returns: logits [B, 900, 12]
+        """
+        B = tokens.shape[0]
+        total_len = self.config.seq_len + self.config.puzzle_emb_len
+
+        tok_emb = self.embed_tokens(tokens) * self.embed_scale
+        puz_emb = self.puzzle_emb.expand(B, -1, -1)
+        x_embed = torch.cat([puz_emb, tok_emb], dim=1)  # [B, 908, H]
+
+        z_H = self.H_init.view(1, 1, -1).expand(B, total_len, -1).clone()
+        z_L = self.L_init.view(1, 1, -1).expand(B, total_len, -1).clone()
+
+        cos = self.cos_cached
+        sin = self.sin_cached
+
+        # Unrolled recursion
+        for _t in range(self.config.H_cycles):
+            for _n in range(self.config.L_cycles):
+                z_L = self.net(z_L, z_H + x_embed, cos, sin)
+            z_H = self.net(z_H, z_L, cos, sin)
+
+        logits = self.lm_head(z_H[:, self.config.puzzle_emb_len:])
+        return logits
+
+    def forward_neurogolf(self, onehot_input: torch.Tensor) -> torch.Tensor:
+        """
+        NeuroGolf I/O: [1,10,30,30] → [1,10,30,30]
+        """
+        grid_ids = onehot_input.argmax(dim=1)       # [1, 30, 30] values 0-9
+        tokens = (grid_ids.view(1, -1) + 2).long()  # [1, 900] values 2-11
+        logits = self.forward(tokens)                # [1, 900, 12]
+        color_logits = logits[:, :, 2:12]            # [1, 900, 10]
+        color_probs = F.softmax(color_logits, dim=-1)
+        return color_probs.view(1, 30, 30, 10).permute(0, 3, 1, 2)
+
+
+def count_params(model: nn.Module) -> int:
+    return sum(p.numel() for p in model.parameters())
+
+
+if __name__ == "__main__":
+    config = TRMNanoConfig()
+    model = TRMNano(config)
+    n = count_params(model)
+    print(f"TRM-Nano params: {n:,}")
+    print(f"Estimated ONNX size: {n * 4 / 1024 / 1024 * 1.1:.2f} MB (limit: 1.44 MB)")
+
+    # Test forward
+    tokens = torch.randint(0, 12, (2, 900))
+    logits = model(tokens)
+    print(f"Forward: {tokens.shape} -> {logits.shape}")
+
+    # Test NeuroGolf
+    onehot = torch.zeros(1, 10, 30, 30)
+    onehot[0, 0] = 1.0
+    out = model.forward_neurogolf(onehot)
+    print(f"NeuroGolf: {onehot.shape} -> {out.shape}")