Phase-Aware Bias-Energy Trust Region / phase_aware.py

✓✓ Beats tuned baseline

Raw ⬇ ZIP
  1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def clip_norm(x, tau):
  8    n = np.linalg.norm(x)
  9    return x if n <= tau else x * (tau / n)
 10
 11
 12def k_formula(alpha, beta, p):
 13    if alpha <= p * beta:
 14        return beta
 15    return (alpha / p) * (alpha * (p - 1) / (p * (alpha - beta))) ** (p - 1)
 16
 17
 18def r_star_formula(alpha, beta, p):
 19    if alpha <= p * beta:
 20        return 1.0
 21    return p * (alpha - beta) / (alpha * (p - 1))
 22
 23
 24def normalized_cost(r, alpha, beta):
 25    return beta * r * r if r <= 1 else alpha * (r - 1) + beta
 26
 27
 28def verify_math():
 29    rows = []
 30    # Dense radial sweep is the exact one-dimensional reduction in the theorem.
 31    rs = np.unique(np.r_[np.linspace(0, 1, 20001), np.geomspace(1.000001, 1e4, 50000)])
 32    for p in (1.25, 1.5, 2.0):
 33        for alpha, beta in ((0.8, 1.0), (2.0, 1.0), (1.0, 1.0), (3.0, .5), (0.25, .5)):
 34            vals = np.array([normalized_cost(r, alpha, beta) / max(r ** p, 1e-300) for r in rs])
 35            j = int(np.argmax(vals))
 36            pred_r = r_star_formula(alpha, beta, p)
 37            # At p=2 in the energy phase the theorem predicts a whole tie interval r<=1.
 38            tie_valid = (p == 2.0 and alpha <= p*beta and rs[j] <= 1.0 + 1e-8)
 39            radius_error = 0.0 if tie_valid else abs(float(rs[j])-pred_r)
 40            rows.append({"p":p, "alpha":alpha, "beta":beta,
 41                         "predicted_phase":"energy" if alpha <= p*beta else "bias",
 42                         "observed_phase":"energy" if rs[j] <= 1.0001 else "bias",
 43                         "boundary_margin":alpha-p*beta,
 44                         "predicted_r":pred_r, "observed_r":float(rs[j]),
 45                         "radius_error":radius_error, "p2_tie_valid":bool(tie_valid),
 46                         "predicted_K":k_formula(alpha,beta,p), "observed_K":float(vals[j])})
 47
 48    # Explicit transition sweep: the maximizer should remain at r=1 through alpha=p beta,
 49    # then move continuously outside with the exact r_star formula.
 50    transition = []
 51    p0, beta0 = 1.5, 1.0
 52    for alpha0 in (1.40, 1.49, 1.50, 1.51, 1.60, 2.00):
 53        vv = np.array([normalized_cost(r, alpha0, beta0) / max(r**p0, 1e-300) for r in rs])
 54        jj = int(np.argmax(vv))
 55        transition.append({"alpha":alpha0, "predicted_r":r_star_formula(alpha0,beta0,p0),
 56                           "observed_r":float(rs[jj]), "predicted_phase":"energy" if alpha0 <= p0*beta0 else "bias",
 57                           "observed_K":float(vv[jj]), "predicted_K":k_formula(alpha0,beta0,p0)})
 58
 59    # Pure residual: max (r-1)_+ / r^p = c_p at p/(p-1), and scaling is tau-independent.
 60    residual_rows = []
 61    for p in (1.25, 1.5, 2.0):
 62        rs2 = np.geomspace(1.0000001, 1e5, 100000)
 63        q = (rs2 - 1) / rs2**p
 64        j = int(np.argmax(q))
 65        cp = (p-1)**(p-1) / p**p
 66        residual_rows.append({"p":p, "predicted_r":p/(p-1), "observed_r":float(rs2[j]),
 67                              "predicted_c":cp, "observed_c":float(q[j])})
 68    # Direct vector scaling check over several tau and dimensions.
 69    scaling = []
 70    rng = np.random.default_rng(7)
 71    for p in (1.5, 2.0):
 72        cp = (p-1)**(p-1) / p**p
 73        for tau in (.2, 1., 5.):
 74            r = p/(p-1)
 75            x = np.array([r*tau, 0.0, 0.0])
 76            resid = np.linalg.norm(x-clip_norm(x,tau))
 77            ratio = resid / (tau**(1-p) * np.linalg.norm(x)**p)
 78            scaling.append({"p":p,"tau":tau,"ratio":float(ratio),"predicted":cp})
 79    return rows, residual_rows, scaling, transition
 80
 81
 82class PhaseAwareController:
 83    def __init__(self, p=2.0, alpha=1.0, beta=1.0, tau=1.0, eta=.035):
 84        self.p, self.alpha, self.beta = p, alpha, beta
 85        self.tau, self.eta = tau, eta
 86        self.phase = "energy" if alpha <= p*beta else "bias"
 87        self.history = []
 88
 89    def step(self, gradients):
 90        norms = np.linalg.norm(gradients, axis=1)
 91        clipped = gradients.copy()
 92        factors = np.minimum(1., self.tau / np.maximum(norms, 1e-12))
 93        clipped *= factors[:, None]
 94        residual = np.mean(np.linalg.norm(gradients-clipped, axis=1))
 95        energy = np.mean(np.sum(clipped*clipped, axis=1)) / self.tau
 96        # Targets are deliberately fixed relative to the initial trust region.
 97        target = .22 if self.phase == "energy" else .10
 98        signal = energy if self.phase == "energy" else residual
 99        delta = np.clip(self.eta * (signal-target), -.08, .08)
