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>

tests/test_hepa.py

@@ -1,7 +1,7 @@
-"""Failing tests for HEPA backbone in train.py.
+"""Failing tests for HEPA backbone + Phase-1 supervised head in train.py.
 
 Run: cd ~/dev/AI/jepa-fx-risk && .venv/bin/python -m pytest tests/test_hepa.py -v
-These tests define what the new backbone must satisfy BEFORE implementation.
+These tests define what the backbone and head must satisfy BEFORE implementation.
 """
 import math
 import torch
@@ -110,3 +110,74 @@ def test_build_hourly_more_windows(train_mod):
     (Xtr, _), _ = train_mod.build()
     # Daily had ~877 train windows; hourly with 2008-2021 should have > 50,000
     assert len(Xtr) > 10_000, f"expected >10k hourly train windows, got {len(Xtr)}"
+
+
+# ── Phase-1: supervised head ──────────────────────────────────────────────────
+
+# 8. SupervisedHead exists and maps (B, D) → (B,)
+def test_supervised_head_shape(train_mod):
+    D = 128
+    head = train_mod.SupervisedHead(D)
+    x = torch.randn(16, D)
+    out = head(x)
+    assert out.shape == (16,), f"expected (16,), got {out.shape}"
+
+
+# 9. SupervisedHead gradient flows (not frozen)
+def test_supervised_head_backward(train_mod):
+    head = train_mod.SupervisedHead(64)
+    x = torch.randn(8, 64)
+    loss = head(x).mean()
+    loss.backward()
+    for name, p in head.named_parameters():
+        assert p.grad is not None, f"no grad on {name}"
+
+
+# 10. Phase-1 beats linear on nonlinear synthetic signal
+def test_phase1_beats_linear_on_nonlinear(train_mod):
+    """MLP head should outperform ridge regression on data with nonlinear structure."""
+    import numpy as np
+    torch.manual_seed(0); np.random.seed(0)
+    N, D = 1000, 32
+    # target = |h|² (quadratic — linear can't fit well)
+    Etr = np.random.randn(N, D).astype(np.float32)
+    ytr = (Etr ** 2).sum(axis=1)
+    Ete = np.random.randn(200, D).astype(np.float32)
+    yte = (Ete ** 2).sum(axis=1)
+
+    # Ridge baseline
+    A = np.hstack([Etr, np.ones((N, 1))])
+    w = np.linalg.solve(A.T @ A + 1e-3 * np.eye(A.shape[1]), A.T @ ytr)
+    pred_lin = np.hstack([Ete, np.ones((200, 1))]) @ w
+    r2_lin = float(1 - ((yte - pred_lin) ** 2).sum() / ((yte - yte.mean()) ** 2).sum())
+
+    # MLP head
+    head = train_mod.SupervisedHead(D)
+    opt  = torch.optim.Adam(head.parameters(), lr=1e-2)
+    Xtr_t = torch.tensor(Etr); ytr_t = torch.tensor(ytr)
+    for _ in range(300):
+        loss = nn.functional.mse_loss(head(Xtr_t), ytr_t)
+        opt.zero_grad(); loss.backward(); opt.step()
+
+    head.eval()
+    with torch.no_grad():
+        pred_mlp = head(torch.tensor(Ete)).numpy()
+    r2_mlp = float(1 - ((yte - pred_mlp) ** 2).sum() / ((yte - yte.mean()) ** 2).sum())
+
+    assert r2_mlp > r2_lin + 0.05, (
+        f"MLP R²={r2_mlp:.3f} should beat ridge R²={r2_lin:.3f} by >0.05 on quadratic target"
+    )
+
+
+# 11. main() returns phase1_r2 in metrics.json (integration — needs real data)
+def test_metrics_json_has_phase1_r2(train_mod):
+    import os, json
+    if not os.path.exists("metrics.json"):
+        pytest.skip("metrics.json not present — run train.py first")
+    with open("metrics.json") as f:
+        m = json.load(f)
+    assert "phase1_r2" in m, f"phase1_r2 missing from metrics.json: {list(m.keys())}"
+    assert m["phase1_r2"] > m["val_vol_r2"], (
+        f"MLP head phase1_r2={m['phase1_r2']:.4f} should beat linear probe "
+        f"val_vol_r2={m['val_vol_r2']:.4f}"
+    )

train.py

@@ -26,9 +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))
-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)
@@ -118,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():
@@ -157,29 +175,37 @@ def main():
     (Xtr, ytr), (Xte, yte) = build()
     n_feats   = Xtr.shape[2]
     n_patches = WINDOW // PATCH_LEN
-    Xtr_t     = torch.tensor(Xtr, device=dev)
+    N_tr      = len(Xtr)
+    bs        = min(BATCH_SIZE, N_tr)
 
     enc  = CausalEncoder(n_feats, PATCH_LEN, D_MODEL, N_HEADS, DEPTH).to(dev)
     pred = HorizonPredictor(D_MODEL).to(dev)
     opt  = torch.optim.AdamW(list(enc.parameters()) + list(pred.parameters()), lr=LR)
 
     for ep in range(EPOCHS):
-        # Sample random context position and horizon; Δt log-biased toward short
+        # Random mini-batch (avoids OOM on large hourly dataset)
+        idx_b = torch.randperm(N_tr)[:bs]
+        Xb    = torch.tensor(Xtr[idx_b.numpy()], device=dev)
+
+        # Sample random context position and horizon
         c  = torch.randint(0, n_patches - 1, ()).item()
         dt = torch.randint(1, max(2, min(DELTA_T_MAX, n_patches - 1 - c) + 1), ()).item()
 
-        tokens = enc(Xtr_t)                     # (B, N, D)
+        tokens = enc(Xb)                        # (bs, N, D)
         h_ctx  = tokens[:, c, :]               # context embedding
         h_tgt  = tokens[:, c + dt, :]          # target embedding (joint training)
-        h_hat  = pred(h_ctx, torch.full((len(Xtr),), float(dt), device=dev))
+        h_hat  = pred(h_ctx, torch.full((bs,), float(dt), device=dev))
         loss   = vicreg_loss(h_hat, h_tgt, alpha=ALPHA)
         opt.zero_grad(); loss.backward(); opt.step()
 
     enc.eval()
     with torch.no_grad():
         def embed(X_np):
-            t = torch.tensor(X_np, device=dev)
-            return enc(t)[:, -1, :].cpu().numpy()   # last token = full-context summary
+            chunks = []
+            for i in range(0, len(X_np), bs):
+                t = torch.tensor(X_np[i:i+bs], device=dev)
+                chunks.append(enc(t)[:, -1, :].cpu().numpy())
+            return np.concatenate(chunks, axis=0)
 
         Etr = embed(Xtr)
         Ete = embed(Xte)
@@ -195,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},
@@ -223,14 +274,18 @@ def main():
             idx = df2.index[year_mask].tolist()
             Xs, dates, rvs = [], [], []
             for t in idx:
-                if t - WINDOW >= 0:
+                if t - WINDOW >= 0 and t + 1 < len(df2):
                     Xs.append(fn2[t - WINDOW:t])
                     dates.append(str(df2["date"].iloc[t].date()))
-                    rvs.append(float(df2["realized_vol"].iloc[t]))
+                    rvs.append(float(df2["realized_vol"].iloc[t + 1]))
             if not Xs:
                 return [], [], []
+            Xa = np.stack(Xs)
+            chunks = []
             with torch.no_grad():
-                E = enc(torch.tensor(np.stack(Xs), device=dev))[:, -1, :].cpu().numpy().tolist()
+                for i in range(0, len(Xa), bs):
+                    chunks.append(enc(torch.tensor(Xa[i:i+bs], device=dev))[:, -1, :].cpu().numpy())
+            E = np.concatenate(chunks, axis=0).tolist()
             return E, dates, rvs
         Etr2, dates_tr, rv_tr = _export_windows(tr_mask)
         Eoos, dates_oos, rv_oos = _export_windows(df2["date"].dt.year >= 2022)