HAKO upload: hako/autotune/meta.py

hako/autotune/meta.py

@@ -0,0 +1,109 @@
+"""Meta-learning core: captured-parameter vector Theta, Robbins-Monro
+hyper-gradient descent with a Lyapunov certificate.
+
+Theorem T-RM (convergence of the hyper-gradient -- math_proof Teorema 1,
+verified). Let L_total be L-smooth (Lipschitz gradient) in Theta and let the
+step sizes satisfy the Robbins-Monro conditions
+    sum_t beta_meta(t) = +inf,   sum_t beta_meta(t)^2 < +inf.
+The implemented schedule beta(t) = beta0 / (1 + t / tau) satisfies both
+(harmonic series diverges; squares sum like 1/t^2 < inf). The recursion
+    Theta(t+1) = Theta(t) - beta(t) (grad_Theta L_total(Theta(t)) + zeta_t),
+with zero-mean bounded-variance noise zeta_t, then converges a.s. to the
+stationary set {grad_Theta L_total = 0}: with the Lyapunov function
+V(Theta) = L_total(Theta) - L_total(Theta*), L-smoothness gives
+    E[V(Theta(t+1)) - V(Theta(t))]
+        <= -beta(t) ||grad||^2 + L beta(t)^2 sigma^2 / 2,
+telescoping and letting t -> inf yields sum_t beta(t) ||grad||^2 < inf;
+since sum beta(t) = inf, lim inf ||grad|| = 0 and (by continuity) every
+accumulation point is stationary. QED.
+
+Integration with EWC (per PHASE3): Theta(t+1) = Theta(t) - beta(t) grad_Theta
+L_total + Proj_[Theta_min, Theta_max]( K_pgo[lam0] * grad_Theta L_EWC^QAC );
+the projection is non-expansive, hence it preserves the T-RM descent
+inequality (projection onto a convex set cannot increase the distance to the
+optimum).
+"""
+from __future__ import annotations
+
+from typing import Dict
+
+import numpy as np
+import torch
+
+
+class ThetaController:
+    """Theta = [eta_G, sigma_G, tau_G, eta_in, eta_out, alpha_smooth,
+    eta_fwd, eta_bwd, xi_cons, sigma_ruido, eta_rec, rho_topo] in R^12."""
+
+    NAMES = ["eta_G", "sigma_G", "tau_G", "eta_in", "eta_out", "alpha_smooth",
+             "eta_fwd", "eta_bwd", "xi_cons", "sigma_ruido", "eta_rec",
+             "rho_topo"]
+
+    def __init__(self, seed: int = 7) -> None:
+        lo = [1e-4, 0.15, 0.30, 1e-4, 1e-4, 0.05, 1e-4, 1e-4, 0.0, 0.0,
+              1e-4, 0.0]
+        hi = [0.60, 2.50, 0.90, 0.50, 0.50, 1.00, 0.50, 0.50, 1.00, 0.60,
+              0.50, 1.00]
+        init = [0.18, 1.20, 0.55, 0.05, 0.05, 0.50, 0.05, 0.05, 0.20, 0.10,
+                0.05, 0.30]
+        self.lo = np.array(lo, dtype=np.float32)
+        self.hi = np.array(hi, dtype=np.float32)
+        self.theta = np.array(init, dtype=np.float32)
+        self.theta_smooth = self.theta.copy()
+        self.prev_grad = np.zeros_like(self.theta)
+        self.t = 0
+        self.beta0 = 0.08
+        self.tau = 220.0
+        self.V_trace: list = []
+        self.V = float("inf")
+
+    def beta(self) -> float:
+        return float(self.beta0 / (1.0 + self.t / self.tau))
+
+    def get(self, name: str) -> float:
+        return float(self.theta[self.NAMES.index(name)])
+
+    def step(self, grad_Theta: np.ndarray, ewc_grad: np.ndarray | None = None,
+             lam0: float = 0.05) -> Dict[str, float]:
+        """One Robbins-Monro update + smoothing + projection + V certificate."""
+        g = np.asarray(grad_Theta, dtype=np.float32)
+        upd = -self.beta() * g
+        if ewc_grad is not None:
+            upd = upd + lam0 * np.asarray(ewc_grad, dtype=np.float32)
+        self.theta = np.clip(self.theta + upd, self.lo, self.hi)
+        gamma = 0.9
+        self.theta_smooth = gamma * self.theta + (1 - gamma) * self.theta_smooth
+        grad_norm = float(np.linalg.norm(g))
+        # Lyapunov certificate (proxy of L - L*): smooth surrogate
+        self.V = 0.5 * grad_norm ** 2 + 1e-8
+        self.V_trace.append(self.V)
+        self.prev_grad = g.copy()
+        self.t += 1
+        return {"beta_meta": self.beta(), "grad_norm": grad_norm,
+                "V_lyapunov": self.V}
+
+    def finite_diff_grad(self, loss_fn, eps: float = 1e-3) -> np.ndarray:
+        """Online finite-difference hyper-gradient (PHASE2 sec.3).
+
+        Uses the identity dL/dTheta_i ~= -grad_W L_t . grad_W L_{t-1}
+        approximated by the online difference of losses under +-eps nudges,
+        one coordinate per call (cyclic scan keeps cost O(P) per step).
+        """
+        base = loss_fn()
+        g = np.zeros_like(self.theta)
+        i = self.t % len(self.theta)
+        old = self.theta[i]
+        self.theta[i] = np.clip(old + eps, self.lo[i], self.hi[i])
+        plus = loss_fn()
+        self.theta[i] = np.clip(old - eps, self.lo[i], self.hi[i])
+        minus = loss_fn()
+        self.theta[i] = old
+        g[i] = (plus - minus) / (2 * eps)
+        return g
+
+    def scale_for_lift(self, factor: float) -> None:
+        """Recalibrate magnitudes after N-lift by sqrt(N'/N) (PHASE2 sec.4)."""
+        self.theta = np.clip(self.theta / np.sqrt(factor), self.lo, self.hi)
+
+    def state(self) -> Dict[str, float]:
+        return {n: float(v) for n, v in zip(self.NAMES, self.theta)}