Tiny Recursive Models (TRM): When 7M Parameters Beat 671B¶

# Each cell is represented as an integer (0-9, plus special tokens)
# For a 9x9 Sudoku:
grid = [
    [5, 3, 0, 0, 7, 0, 0, 0, 0],  # Row 1: 0 means empty
    [6, 0, 0, 1, 9, 5, 0, 0, 0],  # Row 2
    # ... more rows
]

# Flattened to sequence: [5, 3, 0, 0, 7, 0, ...]
# Length: 81 tokens for full 9x9 grid
# Vocabulary size: 12 (0=pad, 1=EOS, 2-11=digits 0-9)

# Simplified TRM forward pass
def trm_forward(question_grid, answer_grid, latent_reasoning):
    # Phase 1: Embed inputs into vector space
    x = embed_question(question_grid)        # [batch, 81, 256]
    y = embed_answer(answer_grid)            # [batch, 81, 256]
    z = initialize_latent(latent_reasoning)  # [batch, 32, 256]

    # Phase 2: Recursive reasoning (8 steps)
    # Only update z (reasoning), keep x (question) and y (answer) fixed
    for step in range(8):
        combined = concatenate([x, y, z])    # [batch, 194, 256]
        output = transformer_block(combined)  # Same weights every time!
        _, _, z = split(output, [81, 81, 32])

    # Phase 3: Refinement (16 steps)
    # Now update y (answer) using the refined z (reasoning)
    for step in range(16):
        combined = concatenate([x, y, z])
        output = transformer_block(combined)  # Still same weights!
        _, y, _ = split(output, [81, 81, 32])

    # Phase 4: Decode to predictions
    predictions = reverse_embed(y)  # [batch, 81, 12] logits
    return predictions

# RNN processes sequentially
hidden = initial_state
for cell in grid:  # Must go one by one
    hidden = rnn_cell(cell, hidden)
    predictions.append(decode(hidden))

# Transformer processes in parallel
embedded_grid = embed(grid)  # All cells at once
attended = self_attention(embedded_grid)  # All-to-all connections
predictions = decode(attended)  # Parallel predictions

# Question stream - fixed throughout
x = embed_question([5, 3, 0, 0, 7, ...])  # Original puzzle
# Shape: [batch_size, 81, 256]
# Never updated during recursive processing

# Answer stream - refined in Phase 2
y = embed_initial_answer([0, 0, 0, ...])  # Start empty or random
# Shape: [batch_size, 81, 256]
# Updated during refinement phase to improve predictions

# Reasoning stream - refined in Phase 1
z = torch.randn(batch_size, 32, 256) * 0.02  # Start random
# Shape: [batch_size, 32, 256]
# 32 is arbitrary "working memory" size
# Updated during reasoning phase to build understanding

def forward_pass(x, y, z):
    # Concatenate the three streams
    combined = torch.cat([x, y, z], dim=1)  # [batch, 194, 256]
    #                     81 + 81 + 32 = 194 tokens

    # Process through transformer
    # Every token can attend to every other token across all streams!
    attended = multi_head_attention(combined)
    processed = feed_forward(attended)

    # Split back into three streams
    x_new = processed[:, :81, :]      # Question (won't use update)
    y_new = processed[:, 81:162, :]   # Answer (may use update)
    z_new = processed[:, 162:, :]     # Reasoning (may use update)

    return x_new, y_new, z_new

# Phase 1: Build understanding
# Update z-stream only, keep x and y fixed
for step in range(8):
    x_new, y_new, z_new = forward_pass(x, y, z)

    # Only keep updated z (reasoning)
    z = z_new
    # Discard x_new and y_new (don't update question or answer yet)

    # What's happening in z:
    # - Identifying constraint relationships
    # - Building logical deduction chains
    # - Preparing to update the answer

# Phase 2: Improve answer
# Update y-stream only, keep x and z fixed
for step in range(16):
    x_new, y_new, z_new = forward_pass(x, y, z)

    # Only keep updated y (answer)
    y = y_new
    # Discard x_new and z_new (question and reasoning are done)

    # What's happening in y:
    # - Filling in cells based on reasoning in z
    # - Correcting mistakes from previous iterations
    # - Increasing confidence in predictions

