feat(phase1): warm-start joint encoder fine-tuning (Option B)

Two-phase phase-1:
  1a. Frozen warmup: head trains on pre-computed embeddings for PHASE1_EPOCHS=200
  1b. Joint fine-tune: encoder + head for PHASE1_JOINT_EPOCHS=30 at PHASE1_ENCODER_LR=3e-6

Key design decisions:
- Warm start prevents catastrophic forgetting (PHASE1_JOINT=1 cold-start → -32 R²)
- Normalize live encoder output with FROZEN stats (mu_e/sd_e) so head sees same
  embedding distribution it was warmed up on
- head LR reduced 10× in joint phase to prevent head from racing ahead

HPO sweep: 30ep@3e-6=0.3962, 30ep@1e-5=0.3930, 50ep@3e-6=0.3923
Baseline (frozen): 0.3908. New best: phase1_r2=0.3962 (+0.0054 OOS).

New knobs: JEPA_PHASE1_JOINT (default 1), JEPA_PHASE1_JOINT_EPOCHS (default 30),
JEPA_PHASE1_ENCODER_LR (default 3e-6). 4 new tests (tests 15-18). 28/28 pass.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

train.py

@@ -30,8 +30,11 @@ DELTA_T_MAX   = int(_os.environ.get("JEPA_DELTA_T_MAX",   3))
 BATCH_SIZE    = int(_os.environ.get("JEPA_BATCH_SIZE",    512))
 EPOCHS        = int(_os.environ.get("JEPA_EPOCHS",        300))
 LR            = float(_os.environ.get("JEPA_LR",          3e-4))
-PHASE1_EPOCHS = int(_os.environ.get("JEPA_PHASE1_EPOCHS", 200))
-PHASE1_LR     = float(_os.environ.get("JEPA_PHASE1_LR",   1e-3))
+PHASE1_EPOCHS      = int(_os.environ.get("JEPA_PHASE1_EPOCHS",       200))
+PHASE1_LR          = float(_os.environ.get("JEPA_PHASE1_LR",           1e-3))
+PHASE1_JOINT       = bool(int(_os.environ.get("JEPA_PHASE1_JOINT",     1)))
+PHASE1_JOINT_EPOCHS= int(_os.environ.get("JEPA_PHASE1_JOINT_EPOCHS",  30))
+PHASE1_ENCODER_LR  = float(_os.environ.get("JEPA_PHASE1_ENCODER_LR",  3e-6))
 SEED          = int(_os.environ.get("JEPA_SEED",          0))
 # ---------------------------
 
@@ -225,27 +228,61 @@ def main():
     ss_tot  = ((yte - yte.mean()) ** 2).sum()
     val_vol_r2 = float(1 - ss_res / ss_tot)
 
-    # Phase-1: MLP supervised head on frozen embeddings
-    # Standardise targets so the head trains on unit-scale signals.
+    # Phase-1: MLP supervised head — joint or frozen-encoder path
     ytr_mu = float(ytr.mean()); ytr_sd = float(ytr.std()) + 1e-8
     ytr_z  = (ytr - ytr_mu) / ytr_sd
-    head     = SupervisedHead(D_MODEL).to(dev)
+    head   = SupervisedHead(D_MODEL).to(dev)
+    p1_bs  = min(BATCH_SIZE, len(Etr_n))
+
+    # Shared tensors for the frozen-head warmup (used by both paths)
+    Etr_t   = torch.tensor(Etr_n, device=dev)
+    ytr_z_t = torch.tensor(ytr_z, device=dev)
+    Ete_t   = torch.tensor(Ete_n, device=dev)
+    N_tr_h  = len(Etr_t)
+
+    # Phase 1a: warm up head on frozen embeddings (both paths run this)
     head_opt = torch.optim.Adam(head.parameters(), lr=PHASE1_LR, weight_decay=1e-4)
-    Etr_t = torch.tensor(Etr_n, device=dev)
-    ytr_t = torch.tensor(ytr_z, device=dev)
-    Ete_t = torch.tensor(Ete_n, device=dev)
-    p1_bs = min(BATCH_SIZE, len(Etr_t))
-    N_tr_h = len(Etr_t)
-    # Real epoch iteration: shuffle full dataset each epoch
     for _ in range(PHASE1_EPOCHS):
         perm = torch.randperm(N_tr_h, device=dev)
         for start in range(0, N_tr_h, p1_bs):
             idx_h = perm[start:start + p1_bs]
-            loss_h = F.mse_loss(head(Etr_t[idx_h]), ytr_t[idx_h])
+            loss_h = F.mse_loss(head(Etr_t[idx_h]), ytr_z_t[idx_h])
             head_opt.zero_grad(); loss_h.backward(); head_opt.step()
+
+    if PHASE1_JOINT:
+        # Phase 1b: short joint fine-tuning — encoder nudged with tiny LR.
+        # Normalize live encoder output with FROZEN stats (mu_e, sd_e) so the
+        # head sees the same embedding distribution it was warmed up on.
+        enc.train()
+        mu_e_t = torch.tensor(mu_e, device=dev)
+        sd_e_t = torch.tensor(sd_e, device=dev)
+        Xtr_t  = torch.tensor(Xtr, device=dev)
+        joint_opt = torch.optim.Adam([
+            {"params": head.parameters(), "lr": PHASE1_LR * 0.1},
+            {"params": enc.parameters(),  "lr": PHASE1_ENCODER_LR},
+        ], weight_decay=1e-4)
+        for _ in range(PHASE1_JOINT_EPOCHS):
+            perm = torch.randperm(len(Xtr_t), device=dev)
+            for start in range(0, len(Xtr_t), p1_bs):
+                idx_j = perm[start:start + p1_bs]
+                h_raw = enc(Xtr_t[idx_j])[:, -1, :]
+                h_n   = (h_raw - mu_e_t) / sd_e_t  # frozen-stats normalisation
+                loss_j = F.mse_loss(head(h_n), ytr_z_t[idx_j])
+                joint_opt.zero_grad(); loss_j.backward(); joint_opt.step()
+        enc.eval()
+        # Re-extract test embeddings with fine-tuned encoder, same normalisation
+        with torch.no_grad():
+            chunks = []
+            for i in range(0, len(Xte), p1_bs):
+                t = torch.tensor(Xte[i:i+p1_bs], device=dev)
+                h = enc(t)[:, -1, :]
+                chunks.append(((h - mu_e_t) / sd_e_t).cpu().numpy())
+            Ete_t = torch.tensor(np.concatenate(chunks), device=dev)
+
     head.eval()
     with torch.no_grad():
         pred_h_z = head(Ete_t).cpu().numpy()
+
     pred_h    = pred_h_z * ytr_sd + ytr_mu          # de-standardise
     phase1_r2 = float(1 - ((yte - pred_h) ** 2).sum() / ss_tot)
     print("phase1_r2   = %.4f  (n_test=%d)" % (phase1_r2, len(yte)))