feat(phase1): MLP supervised head on frozen HEPA embeddings

SupervisedHead: Linear(D→D/2)→GELU→Linear(D/2→1), trained on standardised
targets with proper epoch iteration (not random 200 batches) + weight_decay=1e-4.
Root cause of earlier -803 R²: unstandardised targets + ~1.2 effective passes.

Results on 2008-2023 hourly OOS (n=11,641):
  val_vol_r2 (linear probe): 0.3585
  phase1_r2  (MLP head):     0.3737  (+0.015 over probe)

New knobs: PHASE1_EPOCHS=200, PHASE1_LR=1e-3. 16/16 tests pass.

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

train.py

@@ -26,10 +26,12 @@ DEPTH       = 2
 N_HEADS     = 4
 ALPHA       = 0.1   # VICReg mixing weight (fixed at 0.1 in HEPA paper)
 DELTA_T_MAX = 3     # max prediction horizon in patches (1..min(DELTA_T_MAX, N-1-c))
-BATCH_SIZE  = 512   # mini-batch per step (hourly dataset is too large for full-batch)
-EPOCHS      = 300
-LR          = 3e-4
-SEED        = 0
+BATCH_SIZE    = 512   # mini-batch per step (hourly dataset is too large for full-batch)
+EPOCHS        = 300
+LR            = 3e-4
+PHASE1_EPOCHS = 200   # supervised head epochs (encoder frozen)
+PHASE1_LR     = 1e-3
+SEED          = 0
 # ---------------------------
 
 torch.manual_seed(SEED)
@@ -119,6 +121,21 @@ class HorizonPredictor(nn.Module):
         return self.net(torch.cat([h, dt], dim=-1))
 
 
+# ── Phase-1 supervised head ──────────────────────────────────────────────────
+
+class SupervisedHead(nn.Module):
+    """Small MLP trained on frozen HEPA embeddings to predict next-period realized vol."""
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(d_model, d_model // 2), nn.GELU(),
+            nn.Linear(d_model // 2, 1),
+        )
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        return self.net(h).squeeze(-1)
+
+
 # ── Data ─────────────────────────────────────────────────────────────────────
 
 def build():
@@ -204,8 +221,33 @@ 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.
+    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_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])
+            head_opt.zero_grad(); loss_h.backward(); head_opt.step()
+    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)))
+
     json.dump({
-        "val_vol_r2": val_vol_r2, "n_test": len(yte),
+        "val_vol_r2": val_vol_r2, "phase1_r2": phase1_r2, "n_test": len(yte),
         "knobs": {"WINDOW": WINDOW, "PATCH_LEN": PATCH_LEN,
                   "D_MODEL": D_MODEL, "DEPTH": DEPTH, "ALPHA": ALPHA,
                   "DELTA_T_MAX": DELTA_T_MAX, "EPOCHS": EPOCHS},