model.py

"""
Small transformer model for modular arithmetic experiments.
============================================================
A minimal GPT-style decoder-only transformer designed to:
1. Train from scratch in minutes on a single GPU
2. Expose all internal activations (hidden states, attention patterns)
3. Support checkpoint saving/loading for representation tracking

Architecture matches Nanda et al. 2023 (grokking) configuration
with adjustments for our two-task experiment.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict, Optional, Tuple, List
from dataclasses import dataclass


@dataclass
class TransformerConfig:
    """Configuration for the small transformer."""
    vocab_size: int = 101     # p + NUM_SPECIAL (97 + 4)
    n_layers: int = 2
    d_model: int = 128
    n_heads: int = 4
    d_mlp: int = 512
    max_seq_len: int = 5      # [a, op, b, =, c]
    dropout: float = 0.0
    layer_norm: bool = True


class MultiHeadAttention(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.n_heads = config.n_heads
        self.d_head = config.d_model // config.n_heads
        self.d_model = config.d_model

        self.W_Q = nn.Linear(config.d_model, config.d_model, bias=False)
        self.W_K = nn.Linear(config.d_model, config.d_model, bias=False)
        self.W_V = nn.Linear(config.d_model, config.d_model, bias=False)
        self.W_O = nn.Linear(config.d_model, config.d_model, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor,
                return_attn: bool = False) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
        B, T, D = x.shape

        Q = self.W_Q(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        K = self.W_K(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)
        V = self.W_V(x).view(B, T, self.n_heads, self.d_head).transpose(1, 2)

        # Scaled dot-product attention with causal mask
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_head)
        causal_mask = torch.triu(torch.ones(T, T, device=x.device), diagonal=1).bool()
        scores.masked_fill_(causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        out = (attn_weights @ V).transpose(1, 2).reshape(B, T, D)
        out = self.W_O(out)

        if return_attn:
            return out, attn_weights  # [B, H, T, T]
        return out, None


class MLP(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.W_in = nn.Linear(config.d_model, config.d_mlp)
        self.W_out = nn.Linear(config.d_mlp, config.d_model)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        hidden = F.gelu(self.W_in(x))
        out = self.dropout(self.W_out(hidden))
        return out, hidden  # return pre-projection activations for probing


class TransformerBlock(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.attn = MultiHeadAttention(config)
        self.mlp = MLP(config)
        self.ln1 = nn.LayerNorm(config.d_model) if config.layer_norm else nn.Identity()
        self.ln2 = nn.LayerNorm(config.d_model) if config.layer_norm else nn.Identity()

    def forward(self, x: torch.Tensor,
                return_internals: bool = False) -> Dict[str, torch.Tensor]:
        # Pre-norm residual architecture
        attn_out, attn_weights = self.attn(self.ln1(x), return_attn=return_internals)
        x_post_attn = x + attn_out

        mlp_out, mlp_hidden = self.mlp(self.ln2(x_post_attn))
        x_post_mlp = x_post_attn + mlp_out

        result = {'hidden_state': x_post_mlp}
        if return_internals:
            result['attn_weights'] = attn_weights
            result['mlp_hidden'] = mlp_hidden
            result['residual_post_attn'] = x_post_attn
        return result


class SmallTransformer(nn.Module):
    """
    Minimal GPT for modular arithmetic with full activation access.
    """

    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.config = config
        self.tok_embed = nn.Embedding(config.vocab_size, config.d_model)
        self.pos_embed = nn.Embedding(config.max_seq_len, config.d_model)
        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        self.ln_final = nn.LayerNorm(config.d_model) if config.layer_norm else nn.Identity()
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # Weight tying (embedding ↔ output)
        self.lm_head.weight = self.tok_embed.weight

        self.apply(self._init_weights)

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, std=0.02)

    def forward(self, input_ids: torch.Tensor,
                labels: Optional[torch.Tensor] = None,
                return_internals: bool = False) -> Dict[str, torch.Tensor]:
        B, T = input_ids.shape
        device = input_ids.device

        tok_emb = self.tok_embed(input_ids)
        pos_emb = self.pos_embed(torch.arange(T, device=device))
        x = tok_emb + pos_emb

        # Collect internals
        all_hidden_states = [x.detach()]
        all_attn_weights = []
        all_mlp_hidden = []

        for block in self.blocks:
            block_out = block(x, return_internals=return_internals)
            x = block_out['hidden_state']
            all_hidden_states.append(x.detach())
            if return_internals:
                all_attn_weights.append(block_out['attn_weights'].detach())
                all_mlp_hidden.append(block_out['mlp_hidden'].detach())

        x = self.ln_final(x)
        logits = self.lm_head(x)

        result = {'logits': logits}

        if labels is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   labels.view(-1), ignore_index=-100)
            result['loss'] = loss

        if return_internals:
            result['hidden_states'] = all_hidden_states  # List of [B, T, D]
            result['attn_weights'] = all_attn_weights    # List of [B, H, T, T]
            result['mlp_hidden'] = all_mlp_hidden        # List of [B, T, D_mlp]

        return result

    def get_representations(self, input_ids: torch.Tensor,
                            token_position: int = -1) -> List[torch.Tensor]:
        """
        Get hidden state at each layer for a specific token position.
        Returns list of [batch_size, d_model] tensors.
        """
        with torch.no_grad():
            out = self.forward(input_ids, return_internals=True)
        return [hs[:, token_position, :] for hs in out['hidden_states']]

    def count_parameters(self) -> int:
        return sum(p.numel() for p in self.parameters())