Refactor LLM integration: consolidate training scripts, add extraction strategies

TRAINING SCRIPT CONSOLIDATION:
- Merge train_passthrough.py, train_passthrough_router.py, train_llm.py into unified train.py
- Remove redundant 5-mode system (router, interface, embeddings, llm, enhanced)
- New 3-mode system: router (sanity), interface (sanity), llm (real training)
- Delete orphaned trained_passthrough_router.pt checkpoint
- Move outputs to trained/ subfolder (router.pt, interface.pt, llm.pt)

MODEL ARCHITECTURE UPDATES (model.py):
- Add ArithmeticModel: unified LLM + extractor + frozen circuits
- Add Extractor: attention pooling + per-bit extraction networks
- Add PositionExtractor: position-specific extraction from token positions
- Add DigitExtractor: predict digits (0-9) then convert to bits
- Add AttentionPooling: learnable CLS token attention over sequence
- Add MultiHeadBitExtractor: 8 separate networks for 8 bits
- Add HiddenStateExtractor: simple MLP-based bit extraction
- Remove EmbeddingArithmeticModel (mean pooling failed, ~33% accuracy plateau)
- Remove AugmentedArithmeticModel (merged into ArithmeticModel)

NEW TRAINING FEATURES (baked-in):
- Curriculum learning: 0-9 (epochs 0-20%) -> 0-99 (20-50%) -> 0-255 (50-100%)
- Loss reweighting: 2x multiplier for a/b bit losses (extraction is bottleneck)
- Per-batch progress reporting every 5 batches
- Per-epoch VRAM and timing stats

NEW CLI ARGUMENTS (--mode llm):
- --unfreeze_layers N: fine-tune top N transformer layers (default 0 = frozen)
- --extract_layer N: extract from layer N (-1 = last, try 12 for middle)
- --position_extract: use position-specific extraction instead of pooling
- --digit_pred: predict digits instead of bits (aligns with tokenization)

RATIONALE:
- Embeddings mode removed: mean pooling loses positional info, can't distinguish "47" from "74"
- Operation classification works (97-100%), bit extraction is the bottleneck (~33% accuracy)
- Position-specific and digit-level extraction may better align with LLM representations
- Curriculum learning helps model learn simpler cases before harder ones

USAGE:
python train.py --mode llm --epochs 100 # baseline
python train.py --mode llm --position_extract # position-specific
python train.py --mode llm --digit_pred # digit prediction
python train.py --mode llm --extract_layer 12 # middle layer
python train.py --mode llm --unfreeze_layers 4 # fine-tune LLM

llm_integration/model.py

@@ -1,6 +1,7 @@
 """
 Trainable interface layers for frozen threshold circuits.
 BitEncoder, OpRouter, BitDecoder wrap the frozen circuits.
+HiddenStateExtractor and AugmentedArithmeticModel for LLM integration.
 """
 
 import torch
@@ -8,6 +9,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 from circuits import FrozenThresholdCircuits, heaviside_ste
 
+MODEL_ID = 'HuggingFaceTB/SmolLM2-360M-Instruct'
+OPERATIONS = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
+OP_SYMBOLS = {'add': '+', 'sub': '-', 'mul': '*', 'gt': '>', 'lt': '<', 'eq': '=='}
+
 
 class BitEncoder(nn.Module):
     """
@@ -50,12 +55,6 @@ class OpRouter(nn.Module):
     """
 
     def __init__(self, input_dim: int = 16 + 6, hidden_dim: int = 32, n_ops: int = 6):
-        """
-        Args:
-            input_dim: Input dimension
-            hidden_dim: Hidden layer dimension
-            n_ops: Number of operations to route between
-        """
         super().__init__()
         self.net = nn.Sequential(
             nn.Linear(input_dim, hidden_dim),
@@ -83,21 +82,10 @@ class BitDecoder(nn.Module):
     """
 
     def __init__(self, output_dim: int = 8):
-        """
-        Args:
-            output_dim: Output dimension (8 bits for result)
-        """
         super().__init__()
         self.output_dim = output_dim
 
     def forward(self, result_bits: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            result_bits: [batch, 8] result bits from circuits
-
-        Returns:
-            output: [batch, 8] processed output
-        """
         return result_bits
 
 
@@ -149,15 +137,6 @@ class ThresholdALU(nn.Module):
                        op_onehot: torch.Tensor) -> torch.Tensor:
         """
         Direct forward through circuits (bypass encoder/router for testing).
-        Uses ground truth bits and operation directly.
-
-        Args:
-            a_bits: [batch, 8] operand A bits
-            b_bits: [batch, 8] operand B bits
-            op_onehot: [batch, 6] one-hot operation
-
-        Returns:
-            result_bits: [batch, 8] output bits
         """
         return self.circuits(a_bits, b_bits, op_onehot)
 
@@ -175,10 +154,496 @@ class DirectCircuitModel(nn.Module):
 
     def forward(self, a_bits: torch.Tensor, b_bits: torch.Tensor,
                 op_onehot: torch.Tensor) -> torch.Tensor:
-        """Direct circuit execution."""
         return self.circuits(a_bits, b_bits, op_onehot)
 
 
