HAKO upload: hako/core/heads.py

hako/core/heads.py

@@ -0,0 +1,156 @@
+"""CPNN / BKN / DASOM heads coupled to GHSOM neurons (L3 plugins).
+
+Theorem T-HEADS (convergence with stationary inputs -- frozen sources).
+The frozen ONNX sources give z with d z / d W_src = 0 (LLM01 Phase-1
+exhaustion lemma). The per-neuron update combines Kohonen attraction and
+task gradients,
+    Delta W_j = eta_G h_cj (z - W_j)
+              + eta_D h_cj dL_Denoise/dW_j + eta_B h_cj dL_Bidir/dW_j,
+and analogously for V (CPNN supervised regression) and U (BKN/DASOM).
+With step sizes eta_t -> 0, sum eta_t = inf, sum eta_t^2 < inf (Robbins-Monro
+schedule implemented by the trainer), the recursion converges a.s. to the
+stationary set; for the Kohonen component the stationary points are the
+Voronoi centroids of the input distribution restricted to each winner cell
+(classical SOM result: at stationarity E[ h_cj (z - W_j) ] = 0 forces W_j to
+the local conditional mean). The head layers V, U converge under plain SGD on
+their convex quadratic losses (least squares with fixed features), whose
+objective is strongly convex in (V, U) given bounded activations.
+QED.
+
+Outputs per forward (arch.pdf 2.1):
+    CPNN : y_hat = sum_j h_cj V_j z                (smoothed regression)
+    BKN  : z_rec = sum_j h_cj U^B_j z              (bidirectional reconstruction)
+    DASOM: z_den = sum_j h_cj U^D_j x_noisy        (denoising reconstruction)
+"""
+from __future__ import annotations
+
+from typing import Dict, Iterator, Tuple
+
+import torch
+import torch.nn as nn
+
+
+def neuron_key(nid: int, idx: int) -> str:
+    return f"{nid}:{idx}"
+
+
+class HeadsModule(nn.Module):
+    """Lazy per-neuron head tensors shared by all ensemble replicas.
+
+    Memory bound (aggressive management): V is out_dim x N, U^B and U^D are
+    diagonal + rank-r factorizations r=16, so per neuron we store
+    O(out*N + 2*(N*r + N)) floats, NOT 2*N^2. For N=256, r=16 this is ~50 KiB
+    per neuron, allowing hundreds of neurons within the RAM cap.
+    """
+
+    def __init__(self, N_dim: int, out_dim: int = 64, rank: int = 16,
+                 n_replicas: int = 2, seed: int = 7) -> None:
+        super().__init__()
+        self.N = N_dim
+        self.out_dim = out_dim
+        self.rank = rank
+        self.K = n_replicas
+        g = torch.Generator().manual_seed(seed)
+        self._gen = g
+        self.V: Dict[str, torch.nn.Parameter] = {}
+        self.Udiag: Dict[Tuple[str, int], torch.nn.Parameter] = {}
+        self.Ua: Dict[Tuple[str, int], torch.nn.Parameter] = {}
+        self.Ub: Dict[Tuple[str, int], torch.nn.Parameter] = {}
+
+    # ------------------------------------------------------------- creation
+    def attach(self, key: str) -> None:
+        if key in self.V:
+            return
+        self.V[key] = nn.Parameter(0.02 * torch.randn(self.out_dim, self.N,
+                                                      generator=self._gen))
+        for rep in range(self.K):
+            self.Udiag[(key, rep)] = nn.Parameter(torch.ones(self.N))
+            self.Ua[(key, rep)] = nn.Parameter(
+                0.02 * torch.randn(self.N, self.rank, generator=self._gen))
+            self.Ub[(key, rep)] = nn.Parameter(
+                0.02 * torch.randn(self.rank, self.N, generator=self._gen))
+
+    def attach_parent_mean(self, key: str, W_parent_mean: torch.Tensor) -> None:
+        """New neurons after growth: V initialized from parent-cell projection."""
+        self.attach(key)
+        with torch.no_grad():
+            self.V[key].zero_()
+            # project parent mean onto the dominant output coordinates
+            self.V[key][:, : min(self.N, self.out_dim)] = \
+                W_parent_mean[: min(self.N, self.out_dim), : min(self.N, self.out_dim)].t()
+
+    def parameters(self, recurse: bool = True) -> Iterator[torch.nn.Parameter]:
+        for p in self.V.values():
+            yield p
+        for d in (self.Udiag, self.Ua, self.Ub):
+            for p in d.values():
+                yield p
+
+    # -------------------------------------------------------------- forward
+    def _U(self, key: str, rep: int) -> torch.Tensor:
+        d = self.Udiag[(key, rep)]
+        a = self.Ua[(key, rep)]
+        b = self.Ub[(key, rep)]
+        return torch.diag(d) + a @ b
+
+    def forward(self, z: torch.Tensor, path, h_map: Dict[Tuple[int, int], torch.Tensor],
+                rep: int = 0) -> Dict[str, torch.Tensor]:
+        """z: (N,), path: [(nid, idx)...], h_map: per-node neighborhood matrix.
+
+        h_map[nid] is (rows, cols) SOM neighborhood of the winner at that node.
+        """
+        H_cat: Dict[str, torch.Tensor] = {}
+        y_hat = torch.zeros(self.out_dim, dtype=z.dtype)
+        z_rec = torch.zeros(self.N, dtype=z.dtype)
+        z_den = torch.zeros(self.N, dtype=z.dtype)
+        wsum = torch.zeros(())
+        for (nid, idx) in path:
+            key = neuron_key(nid, idx)
+            self.attach(key)
+            node_h = h_map.get(nid)
+            if node_h is None:
+                w = torch.ones(())
+            else:
+                w = node_h.flatten()[idx]
+            wsum = wsum + w
+            y_hat = y_hat + w * (self.V[key] @ z)
+            z_rec = z_rec + w * (self._U(key, rep) @ z)
+            z_den = z_den + w * (self._U(key, min(rep + 1, self.K - 1)) @ z)
+            H_cat[key] = z
+        if float(wsum) > 0:
+            y_hat, z_rec, z_den = y_hat / wsum, z_rec / wsum, z_den / wsum
+        return {"y_hat": y_hat, "z_rec": z_rec, "z_den": z_den,
+                "H_cat": H_cat, "wsum": wsum}
+
+    # ---------------------------------------------------------- checkpoint
+    def state_numpy(self) -> Dict[str, "object"]:
+        import numpy as np
+        out = {}
+        for k, p in self.V.items():
+            out[f"V[{k}]"] = p.detach().numpy()
+        for (k, rep), p in self.Udiag.items():
+            out[f"Udiag[{k}|{rep}]"] = p.detach().numpy()
+        for (k, rep), p in self.Ua.items():
+            out[f"Ua[{k}|{rep}]"] = p.detach().numpy()
+        for (k, rep), p in self.Ub.items():
+            out[f"Ub[{k}|{rep}]"] = p.detach().numpy()
+        return out
+
+    def load_numpy(self, state) -> None:
+        for name, arr in state.items():
+            if name.startswith("V["):
+                key = name[2:-1]
+                self.attach(key)
+                self.V[key].data = torch.as_tensor(arr)
+            elif name.startswith("Udiag["):
+                key, rep = name[7:-1].split("|")
+                self.attach(key)
+                self.Udiag[(key, int(rep))].data = torch.as_tensor(arr)
+            elif name.startswith("Ua["):
+                key, rep = name[4:-1].split("|")
+                self.attach(key)
+                self.Ua[(key, int(rep))].data = torch.as_tensor(arr)
+            elif name.startswith("Ub["):
+                key, rep = name[4:-1].split("|")
+                self.attach(key)
+                self.Ub[(key, int(rep))].data = torch.as_tensor(arr)