diff --git a/tests/test_hepa.py b/tests/test_hepa.py index 9ec7893..dc4f4f4 100644 --- a/tests/test_hepa.py +++ b/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}" + ) diff --git a/train.py b/train.py index 94ce1f1..2c8fe09 100644 --- a/train.py +++ b/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},