1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 2206
  6np.random.seed(SEED)
  7random.seed(SEED)
  8
  9
 10def toy_sweep():
 11    # Exact quadratic value landscape. A step gamma moves from x=0 to x=gamma.
 12    # V(gamma)=a*gamma - b*gamma^2/2; KL(N(gamma,s2)||N(0,s2))=gamma^2/(2*s2).
 13    # Predictions: KL scales as gamma^2, Delta-V as gamma for small gamma,
 14    # eta=DeltaV/KL scales as gamma^-1, and eta is constant negative for nuisance.
 15    a, b, s2 = 1.7, 2.0, 0.8
 16    gammas = np.logspace(-3, math.log10(3.0), 48)
 17    kl = gammas**2 / (2.0 * s2)
 18    dv = a * gammas - 0.5 * b * gammas**2
 19    eta = dv / kl
 20    nuisance_dv = -0.5 * b * gammas**2
 21    nuisance_eta = nuisance_dv / kl
 22
 23    def slope(x, y, n=10):
 24        return float(np.polyfit(np.log(x[:n]), np.log(np.abs(y[:n]) + 1e-30), 1)[0])
 25
 26    # Exact zero crossing is gamma=2a/b. The grid estimate brackets it.
 27    changes = np.where(np.sign(dv[:-1]) != np.sign(dv[1:]))[0]
 28    bracket = None if len(changes) == 0 else [float(gammas[changes[0]]), float(gammas[changes[0] + 1])]
 29    return {
 30        'predictions': {
 31            'KL_loglog_slope': 2.0,
 32            'small_step_value_loglog_slope': 1.0,
 33            'small_step_eta_loglog_slope': -1.0,
 34            'aligned_value_zero_crossing': 2.0 * a / b,
 35            'nuisance_eta': -b * s2,
 36        },
 37        'observed': {
 38            'KL_loglog_slope': slope(gammas, kl),
 39            'small_step_value_loglog_slope': slope(gammas, dv),
 40            'small_step_eta_loglog_slope': slope(gammas, eta),
 41            'aligned_value_zero_crossing_bracket': bracket,
 42            'nuisance_eta_mean': float(np.mean(nuisance_eta)),
 43            'nuisance_eta_std': float(np.std(nuisance_eta)),
 44        },
 45    }
 46
 47
 48def neural_experiment():
 49    # A tiny linear predictor with useful x and nuisance z. The latter is
 50    # deliberately high dimensional and random, so fitting it acquires KL but
 51    # cannot improve held-out value. The monitor freezes when efficiency is low.
 52    import torch
 53    torch.manual_seed(SEED)
 54    try:
 55        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 56        torch.zeros(1, device=device)
 57    except Exception:
 58        device = torch.device('cpu')
 59
 60    n_train, n_test, p = 512, 1024, 24
 61    rng = np.random.default_rng(SEED)
 62    xtr = rng.normal(size=(n_train, 1)).astype('float32')
 63    xte = rng.normal(size=(n_test, 1)).astype('float32')
 64    ztr = rng.normal(size=(n_train, p)).astype('float32')
 65    zte = rng.normal(size=(n_test, p)).astype('float32')
 66    ytr = (2.0 * xtr[:, 0] + 0.4 * rng.normal(size=n_train)).astype('float32')
 67    yte = (2.0 * xte[:, 0] + 0.4 * rng.normal(size=n_test)).astype('float32')
 68    Xtr = torch.tensor(np.concatenate([xtr, ztr], 1), device=device)
 69    Xte = torch.tensor(np.concatenate([xte, zte], 1), device=device)
 70    Ytr, Yte = torch.tensor(ytr, device=device), torch.tensor(yte, device=device)
 71
 72    def run(monitored):
 73        w = torch.zeros(p + 1, device=device, requires_grad=True)
 74        opt = torch.optim.Adam([w], lr=0.08)
 75        prev = w.detach().clone()
 76        prev_val = None
 77        updates, infos, etas = 0, [], []
 78        for t in range(180):
 79            idx = torch.arange((t * 32) % n_train, (t * 32) % n_train + 32, device=device) % n_train
 80            loss = ((Xtr[idx] @ w - Ytr[idx]) ** 2).mean()
 81            opt.zero_grad(); loss.backward(); opt.step()
 82            new = w.detach().clone()
 83            info = float(((new - prev) ** 2).sum().item() / (2 * 0.25))
 84            val = float(((Xte @ new - Yte) ** 2).mean().item())
 85            # Value is negative test MSE; a block is stopped after repeated
 86            # low/negative value-per-KL updates, while preserving equal steps.
 87            dv = 0.0 if prev_val is None else prev_val - val
 88            eta = dv / (info + 1e-12)
 89            infos.append(info); etas.append(eta)
 90            updates += 1
 91            if monitored and t > 12 and len(etas) >= 5 and np.mean(etas[-5:]) < 0.02:
 92                # trust-region response: undo the low-efficiency update
 93                with torch.no_grad(): w.copy_(prev)
 94            else:
 95                prev = new
 96                prev_val = val
 97        final_mse = float(((Xte @ w - Yte) ** 2).mean().item())
 98        return {'test_mse': final_mse, 'updates': updates,
 99                'total_kl': float(np.sum(infos)),
100                'late_eta': float(np.mean(etas[-30:])),
101                'device': str(device)}
102
103    return {'baseline_adam': run(False), 'monitor_adam': run(True)}
104
105
106def main():
107    out = {'toy': toy_sweep(), 'neural': neural_experiment()}
108    Path('results.json').write_text(json.dumps(out, indent=2))
109    print(json.dumps(out, indent=2))
110
111
112if __name__ == '__main__':
113    main()