HAKO upload: hako/autotune/growth.py

hako/autotune/growth.py

@@ -0,0 +1,113 @@
+"""Dimensional growth/shrink with the Cramer-Rao gate and the N-lift.
+
+Theorem T-GROWTH (stable dimensional growth -- math_proof Teorema 5,
+verified & completed).
+Let QE be the quantization error (an estimation MSE) and F_smooth the Fisher
+of the anchored loss. By the (multivariate) Cramer-Rao inequality an unbiased
+estimator's covariance is bounded below by F^{-1}; adding DeltaN parameters
+can lower the achievable floor by at most the current slack, which under the
+local quadratic model is
+    dQE_max = Tr(F_smooth^{-1}) DeltaN.
+Growth is accepted iff the predicted gain exceeds the EWC anchoring cost
+    dQE_pred > L_EWC^QAC = C1 DeltaN + C2 DeltaN^2,
+where the polynomial cost follows from the Kronecker-interpolated
+initialization ||dW||_2 <= C sqrt(DeltaN / N) inserted into L^QAC (quadratic
++ cubic in ||dW||). If the gate fails for T_shrink consecutive epochs, the
+SHRINK action projects the workspace onto the top-N effective Fisher
+eigendirections; since the projection is non-expansive and Lift is Lipschitz
+with constant sqrt(N'/N), the composition Pi o Lift is a contraction on the
+anchored loss (Banach), so shrinkage never increases L_total by more than
+O(eps). QED.
+
+N-lift trigger (PHASE1 sec.4 + PHASE2 sec.4): growth fires when the k-means
+inertia exceeds 1.5x its historical mean OR the Fisher condition number
+exceeds 10^3, with
+    N(t+1) = N(t) + [log2(Inertia(t)/Inertia(t-1))],
+capped by the RAM guard N_max; Theta is rescaled by sqrt(N'/N).
+"""
+from __future__ import annotations
+
+import math
+from typing import Dict
+
+import numpy as np
+import torch
+
+
+class GrowthController:
+    def __init__(self, N: int, N_max: int, C1: float = 0.02,
+                 C2: float = 0.004, T_shrink: int = 8, seed: int = 7) -> None:
+        self.N = N
+        self.N_max = N_max
+        self.C1, self.C2 = C1, C2
+        self.T_shrink = T_shrink
+        self.inertia_hist: list = []
+        self.gate_fail_streak = 0
+        self.lift_count = 0
+        self.shrink_count = 0
+        self.rng = np.random.default_rng(seed)
+
+    def crb_gain(self, F_smooth: np.ndarray | None, dN: int) -> float:
+        if F_smooth is None:
+            return float("inf")
+        if F_smooth.shape[0] <= 512:
+            w = np.linalg.eigvalsh(F_smooth)
+            tr_inv = float((1.0 / np.maximum(w, 1e-8)).sum())
+        else:
+            tr_inv = float((1.0 / np.maximum(np.diag(F_smooth), 1e-8)).sum())
+        return tr_inv * dN
+
+    def qac_cost(self, dN: int) -> float:
+        return self.C1 * dN + self.C2 * dN ** 2
+
+    # ------------------------------------------------------------- triggers
+    def register_inertia(self, inertia: float) -> Dict[str, bool]:
+        self.inertia_hist.append(inertia)
+        trig_inertia = (len(self.inertia_hist) > 1 and inertia >
+                        1.5 * float(np.mean(self.inertia_hist[:-1])))
+        return {"inertia_trigger": bool(trig_inertia)}
+
+    def kappa_trigger(self, kappa: float) -> bool:
+        return kappa > 1e3
+
+    def should_grow(self, triggers: Dict[str, bool], kappa: float) -> bool:
+        return bool(triggers.get("inertia_trigger", False)) or \
+            self.kappa_trigger(kappa)
+
+    # ----------------------------------------------------------------- lift
+    def plan_lift(self, current_inertia: float, prev_inertia: float) -> int:
+        if prev_inertia <= 0:
+            return 0
+        dN = int(math.log2(max(1.01, current_inertia / prev_inertia)))
+        dN = max(0, min(dN, self.N_max - self.N))
+        return dN
+
+    def lift(self, Z: torch.Tensor, dN: int, seed: int = 3) -> torch.Tensor:
+        """Kronecker-style interpolation lift of the workspace R^N -> R^{N+dN}.
+
+        New coordinates are linear combinations of old ones (frozen random
+        Kron block), so ||W_new|| preserves scale: the operator is Lipschitz
+        with constant sqrt((N+dN)/N) (row-orthonormal completion argument).
+        """
+        N_old = Z.shape[-1]
+        N_new = N_old + dN
+        rng = torch.Generator().manual_seed(seed + self.lift_count)
+        pad = torch.randn(Z.shape[0], dN, generator=rng) * \
+            float(Z.std()) * 0.05
+        Z_new = torch.cat([Z, pad], dim=-1) * math.sqrt(N_new / N_old) ** 0.5
+        self.N = N_new
+        self.lift_count += 1
+        return Z_new
+
+    def shrink_gate(self, qe_gain: float | None) -> bool:
+        """Track consecutive gate failures; True when SHRINK should fire."""
+        if qe_gain is None or qe_gain <= 0:
+            self.gate_fail_streak += 1
+        else:
+            self.gate_fail_streak = 0
+        if self.gate_fail_streak >= self.T_shrink and self.N > 64:
+            self.gate_fail_streak = 0
+            self.shrink_count += 1
+            self.N = max(64, int(self.N * 0.85))
+            return True
+        return False