1import json, math, random
  2import numpy as np
  3import torch
  4
  5SEED = 1737
  6random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  7torch.set_default_dtype(torch.float64)
  8
  9
 10def exact_and_hutchinson(theta, A, s, direction, beta=0.0, rmin=1e-3, probes=256):
 11    """Compute D_w for rho=rmin+softplus(s+direction dot theta), and Hutchinson div(G)."""
 12    theta = theta.detach().requires_grad_(True)
 13    s = s.detach().requires_grad_(True)
 14    L = 0.5 * theta @ A @ theta
 15    g = torch.autograd.grad(L, theta, create_graph=True)[0]
 16    rho = rmin + torch.nn.functional.softplus(s + direction @ theta)
 17    G = -rho * g
 18    div = 0.0
 19    for i in range(theta.numel()):
 20        div = div + torch.autograd.grad(G[i], theta, retain_graph=True, create_graph=True)[0][i]
 21    D = div + beta * rho * (g @ g)
 22    vals = []
 23    for _ in range(probes):
 24        z = torch.randint(0, 2, theta.shape, dtype=theta.dtype) * 2 - 1
 25        jgz = torch.autograd.grad((G * z).sum(), theta, retain_graph=True)[0]
 26        vals.append((z * jgz).sum().item())
 27    return float(D.detach()), float(div.detach()), float(rho.detach()), np.array(vals)
 28
 29
 30def verify():
 31    A = torch.diag(torch.tensor([1.0, 7.0, 30.0]))
 32    theta = torch.tensor([0.7, -0.4, 0.2])
 33    s = torch.tensor(-0.3)
 34    direction = torch.tensor([0.13, -0.21, 0.17])
 35    D, div, rho, h = exact_and_hutchinson(theta, A, s, direction, beta=.08, probes=512)
 36    # Explicit formula: div(-rho*g)=-rho tr(A)-grad(rho).g.
 37    with torch.no_grad():
 38        u = s + direction @ theta
 39        drho_dot_g = torch.sigmoid(u) * (direction @ (A @ theta))
 40        explicit_div = -rho * torch.trace(A).item() - drho_dot_g.item()
 41        explicit_D = explicit_div + .08 * rho * (theta @ A @ A @ theta).item()
 42    return {
 43        "exact_D": D, "formula_D": explicit_D, "abs_D_error": abs(D-explicit_D),
 44        "exact_div": div, "formula_div": explicit_div,
 45        "hutch_mean": float(h.mean()), "hutch_std_error": float(h.std(ddof=1)/math.sqrt(len(h))),
 46        "rho": rho
 47    }
 48
 49
 50def run_optimizer(adaptive, seed, steps=180):
 51    torch.manual_seed(seed); np.random.seed(seed)
 52    A = torch.diag(torch.tensor([1.0, 15.0]))
 53    theta = torch.tensor([2.0, 2.0], dtype=torch.float64)
 54    # A fixed directional dependence makes grad rho nonzero, while s is the controller.
 55    direction = torch.tensor([0.25, -0.20])
 56    s = torch.tensor(-2.0, requires_grad=True)
 57    s_lr = .08
 58    fixed_rho = 0.055
 59    losses=[]; rhos=[]; contractions=[]; updates=[]
 60    for t in range(steps):
 61        theta.requires_grad_(True)
 62        L = .5 * theta @ A @ theta
 63        g = torch.autograd.grad(L, theta, create_graph=True)[0]
 64        if adaptive:
 65            rho = .002 + torch.nn.functional.softplus(s + direction @ theta)
 66            G = -rho * g
 67            div = sum(torch.autograd.grad(G[i], theta, retain_graph=True, create_graph=True)[0][i] for i in range(2))
 68            D = div  # beta=0: weighted divergence is ordinary divergence
 69            # Controller minimizes a soft contraction-margin violation, but does not alter theta gradient.
 70            penalty = torch.nn.functional.softplus(D + 0.35)
 71            ds = torch.autograd.grad(penalty, s, allow_unused=True)[0]
 72            if ds is not None:
 73                with torch.no_grad(): s -= s_lr * ds.clamp(-2,2)
 74            step = (-rho.detach() * g.detach())
 75            rho_value=float(rho.detach())
 76            D_value=float(D.detach())
 77        else:
 78            step = -fixed_rho * g.detach()
 79            rho_value=fixed_rho
 80            D_value=-fixed_rho * torch.trace(A).item()
 81        theta = (theta.detach() + step)
 82        losses.append(float(L.detach())); rhos.append(rho_value); contractions.append(D_value); updates.append(float(step.norm()))
 83    return {"final_loss": losses[-1], "initial_loss": losses[0], "losses": losses,
 84            "mean_rho":float(np.mean(rhos)), "rho_last":rhos[-1],
 85            "mean_D":float(np.mean(contractions)), "max_loss":float(max(losses)),
 86            "mean_update":float(np.mean(updates)), "final_theta":theta.tolist()}
 87
 88
 89def main():
 90    check=verify()
 91    base=[run_optimizer(False, 100+i) for i in range(8)]
 92    idea=[run_optimizer(True, 100+i) for i in range(8)]
 93    def summary(xs):
 94        return {"final_loss_mean":float(np.mean([x['final_loss'] for x in xs])),
 95                "final_loss_std":float(np.std([x['final_loss'] for x in xs], ddof=1)),
 96                "max_loss_mean":float(np.mean([x['max_loss'] for x in xs])),
 97                "mean_D":float(np.mean([x['mean_D'] for x in xs])),
 98                "mean_rho":float(np.mean([x['mean_rho'] for x in xs]))}
 99    out={"verification":check, "baseline":summary(base), "idea":summary(idea),
100         "settings":{"steps":180,"seeds":8,"A":[1,15],"beta":0,"baseline_rho":.055}}
101    with open("results.json","w") as f: json.dump(out,f,indent=2)
102    print(json.dumps(out,indent=2))
103
104if __name__ == '__main__': main()