Add HybridExtractor for digit lookup + word number learning

- HybridExtractor: detects digit tokens (32-41) for hardcoded lookup,
uses learned MLP for word numbers ("forty seven plus eighty six")
- int_to_words(): converts 0-255 to English words
- generate_problem(): randomly mixes digit and word formats
- compute_hybrid_loss(): only trains on word samples (digits are free)
- Hybrid is now the default mode for --mode llm

llm_integration/model.py

@@ -745,6 +745,217 @@ class DigitExtractor(nn.Module):
         return a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits
 
 
+class HybridExtractor(nn.Module):
+    """
+    Hybrid extractor that handles both digit tokens and word numbers.
+
+    For digit tokens (32-41): Direct lookup, no training needed
+    For word numbers: Learned MLP extraction from pooled hidden states
+
+    This is the real training target - learning to extract numbers from
+    natural language like "forty seven plus eighty six".
+    """
+
+    DIGIT_TOKENS = set(range(32, 42))
+    OPERATOR_TOKENS = {
+        1232: 0,   # ' +' -> add
+        731: 1,    # ' -' -> sub
+        1672: 2,   # ' *' -> mul
+        2986: 3,   # ' >' -> gt
+        2067: 4,   # ' <' -> lt
+        1758: 5,   # ' ==' -> eq
+    }
+    WORD_OP_TOKENS = {
+        'plus': 0, 'minus': 1, 'times': 2,
+        'greater': 3, 'less': 4, 'equals': 5, 'equal': 5,
+    }
+
+    def __init__(self, hidden_dim: int = 960, intermediate_dim: int = 256, num_heads: int = 4):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+
+        self.attention_pool = AttentionPooling(hidden_dim, num_heads)
+
+        self.a_predictor = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(intermediate_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 256),
+        )
+
+        self.b_predictor = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Dropout(0.1),
+            nn.Linear(intermediate_dim, intermediate_dim),
+            nn.GELU(),
+            nn.Linear(intermediate_dim, 256),
+        )
+
+        self.op_predictor = nn.Sequential(
+            nn.Linear(hidden_dim, intermediate_dim // 2),
+            nn.GELU(),
+            nn.Linear(intermediate_dim // 2, len(OPERATIONS)),
+        )
+
+    def _has_digit_tokens(self, token_ids: torch.Tensor) -> bool:
+        """Check if input contains digit tokens."""
+        for tid in token_ids.tolist():
+            if tid in self.DIGIT_TOKENS:
+                return True
+        return False
+
+    def _extract_from_digits(self, token_ids: torch.Tensor) -> tuple:
+        """
+        Extract values directly from digit tokens (hardcoded lookup).
+        Returns (a_value, b_value, op_idx) or None if pattern not found.
+        """
+        tokens = token_ids.tolist()
+
+        op_pos = -1
+        op_idx = 0
+        for i, tid in enumerate(tokens):
+            if tid in self.OPERATOR_TOKENS:
+                op_pos = i
+                op_idx = self.OPERATOR_TOKENS[tid]
+                break
+
+        if op_pos == -1:
+            return None
+
+        a_digits = []
+        for i in range(op_pos):
+            if tokens[i] in self.DIGIT_TOKENS:
+                a_digits.append(tokens[i] - 32)
+
+        b_start = op_pos + 1
+        if b_start < len(tokens) and tokens[b_start] == 216:
+            b_start += 1
+
+        b_digits = []
+        for i in range(b_start, len(tokens)):
+            if tokens[i] in self.DIGIT_TOKENS:
+                b_digits.append(tokens[i] - 32)
+
+        if not a_digits or not b_digits:
+            return None
+
+        a_val = 0
+        for d in a_digits:
+            a_val = a_val * 10 + d
+
+        b_val = 0
+        for d in b_digits:
+            b_val = b_val * 10 + d
+
+        return min(a_val, 255), min(b_val, 255), op_idx
+
+    def _value_to_bits(self, value: int, device) -> torch.Tensor:
+        """Convert integer to 8-bit tensor."""
+        bits = torch.zeros(8, device=device)
+        for i in range(8):
+            bits[7 - i] = (value >> i) & 1
+        return bits
+
+    def forward(self, hidden: torch.Tensor, mask: torch.Tensor, token_ids: torch.Tensor = None):
+        """
+        Args:
+            hidden: [batch, seq_len, hidden_dim]
+            mask: [batch, seq_len]
+            token_ids: [batch, seq_len] - optional, enables digit lookup
+
+        Returns:
+            a_bits: [batch, 8]
+            b_bits: [batch, 8]
+            op_logits: [batch, 6]
+            a_values: [batch] predicted values (for loss)
+            b_values: [batch] predicted values (for loss)
+            used_lookup: [batch] bool tensor indicating if lookup was used
+        """
+        batch_size = hidden.shape[0]
+        device = hidden.device
+
+        a_bits_list = []
+        b_bits_list = []
+        op_logits_list = []
+        a_values_list = []
+        b_values_list = []
+        used_lookup_list = []
+
+        pooled = self.attention_pool(hidden, mask)
+
+        for i in range(batch_size):
+            lookup_result = None
+            if token_ids is not None:
+                seq_mask = mask[i].bool()
+                valid_len = seq_mask.sum().item()
+                start_pos = hidden.shape[1] - valid_len
+                valid_tokens = token_ids[i, start_pos:]
+
+                if self._has_digit_tokens(valid_tokens):
+                    lookup_result = self._extract_from_digits(valid_tokens)
+
+            if lookup_result is not None:
+                a_val, b_val, op_idx = lookup_result
+                a_bits = self._value_to_bits(a_val, device)
+                b_bits = self._value_to_bits(b_val, device)
+                op_logits = torch.zeros(len(OPERATIONS), device=device)
+                op_logits[op_idx] = 10.0
+
+                a_bits_list.append(a_bits)
+                b_bits_list.append(b_bits)
+                op_logits_list.append(op_logits)
+                a_values_list.append(float(a_val))
+                b_values_list.append(float(b_val))
+                used_lookup_list.append(True)
+            else:
+                sample_pooled = pooled[i]
+
+                a_logits = self.a_predictor(sample_pooled)
+                b_logits = self.b_predictor(sample_pooled)
+                op_logits = self.op_predictor(sample_pooled)
+
+                a_probs = torch.softmax(a_logits, dim=-1)
+                b_probs = torch.softmax(b_logits, dim=-1)
+
+                values = torch.arange(256, device=device, dtype=torch.float32)
+                a_val = (a_probs * values).sum()
+                b_val = (b_probs * values).sum()
+
+                a_bits = self._soft_value_to_bits(a_val, device)
+                b_bits = self._soft_value_to_bits(b_val, device)
+
+                a_bits_list.append(a_bits)
+                b_bits_list.append(b_bits)
+                op_logits_list.append(op_logits)
+                a_values_list.append(a_val)
+                b_values_list.append(b_val)
+                used_lookup_list.append(False)
+
+        a_bits = torch.stack(a_bits_list)
+        b_bits = torch.stack(b_bits_list)
+        op_logits = torch.stack(op_logits_list)
+        a_values = torch.stack([v if isinstance(v, torch.Tensor) else torch.tensor(v, device=device) for v in a_values_list])
+        b_values = torch.stack([v if isinstance(v, torch.Tensor) else torch.tensor(v, device=device) for v in b_values_list])
+        used_lookup = torch.tensor(used_lookup_list, device=device, dtype=torch.bool)
+
+        return a_bits, b_bits, op_logits, a_values, b_values, used_lookup
+
+    def _soft_value_to_bits(self, value: torch.Tensor, device) -> torch.Tensor:
+        """Convert soft value (0-255) to 8-bit representation differentiably."""
+        value = torch.clamp(value, 0, 255)
+        bits = []
+        remaining = value
+        for i in range(7, -1, -1):
+            threshold = 2 ** i
+            bit = torch.sigmoid((remaining - threshold + 0.5) * 10)
+            bits.append(bit)
+            remaining = remaining - bit * threshold
+        return torch.stack(bits)
+
+
 class ArithmeticModel(nn.Module):
     """
     LLM + extractor + frozen threshold circuits.
@@ -753,7 +964,8 @@ class ArithmeticModel(nn.Module):
 
     def __init__(self, device: str = 'cuda', unfreeze_layers: int = 0,
                  extract_layer: int = -1, position_extract: bool = False,
-                 digit_pred: bool = False, positional_digit: bool = False):
+                 digit_pred: bool = False, positional_digit: bool = False,
+                 hybrid: bool = False):
         super().__init__()
         self.device = device
         self.unfreeze_layers = unfreeze_layers
@@ -761,6 +973,7 @@ class ArithmeticModel(nn.Module):
         self.position_extract = position_extract
         self.digit_pred = digit_pred
         self.positional_digit = positional_digit
+        self.hybrid = hybrid
 
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -801,7 +1014,14 @@ class ArithmeticModel(nn.Module):
         print(f"  Circuits loaded. {len(self.circuits.weights)} tensors", flush=True)
 
         print("[4/4] Initializing extractor...", flush=True)
-        if positional_digit:
+        if hybrid:
+            print("  Using HYBRID extraction (digit lookup + word learning)", flush=True)
+            self.extractor = HybridExtractor(
+                hidden_dim=hidden_dim,
+                intermediate_dim=256,
+                num_heads=4
+            ).to(device)
+        elif positional_digit:
             print("  Using POSITIONAL DIGIT extraction (100% proven)", flush=True)
             self.extractor = PositionalDigitExtractor(
                 hidden_dim=hidden_dim
@@ -875,31 +1095,41 @@ class ArithmeticModel(nn.Module):
         """
         hidden, mask, token_ids = self.get_hidden_states(texts)
 
-        if self.positional_digit or self.position_extract:
+        if self.hybrid or self.positional_digit or self.position_extract:
             extractor_out = self.extractor(hidden, mask, token_ids)
         else:
             extractor_out = self.extractor(hidden, mask)
 
-        if self.positional_digit:
+        if self.hybrid:
+            a_bits, b_bits, op_logits, a_values, b_values, used_lookup = extractor_out
+            op_indices_from_tokens = None
+            a_digit_logits, b_digit_logits = None, None
+        elif self.positional_digit:
             a_bits, b_bits, op_logits, op_indices_from_tokens, a_values, b_values, a_digit_logits, b_digit_logits = extractor_out
+            used_lookup = None
         elif self.digit_pred:
             a_bits, b_bits, op_logits, a_digit_logits, b_digit_logits = extractor_out
             op_indices_from_tokens = None
             a_values, b_values = None, None
+            used_lookup = None
         elif self.position_extract:
             a_bits, b_bits, op_logits, op_indices_from_tokens = extractor_out
             a_digit_logits, b_digit_logits = None, None
             a_values, b_values = None, None
+            used_lookup = None
         else:
             a_bits, b_bits, op_logits = extractor_out
             a_digit_logits, b_digit_logits = None, None
             op_indices_from_tokens = None
             a_values, b_values = None, None
+            used_lookup = None
 
         op_probs = torch.softmax(op_logits, dim=-1)
 
         result_bits = self.circuits(a_bits, b_bits, op_probs)
 
+        if self.hybrid:
+            return result_bits, a_bits, b_bits, op_logits, a_values, b_values, used_lookup
         if self.positional_digit:
             return result_bits, a_bits, b_bits, op_logits, op_indices_from_tokens, a_values, b_values, a_digit_logits, b_digit_logits
         if self.digit_pred:

llm_integration/train.py

@@ -39,6 +39,29 @@ from fitness import generate_batch, compute_fitness, compute_loss
 
 DEVICE = 'cuda'
 
+ONES = ['zero', 'one', 'two', 'three', 'four', 'five', 'six', 'seven', 'eight', 'nine',
+        'ten', 'eleven', 'twelve', 'thirteen', 'fourteen', 'fifteen', 'sixteen',
+        'seventeen', 'eighteen', 'nineteen']
+TENS = ['', '', 'twenty', 'thirty', 'forty', 'fifty', 'sixty', 'seventy', 'eighty', 'ninety']
+
+def int_to_words(n: int) -> str:
+    """Convert integer 0-255 to English words."""
+    if n < 0 or n > 255:
+        return str(n)
+    if n < 20:
+        return ONES[n]
+    if n < 100:
+        if n % 10 == 0:
+            return TENS[n // 10]
+        return f"{TENS[n // 10]} {ONES[n % 10]}"
+    if n % 100 == 0:
+        return f"{ONES[n // 100]} hundred"
+    if n % 100 < 20:
+        return f"{ONES[n // 100]} hundred {ONES[n % 100]}"
+    if n % 10 == 0:
+        return f"{ONES[n // 100]} hundred {TENS[(n % 100) // 10]}"
+    return f"{ONES[n // 100]} hundred {TENS[(n % 100) // 10]} {ONES[n % 10]}"
+
 
 def int_to_bits(val: int, device: str = 'cuda') -> torch.Tensor:
     bits = torch.zeros(8, device=device)
@@ -55,14 +78,84 @@ def bits_to_int(bits: torch.Tensor) -> int:
     return val
 
 
+NL_TEMPLATES = {
+    'add': [
+        "What is {a} plus {b}?",
+        "Calculate {a} + {b}",
+        "Add {a} and {b}",
+        "What's the sum of {a} and {b}?",
+        "If I have {a} and get {b} more, how many total?",
+        "{a} + {b} = ?",
+        "Compute {a} plus {b}",
+    ],
+    'sub': [
+        "What is {a} minus {b}?",
+        "Calculate {a} - {b}",
+        "Subtract {b} from {a}",
+        "What's {a} take away {b}?",
+        "If I have {a} and lose {b}, how many left?",
+        "{a} - {b} = ?",
+        "Compute {a} minus {b}",
+    ],
+    'mul': [
+        "What is {a} times {b}?",
+        "Calculate {a} * {b}",
+        "Multiply {a} by {b}",
+        "What's {a} multiplied by {b}?",
+        "{a} * {b} = ?",
+        "Compute {a} times {b}",
+        "What is the product of {a} and {b}?",
+    ],
+    'gt': [
+        "Is {a} greater than {b}?",
+        "Is {a} > {b}?",
+        "Check if {a} is larger than {b}",
+        "Compare: is {a} more than {b}?",
+        "{a} > {b}?",
+    ],
+    'lt': [
+        "Is {a} less than {b}?",
+        "Is {a} < {b}?",
+        "Check if {a} is smaller than {b}",
+        "Compare: is {a} fewer than {b}?",
+        "{a} < {b}?",
+    ],
+    'eq': [
+        "Is {a} equal to {b}?",
+        "Is {a} == {b}?",
+        "Does {a} equal {b}?",
+        "Check if {a} equals {b}",
+        "Are {a} and {b} the same?",
+    ],
+}
+
+
 def generate_problem(max_val: int = 255):
-    """Generate a random arithmetic problem for LLM training."""
+    """
+    Generate a random arithmetic problem for LLM training.
+    Randomly mixes digit and word formats.
+    """
     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}"
+    fmt = random.choice(['digits', 'words', 'nl_digits', 'nl_words'])
+
+    if fmt == 'digits':
+        sym = OP_SYMBOLS[op]
+        text = f"{a} {sym} {b}"
+    elif fmt == 'words':
+        a_word = int_to_words(a)
+        b_word = int_to_words(b)
+        op_word = {'add': 'plus', 'sub': 'minus', 'mul': 'times',
+                   'gt': 'greater than', 'lt': 'less than', 'eq': 'equals'}[op]
+        text = f"{a_word} {op_word} {b_word}"
+    elif fmt == 'nl_digits':
+        template = random.choice(NL_TEMPLATES[op])
+        text = template.format(a=a, b=b)
+    else:
+        template = random.choice(NL_TEMPLATES[op])
+        text = template.format(a=int_to_words(a), b=int_to_words(b))
 
     if op == 'add':
         result = (a + b) & 0xFF
@@ -457,8 +550,51 @@ def compute_positional_digit_loss(pred_bits, op_logits, a_digit_logits_list, b_d
     }
 
 
+def compute_hybrid_loss(pred_bits, a_values, b_values, op_logits, used_lookup,
+                        target_result, target_a_values, target_b_values, target_op_idx,
+                        device, value_weight: float = 1.0):
+    """
+    Loss for hybrid extraction.
+
+    Only compute value loss for samples where lookup was NOT used (word numbers).
+    Samples using digit lookup are already 100% accurate.
+    """
+    result_loss = nn.functional.binary_cross_entropy_with_logits(
+        pred_bits, target_result
+    )
+
+    op_loss = nn.functional.cross_entropy(op_logits, target_op_idx)
+
+    word_mask = ~used_lookup
+    n_words = word_mask.sum().item()
+
+    if n_words > 0:
+        a_word_values = a_values[word_mask]
+        b_word_values = b_values[word_mask]
+        target_a_word = target_a_values[word_mask]
+        target_b_word = target_b_values[word_mask]
+
+        a_value_loss = nn.functional.mse_loss(a_word_values, target_a_word)
+        b_value_loss = nn.functional.mse_loss(b_word_values, target_b_word)
+    else:
+        a_value_loss = torch.tensor(0.0, device=device)
+        b_value_loss = torch.tensor(0.0, device=device)
+
+    total = result_loss + op_loss + value_weight * (a_value_loss + b_value_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_value': a_value_loss.item() if not torch.isnan(a_value_loss) else 10.0,
+        'b_value': b_value_loss.item() if not torch.isnan(b_value_loss) else 10.0,
+        'op': op_loss.item() if not torch.isnan(op_loss) else 10.0,
+        'n_words': n_words,
+        'n_lookup': used_lookup.sum().item()
+    }
+
+
 def evaluate_llm(model, n_samples: int = 500):
-    """Evaluate LLM model on random problems."""
+    """Evaluate LLM model on random problems (mixed digit/word format)."""
     model.extractor.eval()
     correct = 0
     op_correct = 0
@@ -493,10 +629,12 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
     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
-        positional_digit: Use positional digit extraction (100% proven accuracy)
+        position_extract: Use position-specific extraction (legacy)
+        digit_pred: Predict digits instead of bits (legacy)
+        positional_digit: Use positional digit extraction (legacy, 100% on digits only)
     """
+    hybrid = not (positional_digit or position_extract or digit_pred)
+
     print("=" * 70)
     print(" LLM TRAINING")
     if unfreeze_layers > 0:
@@ -505,12 +643,14 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
         print(" LLM frozen")
     if extract_layer != -1:
         print(f" Extracting from layer {extract_layer}")
-    if positional_digit:
-        print(" POSITIONAL DIGIT extraction (100% proven)")
+    if hybrid:
+        print(" HYBRID extraction (digit lookup + word learning)")
+    elif positional_digit:
+        print(" POSITIONAL DIGIT extraction (legacy, 100% on digits only)")
     elif position_extract:
-        print(" Position-specific extraction")
-    if digit_pred:
-        print(" Digit-level prediction")
+        print(" Position-specific extraction (legacy)")
+    elif digit_pred:
+        print(" Digit-level prediction (legacy)")
     print("=" * 70)
 
     print("\nInitializing model...")
@@ -520,7 +660,8 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
         extract_layer=extract_layer,
         position_extract=position_extract,
         digit_pred=digit_pred,
-        positional_digit=positional_digit
+        positional_digit=positional_digit,
+        hybrid=hybrid
     )
 
     optimizer = optim.AdamW(model.trainable_parameters(), lr=lr)
@@ -534,7 +675,7 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
     print(f"  Samples/epoch: {batch_size * 20}")
 
     print(f"\nInitial evaluation (200 samples)...")
-    acc, op_acc = evaluate_llm(model, n_samples=200)
+    acc, op_acc = evaluate_llm(model, 200)
     print(f"  Accuracy: {acc:.4f}, Op accuracy: {op_acc:.4f}")
 
     print(f"\nStarting training...")
@@ -551,7 +692,9 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
         max_val = get_curriculum_max(epoch, epochs)
 
         epoch_loss = 0
-        if positional_digit:
+        if hybrid:
+            epoch_losses = {'result': 0, 'a_value': 0, 'b_value': 0, 'op': 0, 'n_words': 0, 'n_lookup': 0}
+        elif positional_digit:
             epoch_losses = {'result': 0, 'a_digit': 0, 'b_digit': 0, 'op': 0}
         else:
             epoch_losses = {'result': 0, 'a': 0, 'b': 0, 'op': 0}
@@ -589,7 +732,13 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
             outputs = model(batch_texts)
             pred_bits, a_bits, b_bits, op_logits = outputs[0], outputs[1], outputs[2], outputs[3]
 
-            if positional_digit:
+            if hybrid:
+                a_values, b_values, used_lookup = outputs[4], outputs[5], outputs[6]
+                loss, losses = compute_hybrid_loss(
+                    pred_bits, a_values, b_values, op_logits, used_lookup,
+                    target_result, target_a_values, target_b_values, target_op, device
+                )
+            elif positional_digit:
                 a_digit_logits_list = outputs[7]
                 b_digit_logits_list = outputs[8]
                 loss, losses = compute_positional_digit_loss(
@@ -621,7 +770,7 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
         for k in epoch_losses:
             epoch_losses[k] /= n_batches
 
-        acc, op_acc = evaluate_llm(model, n_samples=300)
+        acc, op_acc = evaluate_llm(model, 300)
         elapsed = time.perf_counter() - start_time
 
         marker = " *" if acc > best_acc else ""
@@ -632,7 +781,11 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
         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")
-        if positional_digit:
+        if hybrid:
+            print(f"          Losses - result:{epoch_losses['result']:.4f} "
+                  f"a_val:{epoch_losses['a_value']:.4f} b_val:{epoch_losses['b_value']:.4f} "
+                  f"op:{epoch_losses['op']:.4f} | words:{epoch_losses['n_words']:.0f} lookup:{epoch_losses['n_lookup']:.0f}")
+        elif positional_digit:
             print(f"          Losses - result:{epoch_losses['result']:.4f} "
                   f"a_digit:{epoch_losses['a_digit']:.4f} b_digit:{epoch_losses['b_digit']:.4f} "
                   f"op:{epoch_losses['op']:.4f}")
@@ -651,7 +804,7 @@ def train_llm(epochs: int = 100, batch_size: int = 256, lr: float = 3e-4,
     print(" FINAL EVALUATION")
     print("=" * 70)
 
-    acc, op_acc = evaluate_llm(model, n_samples=1000)
+    acc, op_acc = evaluate_llm(model, 1000)
     print(f"Final accuracy: {acc:.4f}")
     print(f"Final op accuracy: {op_acc:.4f}")
     print(f"Best accuracy: {best_acc:.4f}")
@@ -717,19 +870,19 @@ Examples:
                         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('--batch_size', type=int, default=512, help='Batch size (default: 512)')
     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('--extract_layer', type=int, default=0,
+                        help='Which layer to extract from (default: 0 = embeddings, best for digits)')
     parser.add_argument('--position_extract', action='store_true',
-                        help='Use position-specific extraction')
+                        help='Use position-specific extraction (legacy)')
     parser.add_argument('--digit_pred', action='store_true',
-                        help='Predict digits instead of bits')
-    parser.add_argument('--positional_digit', action='store_true',
-                        help='Use positional digit extraction (100% proven accuracy)')
+                        help='Predict digits instead of bits (legacy)')
+    parser.add_argument('--positional_digit', action='store_true', default=False,
+                        help='Use positional digit extraction (legacy, 100%% on digits only)')
     parser.add_argument('--device', type=str, default='cuda', help='Device')
     args = parser.parse_args()