100        self.tau *= math.exp(delta)
101        self.tau = float(np.clip(self.tau, .08, 8.0))
102        self.history.append((self.tau, residual, energy, np.mean(factors < 1)))
103        return clipped
104
105
106def percentile_clip(g, tau_state, percentile=90):
107    norms = np.linalg.norm(g, axis=1)
108    threshold = np.percentile(norms, percentile)
109    tau = min(tau_state, threshold)
110    factors = np.minimum(1., tau / np.maximum(norms, 1e-12))
111    return g * factors[:,None], tau
112
113
114def mini_experiment(seed=11, steps=700, batch=64, dim=8):
115    rng = np.random.default_rng(seed)
116    methods = ["fixed", "percentile", "phase_aware"]
117    results = {}
118    for method in methods:
119        w = np.ones(dim) * 3.0
120        tau, taus = 1.0, []
121        ctl = PhaseAwareController(alpha=1., beta=1., tau=1.) if method == "phase_aware" else None
122        losses, updates, clips, residuals, energies = [], [], [], [], []
123        for t in range(steps):
124            # Quadratic objective with occasional Pareto-like, direction-random outlier batches.
125            g = w[None,:] + rng.normal(0, .18, (batch,dim))
126            if t % 70 == 35:
127                g += rng.normal(size=(batch,dim)) * 12.0
128            if method == "fixed":
129                factors = np.minimum(1., tau / np.maximum(np.linalg.norm(g,axis=1),1e-12))
130                cg = g * factors[:,None]
131            elif method == "percentile":
132                cg, tau = percentile_clip(g, tau)
133                factors = np.minimum(1., tau / np.maximum(np.linalg.norm(g,axis=1),1e-12))
134            else:
135                old_tau = ctl.tau
136                cg = ctl.step(g)
137                tau = ctl.tau
138                factors = np.minimum(1., old_tau / np.maximum(np.linalg.norm(g,axis=1),1e-12))
139            update = .045 * np.mean(cg, axis=0)
140            w -= update
141            losses.append(.5*np.dot(w,w)); updates.append(np.linalg.norm(update))
142            clips.append(np.mean(factors < 1)); residuals.append(np.mean(np.linalg.norm(g-cg,axis=1)))
143            energies.append(np.mean(np.sum(cg*cg,axis=1))/max(tau,1e-9)); taus.append(tau)
144        # recovery is steps needed after final outlier to get below .1 loss, or sentinel
145        tail = np.array(losses[596:])
146        hit = np.where(tail < .1)[0]
147        results[method] = {"final_loss":float(losses[-1]), "median_update":float(np.median(updates)),
148          "update_std":float(np.std(updates)), "clip_fraction":float(np.mean(clips)),
149          "mean_residual":float(np.mean(residuals)), "mean_energy":float(np.mean(energies)),
150          "recovery_steps_after_last_outlier":int(hit[0]) if len(hit) else None,
151          "final_tau":float(taus[-1])}
152    return results
153
154
155def main():
156    math_rows, residual_rows, scaling, transition = verify_math()
157    mini = mini_experiment()
158    out = {"math_sweep":math_rows,"transition_sweep":transition,"residual_sweep":residual_rows,"scaling_check":scaling,"mini_experiment":mini}
159    Path("results.json").write_text(json.dumps(out, indent=2))
160    print(json.dumps(out, indent=2))
161
162if __name__ == "__main__":
163    main()