Upload scripts/train_salience_head.py with huggingface_hub

scripts/train_salience_head.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+"""Train the salience head by perturbation of trajectory.
+
+The salience head learns INTRINSICALLY: at each tick with memory active, it
+predicts how much the memory injection will perturb the thought state. The
+signal is ||h_after - h_before||, normalized by a running max. No external
+labels, no REINFORCE — just the dynamical system learning its own sensitivity.
+
+This runs the slow per-tick path (memory is only in `tick`, not `tick_chunk`).
+~25 tok/s on CPU. Budget: 50k tokens (~30 min).
+
+Usage:
+    python scripts/train_salience_head.py [--tokens 50000] [--budget-corpus 200000]
+"""
+import argparse, os, sys, time, math
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+import torch
+import torch.nn.functional as F
+import experiments.edt_ab.ablib as ablib
+from fractus.memory import PersistentMemory
+
+CORPUS = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+                      "data", "communication_corpus.pt")
+SALIENCE_LAMBDA = 0.01  # weight of salience loss vs CE
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--tokens", type=int, default=50_000, help="training tokens")
+    ap.add_argument("--budget-corpus", type=int, default=200_000)
+    ap.add_argument("--lr", type=float, default=3e-4)
+    ap.add_argument("--seed", type=int, default=42)
+    args = ap.parse_args()
+
+    torch.set_num_threads(os.cpu_count() or 6)
+    torch.manual_seed(args.seed)
+
+    print("=== Salience Head Training (perturbation of trajectory) ===", flush=True)
+    split = ablib.load_corpus(CORPUS, n_train=args.budget_corpus,
+                              n_holdout=10_000, n_phase1=int(args.budget_corpus * 0.15))
+    tokens = split["train"][:args.tokens]
+    holdout = split["holdout"]
+
+    # Build engine + attach memory.
+    eng = ablib.build_engine(seed=args.seed)
+    mem = PersistentMemory(d_model=128, max_memories=128)
+    eng.attach_memory(mem)
+    eng.memory_active = True
+
+    # Snapshot salience head weights to detect movement.
+    sal_before = eng.salience_head.weight.detach().clone()
+    sal_bias_before = eng.salience_head.bias.detach().clone()
+
+    opt = torch.optim.AdamW(eng.parameters(), lr=args.lr, weight_decay=0.01)
+    eng.train()
+    eng.reset_thought(batch_size=1)
+
+    t0 = time.time()
+    total_ce, total_sal, n = 0.0, 0.0, 0
+    for t in range(len(tokens) - 1):
+        obs = tokens[t:t + 1]
+        target = tokens[t + 1:t + 2]
+        logits, conf = eng.tick(obs)
+        ce = F.cross_entropy(logits, target)
+        sal_loss = getattr(eng, 'last_salience_loss', torch.tensor(0.0))
+        loss = ce + SALIENCE_LAMBDA * sal_loss
+
+        opt.zero_grad()
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(eng.parameters(), 1.0)
+        opt.step()
+
+        total_ce += ce.item()
+        total_sal += sal_loss.item()
+        n += 1
+
+        if (t + 1) % 2000 == 0:
+            elapsed = time.time() - t0
+            rate = (t + 1) / max(elapsed, 1)
+            print(f"  t={t+1:>6} ce={total_ce/n:.3f} sal={total_sal/n:.4f} "
+                  f"mem={len(mem)} pert_max={eng._pert_max:.3f} "
+                  f"{rate:.0f} tok/s", flush=True)
+
+    elapsed = time.time() - t0
+    # Check if the salience head actually learned.
+    sal_after = eng.salience_head.weight.detach()
+    sal_bias_after = eng.salience_head.bias.detach()
+    weight_delta = (sal_after - sal_before).norm().item()
+    bias_delta = (sal_bias_after - sal_bias_before).norm().item()
+
+    print(f"\n{'='*60}", flush=True)
+    print(f"Training done: {n} tokens in {elapsed/60:.1f}min ({n/elapsed:.0f} tok/s)", flush=True)
+    print(f"Final CE: {total_ce/n:.3f}  Final salience loss: {total_sal/n:.4f}", flush=True)
+    print(f"Memories consolidated: {len(mem)}", flush=True)
+    print(f"Salience head weight delta: {weight_delta:.6f} (should be > 0 if learned)", flush=True)
+    print(f"Salience head bias delta: {bias_delta:.6f}", flush=True)
+    print(f"Perturbation running max: {eng._pert_max:.4f}", flush=True)
+
+    if weight_delta > 1e-6:
+        print("VERDICT: salience head LEARNED (weights moved).", flush=True)
+    else:
+        print("VERDICT: salience head did NOT learn (weights unchanged).", flush=True)
+
+    # Evaluate: does the salience head now predict perturbation?
+    eng.eval()
+    eng.reset_thought(batch_size=1)
+    predicted, actual = [], []
+    with torch.no_grad():
+        for t in range(min(500, len(tokens) - 1)):
+            obs = tokens[t:t + 1]
+            logits, _ = eng.tick(obs)
+            predicted.append(torch.sigmoid(eng.salience_head(
+                eng.thought_state[:, 0, :])).item())
+            actual.append(getattr(eng, '_last_perturbation', 0.0))
+    if max(actual) > 0:
+        corr_n = min(len(predicted), len(actual))
+        pm, am = sum(predicted[:corr_n])/corr_n, sum(actual[:corr_n])/corr_n
+        cov = sum((p-pm)*(a-am) for p,a in zip(predicted[:corr_n], actual[:corr_n]))
+        vp = sum((p-pm)**2 for p in predicted[:corr_n])
+        va = sum((a-am)**2 for a in actual[:corr_n])
+        import math as m
+        denom = m.sqrt(vp * va) if vp > 0 and va > 0 else 0
+        corr = cov / denom if denom > 0 else 0
+        print(f"Correlation(predicted_salience, actual_perturbation) = {corr:.3f}", flush=True)
+    else:
+        print("No perturbations measured during eval — memory may be empty.", flush=True)
+
+
+if __name__ == "__main__":
+    main()