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>
2026-06-26 12:33:28 +02:00
parent 1a17a4c88e
commit dd1004fffc
2 changed files with 120 additions and 7 deletions
@@ -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}"
+    )
@@ -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},