def train_step(question_grid, true_answer_grid):
    x = embed_question(question_grid)
    y = embed_answer(question_grid)  # Start with question as initial guess
    z = torch.randn(batch_size, 32, 256) * 0.02

    total_loss = 0

    # Deep supervision: train on multiple improvement cycles
    for supervision_step in range(16):  # Up to 16 supervision steps
        # Phase 1: Reasoning (8 steps)
        for i in range(8):
            x_new, y_new, z_new = forward_pass(x, y, z)
            z = z_new

        # Phase 2: Refinement (16 steps)
        for i in range(16):
            x_new, y_new, z_new = forward_pass(x, y, z)
            y = y_new

        # Compute loss after this improvement cycle
        predictions = reverse_embed(y)
        loss = cross_entropy(predictions, true_answer_grid)
        total_loss += loss

        # Detach to prevent backprop through all previous cycles
        # Only backprop through the last cycle
        y = y.detach()
        z = z.detach()

    # Backpropagate
    total_loss.backward()
    optimizer.step()

import torch
import torch.nn as nn
import math

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention allows the model to attend to different aspects
    of the input simultaneously.

    For Sudoku:
    - Head 1 might focus on row constraints
    - Head 2 might focus on column constraints
    - Head 3 might focus on box constraints
    - Head 4 might focus on more complex patterns
    """

    def __init__(self, d_model=256, n_heads=4, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"

        self.d_model = d_model      # 256
        self.n_heads = n_heads      # 4
        self.d_k = d_model // n_heads  # 64 per head

        # Linear projections for Q, K, V
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)

        # Output projection
        self.out_proj = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.shape

        # Step 1: Project to Q, K, V
        q = self.q_proj(x)  # [batch, seq_len, 256]
        k = self.k_proj(x)
        v = self.v_proj(x)

        # Step 2: Split into multiple heads
        # Reshape to [batch, n_heads, seq_len, d_k]
        q = q.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = k.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = v.view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Step 3: Compute attention scores
        # scores[i,j] = how much should token i attend to token j?
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        # Shape: [batch, n_heads, seq_len, seq_len]

        # Why divide by sqrt(d_k)?
        # Without scaling, dot products can get very large
        # → softmax becomes peaked → gradients vanish
        # Scaling keeps values in a nice range for softmax

        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # Step 4: Apply softmax to get attention weights
        attn_weights = torch.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Step 5: Apply attention to values
        attn_output = torch.matmul(attn_weights, v)
        # Shape: [batch, n_heads, seq_len, d_k]

        # Step 6: Concatenate heads and project
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, d_model)
        output = self.out_proj(attn_output)

        return output

class FeedForward(nn.Module):
    """
    Two-layer MLP with GELU activation.

    Typical expansion factor is 4x:
    256 → 1024 → 256

    This adds non-linearity and allows the model to process
    the attended information.
    """

    def __init__(self, d_model=256, d_ff=1024, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # Expand → Activate → Compress
        return self.linear2(self.dropout(torch.nn.functional.gelu(self.linear1(x))))

class TransformerBlock(nn.Module):
    """
    A single transformer block combining:
    1. Multi-head attention (tokens talk to each other)
    2. Add & Normalize (residual connection)
    3. Feed-forward (process information)
    4. Add & Normalize (another residual)

    TRM uses 2 of these blocks sequentially, then recurses.
    """

    def __init__(self, d_model=256, n_heads=4, d_ff=1024, dropout=0.1, use_attention=True):
        super().__init__()
        self.use_attention = use_attention

        if use_attention:
            self.attention = MultiHeadAttention(d_model, n_heads, dropout)
            self.norm1 = nn.LayerNorm(d_model)

        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Block 1: Self-attention with residual
        if self.use_attention:
            residual = x
            x = self.attention(x, mask)
            x = self.norm1(residual + self.dropout(x))

        # Block 2: Feed-forward with residual
        residual = x
        x = self.ffn(x)
        x = self.norm2(residual + self.dropout(x))

        return x

# Standard approach: loss only at the end
x, y, z = embed_inputs(question, initial_answer)
for i in range(24):  # All 24 recursive passes
    x, y, z = transformer_block(x, y, z)
predictions = decode(y)
loss = cross_entropy(predictions, ground_truth)
loss.backward()  # Backprop through all 24 passes

# TRM approach: supervision at multiple points
def train_with_deep_supervision(question, ground_truth, n_sup=16):
    x = embed_question(question)
    y = embed_answer(question)  # Start with input as initial guess
    z = torch.randn_like(latent_size)

    total_loss = 0

    for sup_step in range(n_sup):  # 16 supervision cycles
        # Do one complete reasoning + refinement cycle
        # Phase 1: Reasoning (8 steps)
        for i in range(8):
            with torch.no_grad() if sup_step < n_sup - 1 else nullcontext():
                x, y, z_new = forward_pass(x, y, z)
                z = z_new

        # Phase 2: Refinement (16 steps)
        for i in range(16):
            with torch.no_grad() if sup_step < n_sup - 1 else nullcontext():
                x, y_new, z = forward_pass(x, y, z)
                y = y_new

        # Compute loss at this supervision step
        predictions = decode(y)
        loss = cross_entropy(predictions, ground_truth)

        # Weight early supervision less than later supervision
        weight = (sup_step + 1) / n_sup  # 0.0625, 0.125, ..., 1.0
        total_loss += weight * loss

        # Detach for next cycle (prevent backprop through all history)
        y = y.detach()
        z = z.detach()

    # Backpropagate total weighted loss
    total_loss.backward()
    optimizer.step()

class TRMWithACT(nn.Module):
    def __init__(self, d_model=256):
        super().__init__()
        self.trm = TinyRecursiveModel(d_model)
        # Single halting predictor
        self.halt_head = nn.Linear(d_model, 1)

    def forward(self, x, y, z, ground_truth=None):
        # Normal TRM forward pass
        predictions = self.trm(x, y, z)

        # Predict: "Is this answer correct?"
        # Average pool the answer stream
        y_pooled = y.mean(dim=1)  # [batch, d_model]
        halt_logit = self.halt_head(y_pooled)  # [batch, 1]

        if ground_truth is not None:
            # Training: learn to predict correctness
            is_correct = (predictions.argmax(dim=-1) == ground_truth).float().mean()
            halt_loss = nn.functional.binary_cross_entropy_with_logits(
                halt_logit, is_correct.unsqueeze(0)
            )
            return predictions, halt_loss
        else:
            # Inference: use prediction to decide whether to halt
            should_halt = torch.sigmoid(halt_logit) > 0.5
            return predictions, should_halt

def train_with_act(question, ground_truth, max_sup=16):
    x, y, z = initialize_streams(question)
    total_loss = 0

    for sup_step in range(max_sup):
        # Do one reasoning + refinement cycle
        y, z = trm_cycle(x, y, z)

        # Compute loss and halting prediction
        predictions, halt_loss = model(x, y, z, ground_truth)
        prediction_loss = cross_entropy(predictions, ground_truth)

        total_loss += prediction_loss + halt_loss

        # Check if we should halt (during training, to compute statistics)
        should_halt = torch.sigmoid(model.halt_head(y.mean(dim=1))) > 0.5

        if should_halt and sup_step >= 2:  # Minimum 2 cycles
            break

        y = y.detach()
        z = z.detach()

    return total_loss / (sup_step + 1)  # Average loss per cycle used

def solve_sudoku(puzzle_grid):
    """
    Complete inference pipeline for solving a Sudoku puzzle.

    Args:
        puzzle_grid: [9, 9] numpy array with 0 for empty cells

    Returns:
        solution_grid: [9, 9] numpy array with completed solution
    """
    # Step 1: Preprocess
    # Flatten to sequence and add special tokens
    puzzle_flat = puzzle_grid.flatten()  # [81]
    puzzle_tokens = puzzle_flat + 2  # Shift by 2 (0=pad, 1=EOS, 2+=digits)
    puzzle_tensor = torch.tensor(puzzle_tokens).unsqueeze(0)  # [1, 81]

    # Step 2: Embed
    x = model.embed_tokens(puzzle_tensor)  # [1, 81, 256]
    y = model.embed_tokens(puzzle_tensor)  # Start with input as guess
    z = torch.randn(1, 32, 256) * 0.02  # Random reasoning state

    # Step 3: Reasoning Phase (8 passes)
    for i in range(8):
        # Concatenate three streams
        combined = torch.cat([x, y, z], dim=1)  # [1, 194, 256]

        # Pass through 2-layer transformer
        for transformer_block in model.transformer_blocks:
            combined = transformer_block(combined)

        # Split and update only z
        x_new, y_new, z_new = torch.split(combined, [81, 81, 32], dim=1)
        z = z_new

    # Step 4: Refinement Phase (16 passes)
    for i in range(16):
        combined = torch.cat([x, y, z], dim=1)

        for transformer_block in model.transformer_blocks:
            combined = transformer_block(combined)

        x_new, y_new, z_new = torch.split(combined, [81, 81, 32], dim=1)
        y = y_new

    # Step 5: Decode to predictions
    logits = model.reverse_embedding(y)  # [1, 81, 12]
    predictions = torch.argmax(logits, dim=-1)  # [1, 81]

    # Step 6: Post-process
    solution_tokens = predictions[0].cpu().numpy() - 2  # Shift back
    solution_grid = solution_tokens.reshape(9, 9)

    return solution_grid

# Example usage
puzzle = np.array([
    [5, 3, 0, 0, 7, 0, 0, 0, 0],
    [6, 0, 0, 1, 9, 5, 0, 0, 0],
    [0, 9, 8, 0, 0, 0, 0, 6, 0],
    # ... rest of puzzle
])

solution = solve_sudoku(puzzle)
print(solution)

# HRM: Two networks
z_L = f_L(z_L + z_H + x)  # Has x
z_H = f_H(z_L + z_H)      # No x

# TRM: One network, different inputs
z = net(x, y, z)  # Has x → updates reasoning
y = net(y, z)     # No x → updates answer

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class MultiHeadAttention(nn.Module):
    """
    Multi-head attention mechanism.

    Key insight: Different heads learn different relationships.
    For Sudoku: row constraints, column constraints, box constraints, etc.
    """

    def __init__(self, d_model=256, n_heads=4, dropout=0.1):
        super().__init__()
        assert d_model % n_heads == 0

        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        batch_size, seq_len, d_model = x.shape

        # Project and split into heads
        q = self.q_proj(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.n_heads, self.d_k).transpose(1, 2)

        # Attention scores
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.d_k)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Apply attention
        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, d_model)

        return self.out_proj(attn_output)

class FeedForward(nn.Module):
    """
    Position-wise feed-forward network.
    Applies the same transformation to each position independently.
    """

    def __init__(self, d_model=256, d_ff=1024, dropout=0.1):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)
        self.linear2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.linear2(self.dropout(F.gelu(self.linear1(x))))

class TransformerBlock(nn.Module):
    """
    Complete transformer block with attention and feed-forward.
    Includes residual connections and layer normalization.
    """

    def __init__(self, d_model=256, n_heads=4, d_ff=1024, dropout=0.1, use_attention=True):
        super().__init__()
        self.use_attention = use_attention

        if use_attention:
            self.attention = MultiHeadAttention(d_model, n_heads, dropout)
            self.norm1 = nn.LayerNorm(d_model)

        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        # Self-attention block
        if self.use_attention:
            residual = x
            x = self.attention(x, mask)
            x = self.norm1(residual + self.dropout(x))

        # Feed-forward block
        residual = x
        x = self.ffn(x)
        x = self.norm2(residual + self.dropout(x))

        return x

class TRM(nn.Module):
    """
    Tiny Recursive Model - Complete Implementation

    Key components:
    - Three streams: question (x), answer (y), reasoning (z)
    - Recursive processing: same weights applied multiple times
    - Two-phase updates: reasoning then refinement
    """

    def __init__(
        self,
        vocab_size=12,           # 0=pad, 1=EOS, 2-11=digits 0-9
        d_model=256,             # Embedding dimension
        n_heads=4,               # Attention heads
        d_ff=1024,               # Feed-forward dimension
        n_layers=2,              # Transformer layers (2 is optimal!)
        max_seq_len=512,         # Maximum sequence length
        dropout=0.1,
        n_reasoning_steps=8,     # Phase 1: update z
        n_refinement_steps=16,   # Phase 2: update y
        use_attention=True,      # False for MLP-only variant
        tie_embeddings=True      # Share input/output embeddings
    ):
        super().__init__()

        self.d_model = d_model
        self.n_reasoning_steps = n_reasoning_steps
        self.n_refinement_steps = n_refinement_steps
        self.use_attention = use_attention

        # Embeddings
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.position_embedding = nn.Embedding(max_seq_len, d_model)
        self.embedding_dropout = nn.Dropout(dropout)

        # Transformer blocks (reused recursively)
        self.transformer_blocks = nn.ModuleList([
            TransformerBlock(d_model, n_heads, d_ff, dropout, use_attention)
            for _ in range(n_layers)
        ])

        # Output projection
        self.reverse_embedding = nn.Linear(d_model, vocab_size, bias=False)

        # Weight tying: use same weights for input and output embeddings
        if tie_embeddings:
            self.reverse_embedding.weight = self.token_embedding.weight

        self._init_weights()

    def _init_weights(self):
        """Initialize weights with small random values."""
        for module in self.modules():
            if isinstance(module, nn.Linear):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
            elif isinstance(module, nn.Embedding):
                nn.init.normal_(module.weight, mean=0.0, std=0.02)
            elif isinstance(module, nn.LayerNorm):
                nn.init.ones_(module.weight)
                nn.init.zeros_(module.bias)

    def embed_tokens(self, token_ids):
        """
        Convert token IDs to embeddings with positional encoding.

        Args:
            token_ids: [batch, seq_len] integer tensor

        Returns:
            embeddings: [batch, seq_len, d_model] float tensor
        """
        batch_size, seq_len = token_ids.shape

        # Token embeddings
        token_emb = self.token_embedding(token_ids)

        # Positional embeddings
        positions = torch.arange(seq_len, device=token_ids.device)
        positions = positions.unsqueeze(0).expand(batch_size, -1)
        pos_emb = self.position_embedding(positions)

        # Combine
        embeddings = self.embedding_dropout(token_emb + pos_emb)

        return embeddings

    def apply_transformer_blocks(self, x, mask=None):
        """Apply all transformer blocks sequentially."""
        for block in self.transformer_blocks:
            x = block(x, mask)
        return x

    def forward_pass(self, x, y, z, mask=None):
        """
        Single forward pass through the transformer.

        This is the key: all three streams are concatenated,
        processed together, then split back apart.

        Args:
            x: [batch, len_x, d_model] - question stream
            y: [batch, len_y, d_model] - answer stream
            z: [batch, len_z, d_model] - reasoning stream

        Returns:
            x_new, y_new, z_new: Updated streams
        """
        len_x = x.size(1)
        len_y = y.size(1)
        len_z = z.size(1)

        # Concatenate three streams
        combined = torch.cat([x, y, z], dim=1)  # [batch, len_x+len_y+len_z, d_model]

        # Process through transformer
        combined = self.apply_transformer_blocks(combined, mask)

        # Split back
        x_new = combined[:, :len_x, :]
        y_new = combined[:, len_x:len_x+len_y, :]
        z_new = combined[:, len_x+len_y:, :]

        return x_new, y_new, z_new

    def recursive_reasoning(self, x, y, z, mask=None, return_trajectory=False):
        """
        The heart of TRM: recursive reasoning.

        Phase 1 (n_reasoning_steps): Update z (build understanding)
        Phase 2 (n_refinement_steps): Update y (improve answer)

        Args:
            x: Question stream (fixed)
            y: Answer stream (refined in phase 2)
            z: Reasoning stream (refined in phase 1)

        Returns:
            y_final: Final answer predictions
            trajectory: (optional) History of states
        """
        trajectory = {'z_states': [], 'y_states': []} if return_trajectory else None

        # Phase 1: Build reasoning (update z only)
        for step in range(self.n_reasoning_steps):
            x_new, y_new, z_new = self.forward_pass(x, y, z, mask)
            z = z_new  # Only update z

            if return_trajectory:
                trajectory['z_states'].append(z.detach().clone())

        # Phase 2: Refine answer (update y only)
        for step in range(self.n_refinement_steps):
            x_new, y_new, z_new = self.forward_pass(x, y, z, mask)
            y = y_new  # Only update y

            if return_trajectory:
                trajectory['y_states'].append(y.detach().clone())

        return (y, trajectory) if return_trajectory else y

    def forward(self, question_ids, answer_ids=None, latent_len=32, mask=None):
        """
        Complete forward pass.

        Args:
            question_ids: [batch, len_q] - input question as tokens
            answer_ids: [batch, len_a] - target answer (for training)
            latent_len: int - length of reasoning stream

        Returns:
            logits: [batch, len_a, vocab_size] - predicted answer
        """
        batch_size = question_ids.size(0)
        device = question_ids.device

        # Embed question (x stream - fixed)
        x = self.embed_tokens(question_ids)

        # Embed or initialize answer (y stream - will be refined)
        if answer_ids is not None:
            y = self.embed_tokens(answer_ids)  # Training: start from target
        else:
            len_a = 81  # Default for Sudoku
            y = torch.randn(batch_size, len_a, self.d_model, device=device) * 0.02

        # Initialize reasoning (z stream - will be refined)
        z = torch.randn(batch_size, latent_len, self.d_model, device=device) * 0.02

        # Recursive reasoning
        y_final = self.recursive_reasoning(x, y, z, mask)

        # Decode to token predictions
        logits = self.reverse_embedding(y_final)

        return logits

    def generate(self, question_ids, max_length=81, latent_len=32):
        """
        Generate answer for a given question.

        Args:
            question_ids: [batch, len_q] - input question
            max_length: int - maximum answer length
            latent_len: int - reasoning stream length

        Returns:
            generated: [batch, max_length] - predicted tokens
        """
        batch_size = question_ids.size(0)
        device = question_ids.device

        # Initialize empty answer
        y_init = torch.zeros(batch_size, max_length, dtype=torch.long, device=device)

        # Get predictions
        logits = self.forward(question_ids, y_init, latent_len)

        # Take argmax
        generated = torch.argmax(logits, dim=-1)

        return generated

    def count_parameters(self):
        """Count total trainable parameters."""
        return sum(p.numel() for p in self.parameters() if p.requires_grad)

def create_trm_att(vocab_size=12, d_model=256, n_layers=2):
    """
    Create TRM with attention (TRM-Att variant).

    Best for: Maze, ARC-AGI (large, variable-length problems)
    Parameters: ~7M
    """
    return TRM(
        vocab_size=vocab_size,
        d_model=d_model,
        n_heads=4,
        d_ff=d_model * 4,
        n_layers=n_layers,
        n_reasoning_steps=8,
        n_refinement_steps=16,
        use_attention=True
    )

def create_trm_mlp(vocab_size=12, d_model=256, n_layers=2):
    """
    Create TRM without attention (TRM-MLP variant).

    Best for: Sudoku (fixed-length, highly structured problems)
    Parameters: ~5M
    Achieves 87.4% on Sudoku-Extreme!
    """
    return TRM(
        vocab_size=vocab_size,
        d_model=d_model,
        n_heads=4,  # Not used, but kept for compatibility
        d_ff=d_model * 4,
        n_layers=n_layers,
        n_reasoning_steps=8,
        n_refinement_steps=16,
        use_attention=False  # Key difference!
    )

# Example: Create models
model_att = create_trm_att()
model_mlp = create_trm_mlp()

print(f"TRM-Att parameters: {model_att.count_parameters() / 1e6:.2f}M")
print(f"TRM-MLP parameters: {model_mlp.count_parameters() / 1e6:.2f}M")

class ExponentialMovingAverage:
    """
    EMA for model weights - improves stability on small datasets.

    Instead of using current weights directly, maintain a moving average.
    This prevents sharp updates that might break good solutions.
    """

    def __init__(self, model, decay=0.999):
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.backup = {}

        # Initialize shadow weights
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self):
        """Update shadow weights after each optimization step."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.data

    def apply_shadow(self):
        """Replace model weights with shadow weights (for inference)."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.backup[name] = param.data.clone()
                param.data = self.shadow[name]

    def restore(self):
        """Restore original weights (after inference)."""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.data = self.backup[name]

def get_lr_scheduler(optimizer, warmup_steps=1000, total_steps=100000):
    """
    Learning rate scheduler with warmup and cosine decay.

    Warmup: Gradually increase from 0 to target LR
    Cosine: Smoothly decrease to near 0
    """
    def lr_lambda(current_step):
        if current_step < warmup_steps:
            # Linear warmup
            return float(current_step) / float(max(1, warmup_steps))
        # Cosine decay
        progress = float(current_step - warmup_steps) / float(max(1, total_steps - warmup_steps))
        return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress)))

    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)

def train_with_deep_supervision(
    model,
    train_loader,
    optimizer,
    scheduler,
    ema,
    device,
    n_supervision=16,
    max_grad_norm=1.0
):
    """
    Training loop with deep supervision and all the tricks.

    Args:
        model: TRM model
        train_loader: DataLoader for training data
        optimizer: Optimizer (e.g., AdamW)
        scheduler: Learning rate scheduler
        ema: Exponential moving average
        device: torch.device
        n_supervision: Number of supervision cycles (16 in paper)
        max_grad_norm: Gradient clipping threshold

    Returns:
        average_loss: Loss for this epoch
    """
    model.train()
    total_loss = 0
    num_batches = 0

    for question_ids, answer_ids in train_loader:
        question_ids = question_ids.to(device)
        answer_ids = answer_ids.to(device)

        batch_size = question_ids.size(0)

        # Initialize streams
        x = model.embed_tokens(question_ids)
        y = model.embed_tokens(question_ids)  # Start with input
        z = torch.randn(batch_size, 32, model.d_model, device=device) * 0.02

        optimizer.zero_grad()
        cycle_loss = 0

        # Deep supervision: multiple improvement cycles
        for sup_step in range(n_supervision):
            # Phase 1: Reasoning
            for i in range(model.n_reasoning_steps):
                x_new, y_new, z_new = model.forward_pass(x, y, z)
                z = z_new

            # Phase 2: Refinement
            for i in range(model.n_refinement_steps):
                x_new, y_new, z_new = model.forward_pass(x, y, z)
                y = y_new

            # Compute loss for this cycle
            logits = model.reverse_embedding(y)
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                answer_ids.reshape(-1),
                ignore_index=0  # Ignore padding
            )

            # Weight: increase for later supervision steps
            weight = (sup_step + 1) / n_supervision
            cycle_loss += weight * loss

            # Detach to prevent backprop through all history
            y = y.detach()
            z = z.detach()

        # Backpropagate
        cycle_loss.backward()

        # Gradient clipping (crucial for stability!)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)

        # Update weights
        optimizer.step()
        scheduler.step()

        # Update EMA
        ema.update()

        total_loss += cycle_loss.item()
        num_batches += 1

    return total_loss / num_batches

def main():
    """
    Complete training script for TRM.
    """
    # Setup
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")

    # Hyperparameters
    vocab_size = 12
    d_model = 256
    n_layers = 2
    batch_size = 32
    learning_rate = 1e-4
    num_epochs = 100

    # Create model
    model = create_trm_mlp(vocab_size, d_model, n_layers)
    model = model.to(device)

    print(f"Parameters: {model.count_parameters() / 1e6:.2f}M")

    # Optimizer and scheduler
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=0.01)
    scheduler = get_lr_scheduler(optimizer, warmup_steps=1000)

    # EMA for stability
    ema = ExponentialMovingAverage(model, decay=0.999)

    # Loss function
    criterion = nn.CrossEntropyLoss(ignore_index=0)

    # Training loop
    best_val_loss = float('inf')

    for epoch in range(num_epochs):
        # Train
        train_loss = train_with_deep_supervision(
            model, train_loader, optimizer, scheduler, ema, device
        )

        # Validate with EMA weights
        ema.apply_shadow()
        val_loss = evaluate(model, val_loader, criterion, device)
        ema.restore()

        print(f"Epoch {epoch+1}/{num_epochs}")
        print(f"  Train Loss: {train_loss:.4f}")
        print(f"  Val Loss: {val_loss:.4f}")
        print(f"  LR: {optimizer.param_groups[0]['lr']:.6f}")

        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'ema_shadow': ema.shadow,
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'val_loss': val_loss,
            }, 'best_trm_model.pt')
            print("  Saved new best model!")

if __name__ == "__main__":
    main()