+class HiddenStateExtractor(nn.Module):
+    """
+    Extracts operands and operation from LLM hidden states.
+    This is the hard part - must learn to parse numbers from embeddings.
+    """
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256):
+        super().__init__()
+
+        self.a_extractor = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 8),
+        )
+
+        self.b_extractor = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 8),
+        )
+
+        self.op_router = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, len(OPERATIONS)),
+        )
+
+    def forward(self, hidden_states: torch.Tensor):
+        """
+        Args:
+            hidden_states: [batch, hidden_dim] from LLM
+
+        Returns:
+            a_bits: [batch, 8]
+            b_bits: [batch, 8]
+            op_logits: [batch, 6]
+        """
+        a_logits = self.a_extractor(hidden_states)
+        b_logits = self.b_extractor(hidden_states)
+        op_logits = self.op_router(hidden_states)
+
+        a_soft = torch.sigmoid(a_logits)
+        b_soft = torch.sigmoid(b_logits)
+
+        a_hard = heaviside_ste(a_logits)
+        b_hard = heaviside_ste(b_logits)
+
+        a_bits = a_hard - a_soft.detach() + a_soft
+        b_bits = b_hard - b_soft.detach() + b_soft
+
+        return a_bits, b_bits, op_logits
+
+
+class AttentionPooling(nn.Module):
+    """
+    Learnable attention pooling over sequence positions.
+    Replaces mean pooling - learns which tokens matter for extraction.
+    """
+
+    def __init__(self, hidden_dim: int = 960, num_heads: int = 4):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+
+        self.query = nn.Linear(hidden_dim, hidden_dim)
+        self.key = nn.Linear(hidden_dim, hidden_dim)
+        self.value = nn.Linear(hidden_dim, hidden_dim)
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
+
+        self.cls_token = nn.Parameter(torch.randn(1, 1, hidden_dim) * 0.02)
+
+    def forward(self, embeddings: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            embeddings: [batch, seq_len, hidden_dim]
+            mask: [batch, seq_len] attention mask (1 = attend, 0 = ignore)
+
+        Returns:
+            pooled: [batch, hidden_dim]
+        """
+        batch_size, seq_len, hidden_dim = embeddings.shape
+
+        cls_expanded = self.cls_token.expand(batch_size, -1, -1)
+        embeddings = torch.cat([cls_expanded, embeddings], dim=1)
+
+        cls_mask = torch.ones(batch_size, 1, device=mask.device)
+        mask = torch.cat([cls_mask, mask], dim=1)
+
+        Q = self.query(embeddings[:, :1, :])
+        K = self.key(embeddings)
+        V = self.value(embeddings)
+
+        Q = Q.view(batch_size, 1, self.num_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, seq_len + 1, self.num_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, seq_len + 1, self.num_heads, self.head_dim).transpose(1, 2)
+
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)
+
+        mask_expanded = mask.unsqueeze(1).unsqueeze(2)
+        scores = scores.masked_fill(mask_expanded == 0, -1e9)
+
+        attn_weights = torch.softmax(scores, dim=-1)
+        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+
+        context = torch.matmul(attn_weights, V)
+        context = context.transpose(1, 2).contiguous().view(batch_size, 1, hidden_dim)
+
+        pooled = self.out_proj(context).squeeze(1)
+        pooled = torch.nan_to_num(pooled, nan=0.0)
+
+        return pooled
+
+
+class MultiHeadBitExtractor(nn.Module):
+    """
+    8 separate extractors for 8 bits - each bit gets its own specialized network.
+    More expressive than single MLP predicting all 8 bits at once.
+    """
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 128):
+        super().__init__()
+
+        self.bit_extractors = nn.ModuleList([
+            nn.Sequential(
+                nn.Linear(hidden_dim, intermediate_dim),
+                nn.GELU(),
+                nn.Linear(intermediate_dim, 1),
+            )
+            for _ in range(8)
+        ])
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [batch, hidden_dim]
+
+        Returns:
+            bits: [batch, 8] - one bit from each extractor
+        """
+        hidden_states = torch.nan_to_num(hidden_states, nan=0.0)
+
+        bit_logits = [extractor(hidden_states) for extractor in self.bit_extractors]
+        logits = torch.cat(bit_logits, dim=-1)
+        logits = torch.clamp(logits, -20, 20)
+
+        soft = torch.sigmoid(logits)
+        hard = heaviside_ste(logits)
+        bits = hard - soft.detach() + soft
+
+        return bits, logits
+
+
+class Extractor(nn.Module):
+    """
+    Extracts operands and operation from LLM hidden states.
+    Uses attention pooling and per-bit extraction networks.
+    """
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
+        super().__init__()
+
+        self.attention_pool = AttentionPooling(hidden_dim, num_heads)
+
+        self.a_extractor = MultiHeadBitExtractor(hidden_dim, intermediate_dim // 2)
+        self.b_extractor = MultiHeadBitExtractor(hidden_dim, intermediate_dim // 2)
+
+        self.op_router = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, len(OPERATIONS)),
+        )
+
+    def forward(self, embeddings: torch.Tensor, mask: torch.Tensor):
+        """
+        Args:
+            embeddings: [batch, seq_len, hidden_dim]
+            mask: [batch, seq_len]
+
+        Returns:
+            a_bits: [batch, 8]
+            b_bits: [batch, 8]
+            op_logits: [batch, 6]
+        """
+        pooled = self.attention_pool(embeddings, mask)
+
+        a_bits, _ = self.a_extractor(pooled)
+        b_bits, _ = self.b_extractor(pooled)
+        op_logits = self.op_router(pooled)
+
+        return a_bits, b_bits, op_logits
+
+
+class PositionExtractor(nn.Module):
+    """
+    Position-specific extraction.
+    Extracts operand A from first token positions, operand B from later positions.
+    For "47 + 86": positions 0-2 for A, position 3-4 for op, positions 5-7 for B.
+    """
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256):
+        super().__init__()
+
+        self.a_extractor = nn.Sequential(
+            nn.Linear(hidden_dim * 3, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 8),
+        )
+
+        self.b_extractor = nn.Sequential(
+            nn.Linear(hidden_dim * 3, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 8),
+        )
+
+        self.op_router = nn.Sequential(
+            nn.Linear(hidden_dim * 2, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, len(OPERATIONS)),
+        )
+
+    def forward(self, hidden: torch.Tensor, mask: torch.Tensor):
+        """
+        Args:
+            hidden: [batch, seq_len, hidden_dim]
+            mask: [batch, seq_len]
+
+        Returns:
+            a_bits, b_bits, op_logits
+        """
+        batch_size, seq_len, hidden_dim = hidden.shape
+
+        seq_lens = mask.sum(dim=1).long()
+
+        a_features = []
+        b_features = []
+        op_features = []
+
+        for i in range(batch_size):
+            slen = seq_lens[i].item()
+            start = seq_len - slen
+
+            a_pos = hidden[i, start:start+3, :].reshape(-1)
+            if a_pos.shape[0] < hidden_dim * 3:
+                a_pos = F.pad(a_pos, (0, hidden_dim * 3 - a_pos.shape[0]))
+
+            op_pos = hidden[i, start+3:start+5, :].reshape(-1)
+            if op_pos.shape[0] < hidden_dim * 2:
+                op_pos = F.pad(op_pos, (0, hidden_dim * 2 - op_pos.shape[0]))
+
+            b_pos = hidden[i, start+5:start+8, :].reshape(-1)
+            if b_pos.shape[0] < hidden_dim * 3:
+                b_pos = F.pad(b_pos, (0, hidden_dim * 3 - b_pos.shape[0]))
+
+            a_features.append(a_pos)
+            b_features.append(b_pos)
+            op_features.append(op_pos)
+
+        a_features = torch.stack(a_features)
+        b_features = torch.stack(b_features)
+        op_features = torch.stack(op_features)
+
+        a_logits = self.a_extractor(a_features)
+        b_logits = self.b_extractor(b_features)
+        op_logits = self.op_router(op_features)
+
+        a_soft = torch.sigmoid(a_logits)
+        b_soft = torch.sigmoid(b_logits)
+        a_hard = heaviside_ste(a_logits)
+        b_hard = heaviside_ste(b_logits)
+        a_bits = a_hard - a_soft.detach() + a_soft
+        b_bits = b_hard - b_soft.detach() + b_soft
+
+        return a_bits, b_bits, op_logits
+
+
+class DigitExtractor(nn.Module):
+    """
+    Digit-level extraction: predicts digits (0-9) then converts to bits.
+    More aligned with tokenization.
+    """
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
+        super().__init__()
+
+        self.attention_pool = AttentionPooling(hidden_dim, num_heads)
+
+        self.a_digit_pred = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 3 * 10),
+        )
+
+        self.b_digit_pred = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 3 * 10),
+        )
+
+        self.op_router = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, len(OPERATIONS)),
+        )
+
+    def digits_to_bits(self, digit_logits: torch.Tensor) -> torch.Tensor:
+        """
+        Convert 3-digit predictions to 8-bit representation.
+        digit_logits: [batch, 30] (3 digits * 10 classes each)
+        Returns: [batch, 8] bits
+        """
+        batch_size = digit_logits.shape[0]
+
+        logits = digit_logits.view(batch_size, 3, 10)
+        probs = torch.softmax(logits, dim=-1)
+
+        digit_values = torch.arange(10, device=digit_logits.device).float()
+        soft_digits = (probs * digit_values).sum(dim=-1)
+
+        hundreds = soft_digits[:, 0]
+        tens = soft_digits[:, 1]
+        ones = soft_digits[:, 2]
+
+        value = hundreds * 100 + tens * 10 + ones
+        value = torch.clamp(value, 0, 255)
+
+        bits = []
+        for i in range(7, -1, -1):
+            bit = torch.fmod(torch.floor(value / (2 ** i)), 2)
+            bits.append(bit)
+
+        return torch.stack(bits, dim=-1)
+
+    def forward(self, hidden: torch.Tensor, mask: torch.Tensor):
+        """
+        Returns:
+            a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
+        """
+        pooled = self.attention_pool(hidden, mask)
+
+        a_digit_logits = self.a_digit_pred(pooled)
+        b_digit_logits = self.b_digit_pred(pooled)
+        op_logits = self.op_router(pooled)
+
+        a_bits = self.digits_to_bits(a_digit_logits)
+        b_bits = self.digits_to_bits(b_digit_logits)
+
+        return a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
+
+
+class ArithmeticModel(nn.Module):
+    """
+    LLM + extractor + frozen threshold circuits.
+    Optionally unfreeze top N transformer layers with --unfreeze_layers.
+    """
+
+    def __init__(self, device: str = 'cuda', unfreeze_layers: int = 0,
+                 extract_layer: int = -1, position_extract: bool = False,
+                 digit_pred: bool = False):
+        super().__init__()
+        self.device = device
+        self.unfreeze_layers = unfreeze_layers
+        self.extract_layer = extract_layer
+        self.position_extract = position_extract
+        self.digit_pred = digit_pred
+
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        print("[1/4] Loading tokenizer...", flush=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        self.tokenizer.padding_side = 'left'
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        print("  Tokenizer loaded.", flush=True)
+
+        print("[2/4] Loading SmolLM2-360M...", flush=True)
+        self.llm = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.float16,
+            device_map=device,
+            output_hidden_states=True
+        )
+
+        for param in self.llm.parameters():
+            param.requires_grad = False
+
+        if unfreeze_layers > 0:
+            num_layers = len(self.llm.model.layers)
+            layers_to_unfreeze = list(range(num_layers - unfreeze_layers, num_layers))
+            print(f"  Unfreezing layers {layers_to_unfreeze}...", flush=True)
+            for layer_idx in layers_to_unfreeze:
+                for param in self.llm.model.layers[layer_idx].parameters():
+                    param.requires_grad = True
+
+        hidden_dim = self.llm.config.hidden_size
+        llm_params = sum(p.numel() for p in self.llm.parameters())
+        trainable_llm = sum(p.numel() for p in self.llm.parameters() if p.requires_grad)
+        print(f"  LLM loaded. Hidden dim: {hidden_dim}", flush=True)
+        print(f"  LLM params: {llm_params:,} total, {trainable_llm:,} trainable", flush=True)
+
+        print("[3/4] Loading threshold circuits...", flush=True)
+        self.circuits = FrozenThresholdCircuits(device=device)
+        print(f"  Circuits loaded. {len(self.circuits.weights)} tensors", flush=True)
+
+        print("[4/4] Initializing extractor...", flush=True)
+        if position_extract:
+            print("  Using position-specific extraction", flush=True)
+            self.extractor = PositionExtractor(
+                hidden_dim=hidden_dim,
+                intermediate_dim=256
+            ).to(device)
+        elif digit_pred:
+            print("  Using digit-level prediction", flush=True)
+            self.extractor = DigitExtractor(
+                hidden_dim=hidden_dim,
+                intermediate_dim=256,
+                num_heads=4
+            ).to(device)
+        else:
+            self.extractor = Extractor(
+                hidden_dim=hidden_dim,
+                intermediate_dim=256,
+                num_heads=4
+            ).to(device)
+
+        if extract_layer != -1:
+            print(f"  Extracting from layer {extract_layer}", flush=True)
+
+        trainable_ext = sum(p.numel() for p in self.extractor.parameters())
+        total_trainable = trainable_llm + trainable_ext
+        print(f"  Extractor params: {trainable_ext:,}", flush=True)
+        print(f"  Total trainable: {total_trainable:,}", flush=True)
+
+    def get_hidden_states(self, texts: list[str]) -> tuple[torch.Tensor, torch.Tensor]:
+        """Get hidden states from specified layer."""
+        inputs = self.tokenizer(
+            texts,
+            return_tensors='pt',
+            padding=True,
+            truncation=True,
+            max_length=64
+        ).to(self.device)
+
+        if self.unfreeze_layers > 0:
+            outputs = self.llm(**inputs, output_hidden_states=True)
+        else:
+            with torch.no_grad():
+                outputs = self.llm(**inputs, output_hidden_states=True)
+
+        hidden = outputs.hidden_states[self.extract_layer].float()
+        mask = inputs.attention_mask.float()
+
+        return hidden, mask
+
+    def forward(self, texts: list[str]):
+        """
+        Full forward pass: text -> hidden states -> extractor -> circuits -> result
+
+        Returns:
+            result_bits, a_bits, b_bits, op_logits
+            If digit_pred: also returns a_digit_logits, b_digit_logits
+        """
+        hidden, mask = self.get_hidden_states(texts)
+
+        extractor_out = self.extractor(hidden, mask)
+
+        if self.digit_pred:
+            a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits = extractor_out
+        else:
+            a_bits, b_bits, op_logits = extractor_out
+            a_digit_logits, b_digit_logits = None, None
+
+        op_probs = torch.softmax(op_logits, dim=-1)
+
+        result_bits = self.circuits(a_bits, b_bits, op_probs)
+
+        if self.digit_pred:
+            return result_bits, a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
+        return result_bits, a_bits, b_bits, op_logits
+
+    def trainable_parameters(self):
+        """Return all trainable parameters for optimizer."""
+        params = list(self.extractor.parameters())
+        if self.unfreeze_layers > 0:
+            params += [p for p in self.llm.parameters() if p.requires_grad]
+        return params
+
+
 if __name__ == "__main__":
     import sys
     sys.path.insert(0, '.')

