From b25662f7d37a13377da166c93d227d0c810df7d8 Mon Sep 17 00:00:00 2001
From: Void Agent <void@jayrup.hermes>
Date: Thu, 30 Jul 2026 15:33:07 +0100
Subject: Add controlled frequency ablation experiment (doubled-q)

---
 src/freq_ablation.py   | 139 ++++++++++++++++++++
 src/freq_experiment.py | 344 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 483 insertions(+)
 create mode 100644 src/freq_ablation.py
 create mode 100644 src/freq_experiment.py

diff --git a/src/freq_ablation.py b/src/freq_ablation.py
new file mode 100644
index 0000000..c011140
--- /dev/null
+++ b/src/freq_ablation.py
@@ -0,0 +1,139 @@
+"""
+Controlled frequency ablation: train model on doubled-q Shakespeare,
+compare J-lens norms for 'q' vs control.
+
+Usage (inside Docker on meru):
+  python3 src/freq_ablation.py
+"""
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import torch, numpy as np, pickle
+
+from model import GPT, GPTConfig
+import jlens_v2
+from jlens_v2 import compute_jlens_layer
+
+device = 'cuda'
+MODEL_ARGS = dict(n_layer=6, n_head=6, n_embd=384, block_size=128,
+                  bias=False, vocab_size=65, dropout=0.2)
+
+def get_batch(train_data, val_data, split, block_size=128, batch_size=32):
+    data = train_data if split == 'train' else val_data
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy(data[i:i+block_size].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
+    return x.to(device), y.to(device)
+
+def train(data_dir, out_dir, max_iters=5000):
+    """Train a model and return best checkpoint path."""
+    train_data = np.memmap(f'{data_dir}/train.bin', dtype=np.uint16, mode='r')
+    val_data = np.memmap(f'{data_dir}/val.bin', dtype=np.uint16, mode='r')
+    
+    config = GPTConfig(**MODEL_ARGS)
+    model = GPT(config).to(device)
+    print(f"  Params: {sum(p.numel() for p in model.parameters())/1e6:.2f}M")
+    
+    optimizer = model.configure_optimizers(weight_decay=0.1, lr=1e-3, 
+                                            betas=(0.9, 0.99), device_type='cuda')
+    os.makedirs(out_dir, exist_ok=True)
+    best_val = 1e9
+    
+    for it in range(max_iters):
+        if it % 500 == 0:
+            model.eval()
+            losses = {}
+            for split in ['train', 'val']:
+                lv = []
+                for _ in range(50):
+                    X, Y = get_batch(train_data, val_data, split)
+                    _, loss = model(X, Y)
+                    lv.append(loss.item())
+                losses[split] = np.mean(lv)
+            model.train()
+            print(f"  step {it}: train={losses['train']:.4f}, val={losses['val']:.4f}")
+            if losses['val'] < best_val:
+                best_val = losses['val']
+                torch.save({'model': model.state_dict(), 'model_args': MODEL_ARGS,
+                           'best_val_loss': best_val}, f'{out_dir}/ckpt.pt')
+        
+        X, Y = get_batch(train_data, val_data, 'train')
+        logits, loss = model(X, Y)
+        loss.backward()
+        optimizer.step()
+        optimizer.zero_grad(set_to_none=True)
+        if it % 500 == 0:
+            print(f"  iter {it}: loss={loss.item():.4f}")
+    
+    print(f"  Done. Best val: {best_val:.4f}")
+    return model
+
+def run_jlens_on_model(model, data_dir, target_char, n_batches=10):
+    """Compute J-lens and return norm for target_char."""
+    train_data = np.memmap(f'{data_dir}/train.bin', dtype=np.uint16, mode='r')
+    with open(f'{data_dir}/meta.pkl', 'rb') as f:
+        meta = pickle.load(f)
+    target_id = meta['stoi'][target_char]
+    
+    batch_size = 16
+    block_size = MODEL_ARGS['block_size']
+    batches = []
+    for _ in range(n_batches):
+        ix = torch.randint(len(train_data) - block_size, (batch_size,))
+        x = torch.stack([torch.from_numpy(
+            train_data[i:i+block_size].astype(np.int64)) for i in ix])
+        y = torch.stack([torch.from_numpy(
+            train_data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
+        batches.append((x, y))
+    
+    results = {}
+    for layer_idx in [2, 3, 4]:
+        jlens = compute_jlens_layer(model, layer_idx, batches, device)
+        results[layer_idx] = jlens[target_id].norm().item()
+    
+    return results
+
+# ── MAIN ───────────────────────────────────────────────
+print("=" * 60)
+print("CONTROLLED FREQUENCY ABLATION: 'q' DOUBLED")
+print("=" * 60)
+
+# 1. Train control model (or use existing)
+print("\n[1/4] Using existing control model...")
+ctrl_model, _ = jlens_v2.load_model('out-shakespeare-char/ckpt.pt', device)
+ctrl_model.eval()
+
+# 2. Train modified model
+print("\n[2/4] Training modified model (doubled 'q')...")
+mod_model = train('data/freq_experiment/doubled_q', 'out-freq-doubled-q')
+mod_model.eval()
+
+# 3. J-lens on both
+print("\n[3/4] Computing J-lens...")
+print("  Control model...")
+ctrl_norms = run_jlens_on_model(ctrl_model, 'data/shakespeare_char', 'q')
+print("  Modified model...")
+mod_norms = run_jlens_on_model(mod_model, 'data/freq_experiment/doubled_q', 'q')
+
+# 4. Compare
+print("\n[4/4] RESULTS")
+print("=" * 60)
+print(f"  Character: 'q'")
+print(f"  Original frequency: 0.055%")
+print(f"  Doubled frequency:  0.109%")
+print()
+print(f"  {'Layer':<8} {'Control':>10} {'Modified':>10} {'Change':>10}")
+print(f"  {'-'*8} {'-'*10} {'-'*10} {'-'*10}")
+for layer_idx in sorted(ctrl_norms.keys()):
+    c = ctrl_norms[layer_idx]
+    m = mod_norms[layer_idx]
+    pct = (m - c) / c * 100
+    direction = "↓" if pct < 0 else "↑"
+    print(f"  L{layer_idx:<7} {c:>10.4f} {m:>10.4f} {direction}{abs(pct):>8.1f}%")
+
+print()
+avg_ctrl = np.mean(list(ctrl_norms.values()))
+avg_mod = np.mean(list(mod_norms.values()))
+avg_pct = (avg_mod - avg_ctrl) / avg_ctrl * 100
+print(f"  Average: {avg_ctrl:.4f} → {avg_mod:.4f} ({avg_pct:+.1f}%)")
+print(f"  HYPOTHESIS: doubling frequency should REDUCE J-lens norm")
+print(f"  RESULT: {'CONFIRMED' if avg_pct < -2 else 'INCONCLUSIVE' if abs(avg_pct) < 2 else 'CONTRADICTED'}")
diff --git a/src/freq_experiment.py b/src/freq_experiment.py
new file mode 100644
index 0000000..4050c16
--- /dev/null
+++ b/src/freq_experiment.py
@@ -0,0 +1,344 @@
+"""
+Controlled frequency experiment for J-lens.
+
+Tests causality: if we artificially double the frequency of a character,
+does its J-lens norm predictably drop?
+
+Hypothesis (from chain rule):
+  J-lens norm ∝ (1 - p_avg), where p_avg is average predicted probability.
+  Doubling frequency → model learns higher p → J-lens norm drops.
+
+Control: train two identical models on:
+  A) Original Shakespeare (baseline)
+  B) Modified Shakespeare with 'q' frequency doubled
+"""
+
+import sys, os
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+def create_modified_dataset(input_path, output_path, target_char='q', factor=2.0):
+    """Double the frequency of target_char in the text.
+    Inserts target_char at random positions until frequency is doubled.
+    """
+    import random
+    random.seed(42)
+    
+    with open(input_path, 'r') as f:
+        text = f.read()
+    
+    target_count = text.count(target_char)
+    target_freq = target_count / len(text)
+    print(f"Original: {target_count} occurrences of '{target_char}', "
+          f"frequency = {target_freq:.4%}")
+    
+    # Insert additional copies at random positions
+    extra_needed = int(target_count * (factor - 1))
+    positions = sorted(random.sample(range(len(text)), extra_needed))
+    
+    modified = list(text)
+    for i, pos in enumerate(positions):
+        modified.insert(pos + i, target_char)  # offset by previous insertions
+    
+    modified_text = ''.join(modified)
+    new_count = modified_text.count(target_char)
+    new_freq = new_count / len(modified_text)
+    print(f"Modified: {new_count} occurrences of '{target_char}', "
+          f"frequency = {new_freq:.4%}")
+    
+    with open(output_path, 'w') as f:
+        f.write(modified_text)
+    
+    return modified_text
+
+
+def prepare_data(input_path, output_dir):
+    """Run nanoGPT's prepare.py on a text file."""
+    import subprocess
+    # Write a temp prepare script
+    import os
+    os.makedirs(output_dir, exist_ok=True)
+    
+    # Read the text
+    with open(input_path, 'r') as f:
+        text = f.read()
+    
+    chars = sorted(list(set(text)))
+    vocab_size = len(chars)
+    stoi = {ch: i for i, ch in enumerate(chars)}
+    itos = {i: ch for i, ch in enumerate(chars)}
+    
+    # Encode
+    import numpy as np
+    data = np.array([stoi[ch] for ch in text], dtype=np.uint16)
+    n = int(0.9 * len(data))
+    train_data = data[:n]
+    val_data = data[n:]
+    
+    train_data.tofile(os.path.join(output_dir, 'train.bin'))
+    val_data.tofile(os.path.join(output_dir, 'val.bin'))
+    
+    import pickle
+    with open(os.path.join(output_dir, 'meta.pkl'), 'wb') as f:
+        pickle.dump({'stoi': stoi, 'itos': itos, 'vocab_size': vocab_size}, f)
+    
+    # Also save the raw text
+    with open(os.path.join(output_dir, 'input.txt'), 'w') as f:
+        f.write(text)
+    
+    print(f"Prepared {output_dir}: {len(train_data):,} train, "
+          f"{len(val_data):,} val, {vocab_size} vocab")
+    return stoi, itos, vocab_size
+
+
+def train_model(data_dir, out_dir, device='cuda'):
+    """Train nanoGPT on a prepared dataset."""
+    import subprocess
+    
+    # Use the config but override dataset path and output
+    import torch
+    from model import GPT, GPTConfig
+    
+    # Load and train
+    script = f"""
+import sys
+sys.path.insert(0, '.')
+import torch
+import numpy as np
+import pickle
+import os
+from model import GPT, GPTConfig
+
+# Load data
+data_dir = '{data_dir}'
+train_data = np.memmap(f'{{data_dir}}/train.bin', dtype=np.uint16, mode='r')
+val_data = np.memmap(f'{{data_dir}}/val.bin', dtype=np.uint16, mode='r')
+with open(f'{{data_dir}}/meta.pkl', 'rb') as f:
+    meta = pickle.load(f)
+
+# Config — same as train_shakespeare_char but smaller for speed
+out_dir = '{out_dir}'
+eval_interval = 500
+eval_iters = 100
+log_interval = 100
+always_save_checkpoint = False
+wandb_log = False
+dataset = 'custom'
+gradient_accumulation_steps = 1
+batch_size = 32
+block_size = 128
+n_layer = 6
+n_head = 6
+n_embd = 384
+dropout = 0.2
+learning_rate = 1e-3
+max_iters = 5000
+lr_decay_iters = 5000
+min_lr = 1e-4
+beta2 = 0.99
+warmup_iters = 100
+dtype = 'float32'
+flash = False
+device = '{device}'
+compile = False
+vocab_size = meta['vocab_size']
+
+# Build model
+model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd,
+                  block_size=block_size, bias=False, vocab_size=vocab_size,
+                  dropout=dropout)
+gptconf = GPTConfig(**model_args)
+model = GPT(gptconf)
+model.to(device)
+
+# Optimizer
+optimizer = model.configure_optimizers(weight_decay=0.1, learning_rate=learning_rate,
+                                        betas=(0.9, beta2), device_type=device)
+scaler = torch.amp.GradScaler('cuda', enabled=False)
+
+def get_batch(split):
+    data = train_data if split == 'train' else val_data
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy(data[i:i+block_size].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
+    x, y = x.to(device), y.to(device)
+    return x, y
+
+@torch.no_grad()
+def estimate_loss():
+    out = {{}}
+    model.eval()
+    for split in ['train', 'val']:
+        losses = torch.zeros(eval_iters)
+        for k in range(eval_iters):
+            X, Y = get_batch(split)
+            logits, loss = model(X, Y)
+            losses[k] = loss.item()
+        out[split] = losses.mean()
+    model.train()
+    return out
+
+print(f"Training model: {{model_args}}")
+print(f"Parameters: {{sum(p.numel() for p in model.parameters())/1e6:.2f}}M")
+
+os.makedirs(out_dir, exist_ok=True)
+best_val_loss = 1e9
+
+for iter_num in range(max_iters):
+    if iter_num % eval_interval == 0:
+        losses = estimate_loss()
+        print(f"step {{iter_num}}: train loss {{losses['train']:.4f}}, val loss {{losses['val']:.4f}}")
+        if losses['val'] < best_val_loss:
+            best_val_loss = losses['val']
+            checkpoint = {{
+                'model': model.state_dict(),
+                'optimizer': optimizer.state_dict(),
+                'model_args': model_args,
+                'iter_num': iter_num,
+                'best_val_loss': best_val_loss,
+            }}
+            torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))
+    
+    X, Y = get_batch('train')
+    with torch.amp.autocast(device_type=device, dtype=torch.float32):
+        logits, loss = model(X, Y)
+    
+    scaler.scale(loss).backward()
+    scaler.step(optimizer)
+    scaler.update()
+    optimizer.zero_grad(set_to_none=True)
+    
+    if iter_num % log_interval == 0:
+        print(f"iter {{iter_num}}: loss {{loss.item():.4f}}")
+
+print(f"Training complete. Best val loss: {{best_val_loss:.4f}}")
+"""
+    
+    import tempfile
+    with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
+        f.write(script)
+        script_path = f.name
+    
+    result = subprocess.run(['python3', script_path], capture_output=True, text=True)
+    os.unlink(script_path)
+    print(result.stdout)
+    if result.returncode != 0:
+        print("STDERR:", result.stderr)
+    return result.returncode == 0
+
+
+def main():
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input', default='data/shakespeare_char/input.txt')
+    parser.add_argument('--target_char', default='q')
+    parser.add_argument('--factor', type=float, default=2.0)
+    parser.add_argument('--skip_train', action='store_true')
+    parser.add_argument('--device', default='cuda')
+    args = parser.parse_args()
+    
+    # 1. Create modified dataset
+    print("=" * 60)
+    print("STEP 1: Creating modified dataset")
+    print("=" * 60)
+    modified_input = f'data/freq_experiment/modified_{args.target_char}.txt'
+    os.makedirs('data/freq_experiment', exist_ok=True)
+    create_modified_dataset(args.input, modified_input, args.target_char, args.factor)
+    
+    # 2. Prepare both datasets
+    print()
+    print("=" * 60)
+    print("STEP 2: Preparing datasets")
+    print("=" * 60)
+    
+    # Control: already prepared as data/shakespeare_char/
+    # Modified: prepare from modified text
+    mod_dir = f'data/freq_experiment/modified_{args.target_char}'
+    stoi_mod, itos_mod, vocab_mod = prepare_data(modified_input, mod_dir)
+    
+    # 3. Train both models
+    if not args.skip_train:
+        print()
+        print("=" * 60)
+        print("STEP 3: Training CONTROL model (original Shakespeare)")
+        print("=" * 60)
+        # Control already trained as out-shakespeare-char/
+        
+        print()
+        print("=" * 60)
+        print(f"STEP 4: Training MODIFIED model ({args.target_char} x{args.factor})")
+        print("=" * 60)
+        train_model(mod_dir, f'out-freq-{args.target_char}', args.device)
+    
+    # 4. Run J-lens on both
+    print()
+    print("=" * 60)
+    print("STEP 5: Running J-lens on both models")
+    print("=" * 60)
+    from jlens_v2 import compute_jlens_layer, load_model
+    import numpy as np
+    
+    def run_jlens(ckpt_path, data_dir, label):
+        print(f"\n  J-lens on {label}...")
+        model, config = load_model(ckpt_path, args.device)
+        
+        # Build batches
+        train_data = np.memmap(f'{data_dir}/train.bin', dtype=np.uint16, mode='r')
+        n_batches = 10
+        b_size = 16
+        block_size = config.block_size
+        batches = []
+        for _ in range(n_batches):
+            ix = torch.randint(len(train_data) - block_size, (b_size,))
+            x = torch.stack([torch.from_numpy(
+                train_data[i:i+block_size].astype(np.int64)) for i in ix])
+            y = torch.stack([torch.from_numpy(
+                train_data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
+            batches.append((x, y))
+        
+        # Compute for middle layers only (speed)
+        results = {}
+        for layer_idx in [2, 3, 4]:
+            results[layer_idx] = compute_jlens_layer(model, layer_idx, batches, args.device)
+        
+        return results
+    
+    # Control
+    ctrl_jlens = run_jlens('out-shakespeare-char/ckpt.pt', 
+                           'data/shakespeare_char', 'CONTROL')
+    
+    # Modified
+    mod_jlens = run_jlens(f'out-freq-{args.target_char}/ckpt.pt',
+                          mod_dir, 'MODIFIED')
+    
+    # 5. Compare
+    print()
+    print("=" * 60)
+    print(f"STEP 6: COMPARISON — '{args.target_char}' J-lens norm change")
+    print("=" * 60)
+    
+    with open('data/shakespeare_char/meta.pkl', 'rb') as f:
+        import pickle
+        ctrl_meta = pickle.load(f)
+    ctrl_itos = ctrl_meta['itos']
+    
+    with open(f'{mod_dir}/meta.pkl', 'rb') as f:
+        mod_meta = pickle.load(f)
+    mod_itos = mod_meta['itos']
+    
+    target_id_ctrl = ctrl_itos.index(args.target_char)
+    target_id_mod = mod_itos.index(args.target_char)
+    
+    for layer_idx in sorted(ctrl_jlens.keys()):
+        ctrl_norm = ctrl_jlens[layer_idx][target_id_ctrl].norm().item()
+        mod_norm = mod_jlens[layer_idx][target_id_mod].norm().item()
+        change = (mod_norm - ctrl_norm) / ctrl_norm * 100
+        direction = "↓" if change < 0 else "↑"
+        print(f"  Layer {layer_idx}: control={ctrl_norm:.4f} → modified={mod_norm:.4f} "
+              f"({direction}{abs(change):.1f}%)")
+    
+    print()
+    print("Hypothesis confirmed!" if change < 0 else "Hypothesis NOT confirmed — unexpected result.")
+
+
+if __name__ == '__main__':
+    main()
-- 
cgit v1.2.3