llm_integration/train.py

@@ -0,0 +1,673 @@
+"""
+Unified training script for threshold circuit LLM integration.
+
+Modes:
+  --mode router    : Train only OpRouter with ground truth bits (sanity check)
+  --mode interface : Train BitEncoder + OpRouter with ground truth bits (sanity check)
+  --mode llm       : Train extractor with LLM hidden states (the real training)
+
+LLM mode options:
+  --unfreeze_layers N : Unfreeze top N transformer layers (default 0 = fully frozen)
+
+Hardware Profile (NVIDIA RTX 6000 Ada 48GB):
+  VRAM Scaling (unfreeze_layers=4):
+    batch_size |  VRAM   |  %
+    -----------+---------+------
+         512   |  5,784  | 11.8%
+       1,024   |  7,384  | 15.0%
+       4,096   | 16,534  | 33.6%
+      13,000   | 39,000  | 79.4%  <-- recommended for 80% target
+
+Examples:
+  python train.py --mode llm --epochs 100 --batch_size 256
+  python train.py --mode llm --epochs 100 --batch_size 4096 --unfreeze_layers 4
+"""
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import time
+import argparse
+import random
+
+from model import (
+    ThresholdALU, DirectCircuitModel, OpRouter,
+    ArithmeticModel, OPERATIONS, OP_SYMBOLS
+)
+from circuits import FrozenThresholdCircuits
+from fitness import generate_batch, compute_fitness, compute_loss
+
+DEVICE = 'cuda'
+
+
+def int_to_bits(val: int, device: str = 'cuda') -> torch.Tensor:
+    bits = torch.zeros(8, device=device)
+    for i in range(8):
+        bits[7-i] = (val >> i) & 1
+    return bits
+
+
+def bits_to_int(bits: torch.Tensor) -> int:
+    val = 0
+    for i in range(8):
+        if bits[i].item() > 0.5:
+            val += 1 << (7-i)
+    return val
+
+
+def generate_problem(max_val: int = 255):
+    """Generate a random arithmetic problem for LLM training."""
+    a = random.randint(0, max_val)
+    b = random.randint(0, max_val)
+    op = random.choice(OPERATIONS)
+
+    sym = OP_SYMBOLS[op]
+    text = f"{a} {sym} {b}"
+
+    if op == 'add':
+        result = (a + b) & 0xFF
+    elif op == 'sub':
+        result = (a - b) & 0xFF
+    elif op == 'mul':
+        result = (a * b) & 0xFF
+    elif op == 'gt':
+        result = 1 if a > b else 0
+    elif op == 'lt':
+        result = 1 if a < b else 0
+    elif op == 'eq':
+        result = 1 if a == b else 0
+
+    return text, a, b, op, result
+
+
+def get_curriculum_max(epoch: int, total_epochs: int) -> int:
+    """
+    Curriculum learning: start with small numbers, gradually increase.
+    Epoch 0-20%: 0-9 (single digit)
+    Epoch 20-50%: 0-99 (two digit)
+    Epoch 50-100%: 0-255 (full range)
+    """
+    progress = epoch / total_epochs
+    if progress < 0.2:
+        return 9
+    elif progress < 0.5:
+        return 99
+    else:
+        return 255
+
+
+def train_router(epochs: int = 100, batch_size: int = 256, lr: float = 1e-2, device: str = 'cuda'):
+    """Train only the router with ground truth bits."""
+    print("=" * 70)
+    print(" ROUTER-ONLY TRAINING (Ground Truth Bits)")
+    print("=" * 70)
+
+    circuits = FrozenThresholdCircuits(device=device)
+    router = OpRouter(input_dim=16 + 6, hidden_dim=64, n_ops=6).to(device)
+
+    print(f"\nRouter parameters: {sum(p.numel() for p in router.parameters()):,}")
+
+    def model_fn(a_bits, b_bits, op_onehot):
+        x = torch.cat([a_bits, b_bits, op_onehot], dim=-1)
+        op_weights = router(x)
+        return circuits(a_bits, b_bits, op_weights)
+
+    initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
+    print(f"Initial fitness: {initial_fitness:.4f}")
+
+    optimizer = optim.AdamW(router.parameters(), lr=lr)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    print("\nTraining...")
+    print("-" * 70)
+
+    best_fitness = initial_fitness
+    start_time = time.perf_counter()
+
+    for epoch in range(epochs):
+        router.train()
+        epoch_loss = 0.0
+
+        for _ in range(100):
+            batch = generate_batch(batch_size, device)
+
+            optimizer.zero_grad()
+
+            x = torch.cat([batch['a_bits'], batch['b_bits'], batch['op_onehot']], dim=-1)
+            op_weights = router(x)
+            pred_bits = circuits(batch['a_bits'], batch['b_bits'], op_weights)
+
+            loss = compute_loss(pred_bits, batch['expected_bits'])
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+        scheduler.step()
+
+        if (epoch + 1) % 10 == 0 or epoch == 0:
+            router.eval()
+            fitness, details = compute_fitness(model_fn, n_samples=2000, device=device, return_details=True)
+            elapsed = time.perf_counter() - start_time
+
+            if fitness > best_fitness:
+                best_fitness = fitness
+                marker = " *"
+            else:
+                marker = ""
+
+            print(f"Epoch {epoch+1:3d} | Loss: {epoch_loss/100:.4f} | "
+                  f"Fitness: {fitness:.4f}{marker} | Time: {elapsed:.1f}s")
+
+            if fitness >= 0.9999:
+                print("\n TARGET: 100% FITNESS ACHIEVED")
+                break
+
+    print("\n" + "=" * 70)
+    print(" RESULTS")
+    print("=" * 70)
+
+    router.eval()
+    final_fitness, details = compute_fitness(model_fn, n_samples=5000, device=device, return_details=True)
+
+    print(f"\nFinal fitness: {final_fitness:.4f}")
+    print(f"\nPer-operation:")
+    for op in OPERATIONS:
+        acc = details['by_op'][op]['accuracy']
+        print(f"  {op}: {acc:.4f}")
+
+    print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
+
+    if final_fitness >= 0.99:
+        print("\nCONCLUSION: Router successfully learned operation dispatch.")
+        print("           With correct bit encoding, 100% is achievable.")
+
+    save_path = "D:/8bit-threshold-computer/llm_integration/trained/router.pt"
+    torch.save({
+        'router_state_dict': router.state_dict(),
+        'final_fitness': final_fitness,
+        'params': sum(p.numel() for p in router.parameters()),
+    }, save_path)
+    print(f"\nSaved trained router to: {save_path}")
+
+    return router, final_fitness
+
+
+def get_gpu_memory():
+    """Get GPU memory usage in MB."""
+    if torch.cuda.is_available():
+        return torch.cuda.memory_allocated() / 1024 / 1024, torch.cuda.max_memory_allocated() / 1024 / 1024
+    return 0, 0
+
+
+def train_interface(epochs: int = 200, batch_size: int = 512, lr: float = 1e-3,
+                    eval_interval: int = 10, device: str = 'cuda'):
+    """Train BitEncoder + OpRouter with ground truth bits."""
+    print("=" * 70)
+    print(" INTERFACE TRAINING (Encoder + Router)")
+    print("=" * 70)
+    print(f"  Started at: {time.strftime('%H:%M:%S')}")
+
+    print("\n[1/4] Verifying frozen circuits...")
+    print("  Creating DirectCircuitModel...", end=" ", flush=True)
+    direct_model = DirectCircuitModel(device=device)
+    mem, max_mem = get_gpu_memory()
+    print(f"done. VRAM: {mem:.0f}MB")
+
+    def direct_fn(a, b, op):
+        return direct_model(a, b, op)
+
+    print("  Running fitness check (1000 samples)...", end=" ", flush=True)
+    circuit_fitness = compute_fitness(direct_fn, n_samples=1000, device=device)
+    print(f"done. Fitness: {circuit_fitness:.4f}")
+    if circuit_fitness < 0.999:
+        print("  ERROR: Circuits not achieving 100%. Aborting.")
+        return None, 0.0
+    print("  STATUS: PASS")
+
+    print("\n[2/4] Initializing model...")
+    print("  Creating ThresholdALU...", end=" ", flush=True)
+    model = ThresholdALU(device=device)
+    mem, max_mem = get_gpu_memory()
+    print(f"done. VRAM: {mem:.0f}MB")
+
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"  Trainable parameters: {trainable_params:,}")
+
+    def model_fn(a, b, op):
+        return model(a, b, op)
+
+    print("  Running initial fitness check...", end=" ", flush=True)
+    initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
+    print(f"done. Fitness: {initial_fitness:.4f}")
+
+    print("\n[3/4] Setting up training...")
+    print("  Creating optimizer...", end=" ", flush=True)
+    optimizer = optim.AdamW(model.parameters(), lr=lr)
+    print("done.")
+    print("  Creating scheduler...", end=" ", flush=True)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    print("done.")
+
+    print(f"  Config: lr={lr}, batch_size={batch_size}, epochs={epochs}")
+
+    print("\n[4/4] Training...")
+    print("  Generating first batch to warm up...", end=" ", flush=True)
+    warmup_batch = generate_batch(batch_size, device)
+    mem, max_mem = get_gpu_memory()
+    print(f"done. VRAM: {mem:.0f}MB (max: {max_mem:.0f}MB)")
+
+    print("-" * 70)
+
+    best_fitness = initial_fitness
+    start_time = time.perf_counter()
+    n_batches = 100
+
+    for epoch in range(epochs):
+        model.train()
+        epoch_loss = 0.0
+        epoch_start = time.perf_counter()
+
+        for batch_idx in range(n_batches):
+            batch = generate_batch(batch_size, device)
+
+            optimizer.zero_grad()
+
+            pred_bits = model(batch['a_bits'], batch['b_bits'], batch['op_onehot'])
+
+            loss = compute_loss(pred_bits, batch['expected_bits'])
+
+            loss.backward()
+            optimizer.step()
+
+            epoch_loss += loss.item()
+
+            if batch_idx == 0 and epoch == 0:
+                mem, max_mem = get_gpu_memory()
+                print(f"  First forward/backward done. VRAM: {mem:.0f}MB (max: {max_mem:.0f}MB)")
+
+            if (batch_idx + 1) % 25 == 0:
+                avg_so_far = epoch_loss / (batch_idx + 1)
+                print(f"  Epoch {epoch+1} batch {batch_idx+1}/{n_batches} | loss: {avg_so_far:.4f}", flush=True)
+
+        scheduler.step()
+
+        avg_loss = epoch_loss / n_batches
+        epoch_time = time.perf_counter() - epoch_start
+
+        if (epoch + 1) % 5 == 0 or epoch == 0:  # Eval every 5 epochs
+            model.eval()
+            fitness, details = compute_fitness(
+                model_fn, n_samples=2000, device=device, return_details=True
+            )
+
+            elapsed = time.perf_counter() - start_time
+
+            if fitness > best_fitness:
+                best_fitness = fitness
+                marker = " *"
+            else:
+                marker = ""
+
+            mem, _ = get_gpu_memory()
+            print(f"Epoch {epoch+1:4d} | Loss: {avg_loss:.4f} | "
+                  f"Fitness: {fitness:.4f}{marker} | "
+                  f"LR: {scheduler.get_last_lr()[0]:.2e} | "
+                  f"VRAM: {mem:.0f}MB | "
+                  f"Time: {elapsed:.1f}s ({epoch_time:.1f}s/epoch)")
+
+            if fitness >= 0.9999:
+                print("\n" + "=" * 70)
+                print(" TARGET ACHIEVED: 100% FITNESS")
+                print("=" * 70)
+                break
+
+    print("\n" + "=" * 70)
+    print(" TRAINING COMPLETE")
+    print("=" * 70)
+
+    model.eval()
+    final_fitness, details = compute_fitness(
+        model_fn, n_samples=5000, device=device, return_details=True
+    )
+
+    print(f"\nFinal fitness: {final_fitness:.4f}")
+    print(f"Best fitness:  {best_fitness:.4f}")
+    print(f"\nPer-operation breakdown:")
+    for op in OPERATIONS:
+        acc = details['by_op'][op]['accuracy']
+        print(f"  {op:6}: {acc:.4f}")
+
+    print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
+
+    save_path = "D:/8bit-threshold-computer/llm_integration/trained/interface.pt"
+    torch.save({
+        'encoder_state_dict': model.encoder.state_dict(),
+        'router_state_dict': model.router.state_dict(),
+        'final_fitness': final_fitness,
+        'best_fitness': best_fitness,
+    }, save_path)
+    print(f"\nSaved trained model to: {save_path}")
+
+    return model, final_fitness
+
+
+def compute_llm_loss(pred_bits, a_bits, b_bits, op_logits,
+                     target_result, target_a, target_b, target_op_idx,
+                     bit_weight: float = 2.0):
+    """
+    Multi-component loss for LLM training.
+    bit_weight: multiplier for a/b bit losses (default 2x since extraction is the bottleneck)
+    """
+    result_loss = nn.functional.binary_cross_entropy_with_logits(
+        pred_bits, target_result
+    )
+
+    a_bits_safe = torch.clamp(a_bits, 0.0, 1.0)
+    b_bits_safe = torch.clamp(b_bits, 0.0, 1.0)
+    a_bits_safe = torch.nan_to_num(a_bits_safe, nan=0.5, posinf=1.0, neginf=0.0)
+    b_bits_safe = torch.nan_to_num(b_bits_safe, nan=0.5, posinf=1.0, neginf=0.0)
+
+    a_loss = nn.functional.binary_cross_entropy(
+        torch.clamp(a_bits_safe, 1e-6, 1-1e-6), target_a
+    )
+    b_loss = nn.functional.binary_cross_entropy(
+        torch.clamp(b_bits_safe, 1e-6, 1-1e-6), target_b
+    )
+
+    op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
+
+    total = result_loss + bit_weight * a_loss + bit_weight * b_loss + op_loss
+    total = torch.nan_to_num(total, nan=10.0, posinf=10.0, neginf=0.0)
+
+    return total, {
+        'result': result_loss.item() if not torch.isnan(result_loss) else 10.0,
+        'a': a_loss.item() if not torch.isnan(a_loss) else 10.0,
+        'b': b_loss.item() if not torch.isnan(b_loss) else 10.0,
+        'op': op_loss.item() if not torch.isnan(op_loss) else 10.0
+    }
+
+
+def evaluate_llm(model, n_samples: int = 500):
+    """Evaluate LLM model on random problems."""
+    model.extractor.eval()
+    correct = 0
+    op_correct = 0
+
+    for _ in range(n_samples):
+        text, a, b, op, expected = generate_problem()
+
+        with torch.no_grad():
+            result_bits, a_bits, b_bits, op_logits = model([text])
+
+        pred_result = bits_to_int(result_bits[0])
+        pred_op = OPERATIONS[op_logits[0].argmax().item()]
+
+        if pred_result == expected:
+            correct += 1
+        if pred_op == op:
+            op_correct += 1
+
+    model.extractor.train()
+    return correct / n_samples, op_correct / n_samples
+
+
+def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
+              unfreeze_layers: int = 0, extract_layer: int = -1,
+              position_extract: bool = False, digit_pred: bool = False,
+              device: str = 'cuda'):
+    """
+    Train extractor with LLM hidden states.
+
+    Args:
+        unfreeze_layers: Number of top transformer layers to unfreeze (0 = fully frozen)
+        extract_layer: Which layer to extract from (-1 = last)
+        position_extract: Use position-specific extraction
+        digit_pred: Predict digits instead of bits
+    """
+    print("=" * 70)
+    print(" LLM TRAINING")
+    if unfreeze_layers > 0:
+        print(f" {unfreeze_layers} transformer layers unfrozen")
+    else:
+        print(" LLM frozen")
+    if extract_layer != -1:
+        print(f" Extracting from layer {extract_layer}")
+    if position_extract:
+        print(" Position-specific extraction")
+    if digit_pred:
+        print(" Digit-level prediction")
+    print("=" * 70)
+
+    print("\nInitializing model...")
+    model = ArithmeticModel(
+        device=device,
+        unfreeze_layers=unfreeze_layers,
+        extract_layer=extract_layer,
+        position_extract=position_extract,
+        digit_pred=digit_pred
+    )
+
+    optimizer = optim.AdamW(model.trainable_parameters(), lr=lr)
+    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    print(f"\nTraining config:")
+    print(f"  Epochs: {epochs}")
+    print(f"  Batch size: {batch_size}")
+    print(f"  Learning rate: {lr}")
+    print(f"  Unfreeze layers: {unfreeze_layers}")
+    print(f"  Samples/epoch: {batch_size * 20}")
+
+    print(f"\nInitial evaluation (200 samples)...")
+    acc, op_acc = evaluate_llm(model, n_samples=200)
+    print(f"  Accuracy: {acc:.4f}, Op accuracy: {op_acc:.4f}")
+
+    print(f"\nStarting training...")
+    print("-" * 70)
+
+    best_acc = acc
+    start_time = time.perf_counter()
+
+    for epoch in range(epochs):
+        model.extractor.train()
+        if unfreeze_layers > 0:
+            model.llm.train()
+
+        max_val = get_curriculum_max(epoch, epochs)
+
+        epoch_loss = 0
+        epoch_losses = {'result': 0, 'a': 0, 'b': 0, 'op': 0}
+        n_batches = 20
+        epoch_start = time.perf_counter()
+
+        for batch_idx in range(n_batches):
+            batch_texts = []
+            batch_a = []
+            batch_b = []
+            batch_op = []
+            batch_result = []
+
+            for _ in range(batch_size):
+                text, a, b, op, result = generate_problem(max_val)
+                batch_texts.append(text)
+                batch_a.append(int_to_bits(a, device))
+                batch_b.append(int_to_bits(b, device))
+                batch_op.append(OPERATIONS.index(op))
+                batch_result.append(int_to_bits(result, device))
+
+            target_a = torch.stack(batch_a)
+            target_b = torch.stack(batch_b)
+            target_op = torch.tensor(batch_op, device=device)
+            target_result = torch.stack(batch_result)
+
+            optimizer.zero_grad()
+
+            pred_bits, a_bits, b_bits, op_logits = model(batch_texts)
+
+            loss, losses = compute_llm_loss(
+                pred_bits, a_bits, b_bits, op_logits,
+                target_result, target_a, target_b, target_op
+            )
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.trainable_parameters(), 1.0)
+            optimizer.step()
+
+            epoch_loss += loss.item()
+            for k in epoch_losses:
+                epoch_losses[k] += losses[k]
+
+            if (batch_idx + 1) % 5 == 0:
+                avg_so_far = epoch_loss / (batch_idx + 1)
+                print(f"  Epoch {epoch+1} batch {batch_idx+1}/{n_batches} | loss: {avg_so_far:.4f}", flush=True)
+
+        epoch_time = time.perf_counter() - epoch_start
+        scheduler.step()
+
+        avg_loss = epoch_loss / n_batches
+        for k in epoch_losses:
+            epoch_losses[k] /= n_batches
+
+        acc, op_acc = evaluate_llm(model, n_samples=300)
+        elapsed = time.perf_counter() - start_time
+
+        marker = " *" if acc > best_acc else ""
+        if acc > best_acc:
+            best_acc = acc
+
+        mem, _ = get_gpu_memory()
+        print(f"Epoch {epoch+1:3d} | Loss: {avg_loss:.4f} | "
+              f"Acc: {acc:.4f}{marker} | OpAcc: {op_acc:.4f} | "
+              f"Range: 0-{max_val} | VRAM: {mem:.0f}MB | Time: {elapsed:.0f}s")
+        print(f"          Losses - result:{epoch_losses['result']:.4f} "
+              f"a:{epoch_losses['a']:.4f} b:{epoch_losses['b']:.4f} "
+              f"op:{epoch_losses['op']:.4f}")
+
+    print("\n" + "=" * 70)
+    print(" FINAL EVALUATION")
+    print("=" * 70)
+
+    acc, op_acc = evaluate_llm(model, n_samples=1000)
+    print(f"Final accuracy: {acc:.4f}")
+    print(f"Final op accuracy: {op_acc:.4f}")
+    print(f"Best accuracy: {best_acc:.4f}")
+
+    print("\nSample predictions:")
+    for _ in range(10):
+        text, a, b, op, expected = generate_problem()
+        with torch.no_grad():
+            result_bits, a_bits, b_bits, op_logits = model([text])
+        pred = bits_to_int(result_bits[0])
+        pred_a = bits_to_int(a_bits[0])
+        pred_b = bits_to_int(b_bits[0])
+        pred_op = OPERATIONS[op_logits[0].argmax().item()]
+
+        status = "OK" if pred == expected else "WRONG"
+        print(f"  '{text}' = {expected} | pred={pred} (a={pred_a}, b={pred_b}, op={pred_op}) [{status}]")
+
+    save_path = "D:/8bit-threshold-computer/llm_integration/trained/llm.pt"
+    save_dict = {
+        'extractor_state_dict': model.extractor.state_dict(),
+        'final_accuracy': acc,
+        'best_accuracy': best_acc,
+        'unfreeze_layers': unfreeze_layers,
+    }
+    if unfreeze_layers > 0:
+        save_dict['llm_state_dict'] = {
+            k: v for k, v in model.llm.state_dict().items()
+            if any(f'layers.{i}.' in k for i in range(len(model.llm.model.layers) - unfreeze_layers, len(model.llm.model.layers)))
+        }
+    torch.save(save_dict, save_path)
+    print(f"\nSaved to: {save_path}")
+
+    return model, acc
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description='Unified training for threshold circuit LLM integration',
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Modes:
+  router    - Train only OpRouter with ground truth bits (sanity check)
+  interface - Train BitEncoder + OpRouter with ground truth bits (sanity check)
+  llm       - Train extractor with LLM hidden states (the real training)
+
+LLM options:
+  --unfreeze_layers N  Fine-tune top N transformer layers
+  --extract_layer N    Extract from layer N (-1 = last)
+  --position_extract   Use position-specific extraction
+  --digit_pred         Predict digits instead of bits
+
+Baked-in: curriculum learning (0-9 -> 0-99 -> 0-255), 2x loss weight for a/b
+
+Examples:
+  python train.py --mode llm --epochs 100
+  python train.py --mode llm --position_extract
+  python train.py --mode llm --digit_pred --extract_layer 12
+  python train.py --mode llm --unfreeze_layers 4 --batch_size 4096
+        """
+    )
+    parser.add_argument('--mode', type=str, required=True,
+                        choices=['router', 'interface', 'llm'],
+                        help='Training mode')
+    parser.add_argument('--epochs', type=int, default=100, help='Number of epochs')
+    parser.add_argument('--batch_size', type=int, default=256, help='Batch size')
+    parser.add_argument('--lr', type=float, default=None,
+                        help='Learning rate (default: mode-specific)')
+    parser.add_argument('--unfreeze_layers', type=int, default=0,
+                        help='Unfreeze top N transformer layers (default 0 = frozen)')
+    parser.add_argument('--extract_layer', type=int, default=-1,
+                        help='Which layer to extract from (-1 = last)')
+    parser.add_argument('--position_extract', action='store_true',
+                        help='Use position-specific extraction')
+    parser.add_argument('--digit_pred', action='store_true',
+                        help='Predict digits instead of bits')
+    parser.add_argument('--device', type=str, default='cuda', help='Device')
+    args = parser.parse_args()
+
+    torch.manual_seed(42)
+    random.seed(42)
+
+    if args.mode == 'router':
+        lr = args.lr if args.lr is not None else 1e-2
+        train_router(epochs=args.epochs, batch_size=args.batch_size, lr=lr, device=args.device)
+
+    elif args.mode == 'interface':
+        lr = args.lr if args.lr is not None else 1e-3
+        model, fitness = train_interface(
+            epochs=args.epochs, batch_size=args.batch_size, lr=lr, device=args.device
+        )
+
+        print("\n" + "=" * 70)
+        print(" EXPERIMENT SUMMARY")
+        print("=" * 70)
+        print(f"\n  Control (Vanilla SmolLM2-360M):     11.90%")
+        print(f"  Experimental (Trained Interface):   {100*fitness:.2f}%")
+        if fitness > 0:
+            print(f"\n  Improvement: {100*(fitness - 0.119)/0.119:.1f}%")
+
+        if fitness >= 0.99:
+            print("\n  CONCLUSION: Frozen threshold circuits + trained interface")
+            print("              achieves near-perfect arithmetic accuracy.")
+            print("              Core thesis VALIDATED.")
+        else:
+            print(f"\n  CONCLUSION: Further training or architecture changes needed.")
+            print(f"              Current gap: {100*(1.0 - fitness):.2f}%")
+
+    elif args.mode == 'llm':
+        lr = args.lr if args.lr is not None else 3e-4
+        train_llm(
+            epochs=args.epochs,
+            batch_size=args.batch_size,
+            lr=lr,
+            unfreeze_layers=args.unfreeze_layers,
+            extract_layer=args.extract_layer,
+            position_extract=args.position_extract,
+            digit_pred=args.digit_pred,
+            device=args.device
+        )
+
+
+if __name__ == "__main__":
+    main()

llm_integration/train_llm.py

@@ -1,387 +0,0 @@
-"""
-LLM Integration Training
-
-Train interface layers to extract operands from SmolLM2 hidden states.
-The hard part: learning to parse "47 + 86" into bits from embeddings.
-"""
-
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import random
-import time
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from circuits import FrozenThresholdCircuits, heaviside_ste
-
-DEVICE = 'cuda'
-MODEL_ID = 'HuggingFaceTB/SmolLM2-360M-Instruct'
-
-OPERATIONS = ['add', 'sub', 'mul', 'gt', 'lt', 'eq']
-OP_SYMBOLS = {'add': '+', 'sub': '-', 'mul': '*', 'gt': '>', 'lt': '<', 'eq': '=='}
-
-
-class HiddenStateExtractor(nn.Module):
-    """
-    Extracts operands and operation from LLM hidden states.
-    This is the hard part - must learn to parse numbers from embeddings.
-    """
-
-    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256):
-        super().__init__()
-
-        self.a_extractor = nn.Sequential(
-            nn.Linear(hidden_dim, intermediate_dim),
-            nn.GELU(),
-            nn.Linear(intermediate_dim, 8),
-        )
-
-        self.b_extractor = nn.Sequential(
-            nn.Linear(hidden_dim, intermediate_dim),
-            nn.GELU(),
-            nn.Linear(intermediate_dim, 8),
-        )
-
-        self.op_router = nn.Sequential(
-            nn.Linear(hidden_dim, intermediate_dim),
-            nn.GELU(),
-            nn.Linear(intermediate_dim, len(OPERATIONS)),
-        )
-
-    def forward(self, hidden_states: torch.Tensor):
-        """
-        Args:
-            hidden_states: [batch, hidden_dim] from LLM
-
-        Returns:
-            a_bits: [batch, 8]
-            b_bits: [batch, 8]
-            op_logits: [batch, 6]
-        """
-        a_logits = self.a_extractor(hidden_states)
-        b_logits = self.b_extractor(hidden_states)
-        op_logits = self.op_router(hidden_states)
-
-        a_soft = torch.sigmoid(a_logits)
-        b_soft = torch.sigmoid(b_logits)
-
-        a_hard = heaviside_ste(a_logits)
-        b_hard = heaviside_ste(b_logits)
-
-        a_bits = a_hard - a_soft.detach() + a_soft
-        b_bits = b_hard - b_soft.detach() + b_soft
-
-        return a_bits, b_bits, op_logits
-
-
-class AugmentedArithmeticModel(nn.Module):
-    """
-    SmolLM2 + frozen threshold circuits.
-    Trains only the extraction interface.
-    """
-
-    def __init__(self, device: str = 'cuda'):
-        super().__init__()
-        self.device = device
-
-        print("[1/4] Loading tokenizer...", flush=True)
-        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        self.tokenizer.padding_side = 'left'
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-        print("  Tokenizer loaded.", flush=True)
-
-        print("[2/4] Loading SmolLM2-360M...", flush=True)
-        self.llm = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.float16,
-            device_map=device,
-            output_hidden_states=True
-        )
-        self.llm.eval()
-        for param in self.llm.parameters():
-            param.requires_grad = False
-
-        hidden_dim = self.llm.config.hidden_size
-        llm_params = sum(p.numel() for p in self.llm.parameters())
-        print(f"  LLM loaded. Hidden dim: {hidden_dim}, Params: {llm_params:,}", flush=True)
-
-        print("[3/4] Loading threshold circuits...", flush=True)
-        self.circuits = FrozenThresholdCircuits(device=device)
-        print(f"  Circuits loaded. {len(self.circuits.weights)} tensors", flush=True)
-
-        print("[4/4] Initializing extractor...", flush=True)
-        self.extractor = HiddenStateExtractor(
-            hidden_dim=hidden_dim,
-            intermediate_dim=256
-        ).to(device)
-
-        trainable = sum(p.numel() for p in self.extractor.parameters())
-        print(f"  Extractor ready. Trainable params: {trainable:,}", flush=True)
-
-    def get_hidden_states(self, texts: list[str]) -> torch.Tensor:
-        """Get hidden states from last layer for each input."""
-        inputs = self.tokenizer(
-            texts,
-            return_tensors='pt',
-            padding=True,
-            truncation=True,
-            max_length=64
-        ).to(self.device)
-
-        with torch.no_grad():
-            outputs = self.llm(**inputs, output_hidden_states=True)
-
-        last_hidden = outputs.hidden_states[-1]
-        mask = inputs.attention_mask
-        seq_lens = mask.sum(dim=1) - 1
-        batch_size = last_hidden.shape[0]
-
-        final_hidden = torch.stack([
-            last_hidden[i, seq_lens[i], :]
-            for i in range(batch_size)
-        ])
-
-        return final_hidden.float()
-
-    def forward(self, texts: list[str]):
-        """
-        Full forward pass: text → hidden states → extractor → circuits → result
-        """
-        hidden = self.get_hidden_states(texts)
-
-        a_bits, b_bits, op_logits = self.extractor(hidden)
-
-        op_probs = torch.softmax(op_logits, dim=-1)
-
-        result_bits = self.circuits(a_bits, b_bits, op_probs)
-
-        return result_bits, a_bits, b_bits, op_logits
-
-
-def generate_problem():
-    """Generate a random arithmetic problem."""
-    a = random.randint(0, 255)
-    b = random.randint(0, 255)
-    op = random.choice(OPERATIONS)
-
-    sym = OP_SYMBOLS[op]
-    text = f"{a} {sym} {b}"
-
-    if op == 'add':
-        result = (a + b) & 0xFF
-    elif op == 'sub':
-        result = (a - b) & 0xFF
-    elif op == 'mul':
-        result = (a * b) & 0xFF
-    elif op == 'gt':
-        result = 1 if a > b else 0
-    elif op == 'lt':
-        result = 1 if a < b else 0
-    elif op == 'eq':
-        result = 1 if a == b else 0
-
-    return text, a, b, op, result
-
-
-def int_to_bits(val: int, device: str = 'cuda') -> torch.Tensor:
-    bits = torch.zeros(8, device=device)
-    for i in range(8):
-        bits[7-i] = (val >> i) & 1
-    return bits
-
-
-def bits_to_int(bits: torch.Tensor) -> int:
-    val = 0
-    for i in range(8):
-        if bits[i].item() > 0.5:
-            val += 1 << (7-i)
-    return val
-
-
-def compute_loss(pred_bits, a_bits, b_bits, op_logits,
-                 target_result, target_a, target_b, target_op_idx):
-    """
-    Multi-component loss:
-    1. Result bits match expected
-    2. Extracted A bits match actual A
-    3. Extracted B bits match actual B
-    4. Operation classification correct
-    """
-    result_loss = nn.functional.binary_cross_entropy_with_logits(
-        pred_bits, target_result
-    )
-
-    a_loss = nn.functional.binary_cross_entropy(
-        torch.clamp(a_bits, 1e-7, 1-1e-7), target_a
-    )
-    b_loss = nn.functional.binary_cross_entropy(
-        torch.clamp(b_bits, 1e-7, 1-1e-7), target_b
-    )
-
-    op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
-
-    total = result_loss + a_loss + b_loss + op_loss
-
-    return total, {
-        'result': result_loss.item(),
-        'a': a_loss.item(),
-        'b': b_loss.item(),
-        'op': op_loss.item()
-    }
-
-
-def evaluate(model, n_samples: int = 500):
-    """Evaluate on random problems."""
-    model.extractor.eval()
-    correct = 0
-    op_correct = 0
-
-    for _ in range(n_samples):
-        text, a, b, op, expected = generate_problem()
-
-        with torch.no_grad():
-            result_bits, a_bits, b_bits, op_logits = model([text])
-
-        pred_result = bits_to_int(result_bits[0])
-        pred_op = OPERATIONS[op_logits[0].argmax().item()]
-
-        if pred_result == expected:
-            correct += 1
-        if pred_op == op:
-            op_correct += 1
-
-    model.extractor.train()
-    return correct / n_samples, op_correct / n_samples
-
-
-def train(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4):
-    print("=" * 70, flush=True)
-    print(" LLM INTEGRATION TRAINING", flush=True)
-    print(" Learning to extract operands from hidden states", flush=True)
-    print("=" * 70, flush=True)
-
-    print("\nInitializing model...", flush=True)
-    model = AugmentedArithmeticModel(device=DEVICE)
-
-    optimizer = optim.AdamW(model.extractor.parameters(), lr=lr)
-    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
-
-    print(f"\nTraining config:", flush=True)
-    print(f"  Epochs: {epochs}", flush=True)
-    print(f"  Batch size: {batch_size}", flush=True)
-    print(f"  Learning rate: {lr}", flush=True)
-    print(f"  Samples/epoch: {batch_size * 20}", flush=True)
-
-    print(f"\nInitial evaluation (200 samples)...", flush=True)
-    acc, op_acc = evaluate(model, n_samples=200)
-    print(f"  Accuracy: {acc:.4f}, Op accuracy: {op_acc:.4f}", flush=True)
-
-    print(f"\nStarting training...", flush=True)
-    print("-" * 70, flush=True)
-
-    best_acc = acc
-    start_time = time.perf_counter()
-
-    for epoch in range(epochs):
-        model.extractor.train()
-        epoch_loss = 0
-        epoch_losses = {'result': 0, 'a': 0, 'b': 0, 'op': 0}
-        n_batches = 20  # 20 batches * 128 = 2,560 samples/epoch
-
-        for batch_idx in range(n_batches):
-            batch_texts = []
-            batch_a = []
-            batch_b = []
-            batch_op = []
-            batch_result = []
-
-            for _ in range(batch_size):
-                text, a, b, op, result = generate_problem()
-                batch_texts.append(text)
-                batch_a.append(int_to_bits(a, DEVICE))
-                batch_b.append(int_to_bits(b, DEVICE))
-                batch_op.append(OPERATIONS.index(op))
-                batch_result.append(int_to_bits(result, DEVICE))
-
-            target_a = torch.stack(batch_a)
-            target_b = torch.stack(batch_b)
-            target_op = torch.tensor(batch_op, device=DEVICE)
-            target_result = torch.stack(batch_result)
-
-            optimizer.zero_grad()
-
-            pred_bits, a_bits, b_bits, op_logits = model(batch_texts)
-
-            loss, losses = compute_loss(
-                pred_bits, a_bits, b_bits, op_logits,
-                target_result, target_a, target_b, target_op
-            )
-
-            loss.backward()
-            torch.nn.utils.clip_grad_norm_(model.extractor.parameters(), 1.0)
-            optimizer.step()
-
-            epoch_loss += loss.item()
-            for k in epoch_losses:
-                epoch_losses[k] += losses[k]
-
-
-        scheduler.step()
-
-        avg_loss = epoch_loss / n_batches
-        for k in epoch_losses:
-            epoch_losses[k] /= n_batches
-
-        if (epoch + 1) % 5 == 0 or epoch == 0:
-            acc, op_acc = evaluate(model, n_samples=300)
-            elapsed = time.perf_counter() - start_time
-
-            marker = " *" if acc > best_acc else ""
-            if acc > best_acc:
-                best_acc = acc
-
-            print(f"Epoch {epoch+1:3d} | Loss: {avg_loss:.4f} | "
-                  f"Acc: {acc:.4f}{marker} | OpAcc: {op_acc:.4f} | "
-                  f"Time: {elapsed:.0f}s")
-            print(f"          Losses - result:{epoch_losses['result']:.4f} "
-                  f"a:{epoch_losses['a']:.4f} b:{epoch_losses['b']:.4f} "
-                  f"op:{epoch_losses['op']:.4f}")
-
-    print("\n" + "=" * 70)
-    print(" FINAL EVALUATION")
-    print("=" * 70)
-
-    acc, op_acc = evaluate(model, n_samples=1000)
-    print(f"Final accuracy: {acc:.4f}")
-    print(f"Final op accuracy: {op_acc:.4f}")
-    print(f"Best accuracy: {best_acc:.4f}")
-
-    print("\nSample predictions:")
-    for _ in range(10):
-        text, a, b, op, expected = generate_problem()
-        with torch.no_grad():
-            result_bits, a_bits, b_bits, op_logits = model([text])
-        pred = bits_to_int(result_bits[0])
-        pred_a = bits_to_int(a_bits[0])
-        pred_b = bits_to_int(b_bits[0])
-        pred_op = OPERATIONS[op_logits[0].argmax().item()]
-
-        status = "OK" if pred == expected else "WRONG"
-        print(f"  '{text}' = {expected} | pred={pred} (a={pred_a}, b={pred_b}, op={pred_op}) [{status}]")
-
-    save_path = "D:/8bit-threshold-computer/llm_integration/trained_extractor.pt"
-    torch.save({
-        'extractor_state_dict': model.extractor.state_dict(),
-        'final_accuracy': acc,
-        'best_accuracy': best_acc,
-    }, save_path)
-    print(f"\nSaved to: {save_path}")
-
-    return model, acc
-
-
-if __name__ == "__main__":
-    random.seed(42)
-    torch.manual_seed(42)
-
-    train(epochs=100, batch_size=384, lr=3e-4)

llm_integration/train_passthrough.py

@@ -1,182 +0,0 @@
-"""
-Training script for ThresholdALU interface layers.
-Trains encoder/router to correctly use frozen threshold circuits.
-"""
-
-import torch
-import torch.nn as nn
-import torch.optim as optim
-import time
-import argparse
-from model import ThresholdALU, DirectCircuitModel
-from fitness import generate_batch, compute_fitness, compute_loss, OPERATIONS
-
-
-def train(
-    epochs: int = 100,
-    batch_size: int = 512,
-    lr: float = 1e-3,
-    eval_interval: int = 10,
-    eval_samples: int = 2000,
-    device: str = 'cuda'
-):
-    print("=" * 70)
-    print(" THRESHOLD ALU INTERFACE TRAINING")
-    print("=" * 70)
-
-    print("\n[1/4] Verifying frozen circuits...")
-    direct_model = DirectCircuitModel(device=device)
-
-    def direct_fn(a, b, op):
-        return direct_model(a, b, op)
-
-    circuit_fitness = compute_fitness(direct_fn, n_samples=1000, device=device)
-    print(f"  Frozen circuit fitness: {circuit_fitness:.4f}")
-    if circuit_fitness < 0.999:
-        print("  ERROR: Circuits not achieving 100%. Aborting.")
-        return
-    print("  STATUS: PASS")
-
-    print("\n[2/4] Initializing model...")
-    model = ThresholdALU(device=device)
-
-    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
-    print(f"  Trainable parameters: {trainable_params:,}")
-
-    def model_fn(a, b, op):
-        return model(a, b, op)
-
-    initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
-    print(f"  Initial fitness: {initial_fitness:.4f}")
-
-    print("\n[3/4] Setting up training...")
-    optimizer = optim.AdamW(model.parameters(), lr=lr)
-    scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
-
-    print(f"  Optimizer: AdamW")
-    print(f"  Learning rate: {lr}")
-    print(f"  Batch size: {batch_size}")
-    print(f"  Epochs: {epochs}")
-
-    print("\n[4/4] Training...")
-    print("-" * 70)
-
-    best_fitness = initial_fitness
-    start_time = time.perf_counter()
-
-    for epoch in range(epochs):
-        model.train()
-        epoch_loss = 0.0
-        n_batches = 100
-
-        for _ in range(n_batches):
-            batch = generate_batch(batch_size, device)
-
-            optimizer.zero_grad()
-
-            pred_bits = model(batch['a_bits'], batch['b_bits'], batch['op_onehot'])
-
-            loss = compute_loss(pred_bits, batch['expected_bits'])
-
-            loss.backward()
-            optimizer.step()
-
-            epoch_loss += loss.item()
-
-        scheduler.step()
-
-        avg_loss = epoch_loss / n_batches
-
-        if (epoch + 1) % eval_interval == 0 or epoch == 0:
-            model.eval()
-            fitness, details = compute_fitness(
-                model_fn, n_samples=eval_samples, device=device, return_details=True
-            )
-
-            elapsed = time.perf_counter() - start_time
-
-            if fitness > best_fitness:
-                best_fitness = fitness
-                marker = " *"
-            else:
-                marker = ""
-
-            print(f"Epoch {epoch+1:4d} | Loss: {avg_loss:.4f} | "
-                  f"Fitness: {fitness:.4f}{marker} | "
-                  f"LR: {scheduler.get_last_lr()[0]:.2e} | "
-                  f"Time: {elapsed:.1f}s")
-
-            if fitness >= 0.9999:
-                print("\n" + "=" * 70)
-                print(" TARGET ACHIEVED: 100% FITNESS")
-                print("=" * 70)
-                break
-
-    print("\n" + "=" * 70)
-    print(" TRAINING COMPLETE")
-    print("=" * 70)
-
-    model.eval()
-    final_fitness, details = compute_fitness(
-        model_fn, n_samples=5000, device=device, return_details=True
-    )
-
-    print(f"\nFinal fitness: {final_fitness:.4f}")
-    print(f"Best fitness:  {best_fitness:.4f}")
-    print(f"\nPer-operation breakdown:")
-    for op in OPERATIONS:
-        acc = details['by_op'][op]['accuracy']
-        print(f"  {op:6}: {acc:.4f}")
-
-    print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
-
-    # Save trained model
-    save_path = "D:/8bit-threshold-computer/llm_integration/trained_model.pt"
-    torch.save({
-        'encoder_state_dict': model.encoder.state_dict(),
-        'router_state_dict': model.router.state_dict(),
-        'final_fitness': final_fitness,
-        'best_fitness': best_fitness,
-    }, save_path)
-    print(f"\nSaved trained model to: {save_path}")
-
-    return model, final_fitness
-
-
-def main():
-    parser = argparse.ArgumentParser(description='Train ThresholdALU interface')
-    parser.add_argument('--epochs', type=int, default=200, help='Number of epochs')
-    parser.add_argument('--batch_size', type=int, default=512, help='Batch size')
-    parser.add_argument('--lr', type=float, default=1e-3, help='Learning rate')
-    parser.add_argument('--eval_interval', type=int, default=10, help='Eval every N epochs')
-    parser.add_argument('--device', type=str, default='cuda', help='Device')
-    args = parser.parse_args()
-
-    torch.manual_seed(42)
-
-    model, fitness = train(
-        epochs=args.epochs,
-        batch_size=args.batch_size,
-        lr=args.lr,
-        eval_interval=args.eval_interval,
-        device=args.device
-    )
-
-    print("\n" + "=" * 70)
-    print(" EXPERIMENT SUMMARY")
-    print("=" * 70)
-    print(f"\n  Control (Vanilla SmolLM2-360M):     11.90%")
-    print(f"  Experimental (Trained Interface):   {100*fitness:.2f}%")
-    print(f"\n  Improvement: {100*(fitness - 0.119)/0.119:.1f}%")
-
-    if fitness >= 0.99:
-        print("\n  CONCLUSION: Frozen threshold circuits + trained interface")
-        print("              achieves near-perfect arithmetic accuracy.")
-        print("              Core thesis VALIDATED.")
-    else:
-        print(f"\n  CONCLUSION: Further training or architecture changes needed.")
-        print(f"              Current gap: {100*(1.0 - fitness):.2f}%")
-
-
-if __name__ == "__main__":
-    main()

llm_integration/train_passthrough_router.py

@@ -1,106 +0,0 @@
-"""
-Train only the router with ground truth bits.
-Proves that operation routing can be learned perfectly.
-"""
-
-import torch
-import torch.optim as optim
-import time
-from model import OpRouter
-from circuits import FrozenThresholdCircuits
-from fitness import generate_batch, compute_fitness, compute_loss, OPERATIONS
-
-device = 'cuda'
-
-print("=" * 70)
-print(" ROUTER-ONLY TRAINING (Ground Truth Bits)")
-print("=" * 70)
-
-circuits = FrozenThresholdCircuits(device=device)
-router = OpRouter(input_dim=16 + 6, hidden_dim=64, n_ops=6).to(device)
-
-print(f"\nRouter parameters: {sum(p.numel() for p in router.parameters()):,}")
-
-def model_fn(a_bits, b_bits, op_onehot):
-    x = torch.cat([a_bits, b_bits, op_onehot], dim=-1)
-    op_weights = router(x)
-    return circuits(a_bits, b_bits, op_weights)
-
-initial_fitness = compute_fitness(model_fn, n_samples=1000, device=device)
-print(f"Initial fitness: {initial_fitness:.4f}")
-
-optimizer = optim.AdamW(router.parameters(), lr=1e-2)
-scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)
-
-print("\nTraining...")
-print("-" * 70)
-
-best_fitness = initial_fitness
-start_time = time.perf_counter()
-
-for epoch in range(100):
-    router.train()
-    epoch_loss = 0.0
-
-    for _ in range(100):
-        batch = generate_batch(256, device)
-
-        optimizer.zero_grad()
-
-        x = torch.cat([batch['a_bits'], batch['b_bits'], batch['op_onehot']], dim=-1)
-        op_weights = router(x)
-        pred_bits = circuits(batch['a_bits'], batch['b_bits'], op_weights)
-
-        loss = compute_loss(pred_bits, batch['expected_bits'])
-        loss.backward()
-        optimizer.step()
-
-        epoch_loss += loss.item()
-
-    scheduler.step()
-
-    if (epoch + 1) % 10 == 0 or epoch == 0:
-        router.eval()
-        fitness, details = compute_fitness(model_fn, n_samples=2000, device=device, return_details=True)
-        elapsed = time.perf_counter() - start_time
-
-        if fitness > best_fitness:
-            best_fitness = fitness
-            marker = " *"
-        else:
-            marker = ""
-
-        print(f"Epoch {epoch+1:3d} | Loss: {epoch_loss/100:.4f} | "
-              f"Fitness: {fitness:.4f}{marker} | Time: {elapsed:.1f}s")
-
-        if fitness >= 0.9999:
-            print("\n TARGET: 100% FITNESS ACHIEVED")
-            break
-
-print("\n" + "=" * 70)
-print(" RESULTS")
-print("=" * 70)
-
-router.eval()
-final_fitness, details = compute_fitness(model_fn, n_samples=5000, device=device, return_details=True)
-
-print(f"\nFinal fitness: {final_fitness:.4f}")
-print(f"\nPer-operation:")
-for op in OPERATIONS:
-    acc = details['by_op'][op]['accuracy']
-    print(f"  {op}: {acc:.4f}")
-
-print(f"\nTotal time: {time.perf_counter() - start_time:.1f}s")
-
-if final_fitness >= 0.99:
-    print("\nCONCLUSION: Router successfully learned operation dispatch.")
-    print("           With correct bit encoding, 100% is achievable.")
-
-# Save trained router weights
-save_path = "D:/8bit-threshold-computer/llm_integration/trained_router.pt"
-torch.save({
-    'router_state_dict': router.state_dict(),
-    'final_fitness': final_fitness,
-    'params': sum(p.numel() for p in router.parameters()),
-}, save_path)
-print(f"\nSaved trained router to: {save_path}")

llm_integration/{trained_passthrough_router.pt → trained/router.pt}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2b33772a74d3891031225298d33d57663c36719e438b5bc9f9039f9e57d636df
-size 10147
+oid sha256:0ddfc24cd4a98b65de8d434bb843ebd24f8c902d067201fd7954e7b623a8ebcd
+size 9811