HAKO upload: hako/core/kmeanspp.py

hako/core/kmeanspp.py

@@ -0,0 +1,89 @@
+"""K-Means++ seeding + Lloyd refinement.
+
+Theorem T-KM (existence, monotonicity, approximation).
+(1) Boundedness: the frozen source features satisfy ||z|| <= C (layer-normed
+    activations / bounded quantized weights), hence the k-means cost
+        J_KM({mu}) = sum_i min_k ||z_i - mu_k||^2
+    is finite and attains its infimum over the compact set
+    {mu : ||mu_k|| <= 2C} (Weierstrass).
+(2) Monotone decrease: the assignment step (point -> nearest centroid) never
+    increases J_KM for fixed centroids; the update step (centroid -> mean of
+    its assigned points) is the exact argmin of J_KM for a fixed assignment
+    (d/dmu_k sum ||z - mu_k||^2 = 0 gives mu_k = mean). Therefore every Lloyd
+    sweep weakly decreases J_KM. There are finitely many assignments (K^N),
+    no assignment repeats under strict decrease, so Lloyd converges to a
+    local minimum in finitely many sweeps. QED.
+(3) Approximation: k-means++ seeding (D^2 sampling) yields expected cost
+        E[J_seed] <= 8 (ln K + 2) * OPT
+    (Arthur & Vassilvitskii, 2007); Lloyd afterwards only improves it by (2).
+
+Implementation is numpy-vectorized; K seeds by D^2 sampling with O(N K d)
+total work, then a fixed number of Lloyd sweeps (early-stop on delta).
+"""
+from __future__ import annotations
+
+from typing import Tuple
+
+import numpy as np
+
+
+def _d2_sample(X: np.ndarray, K: int, rng: np.random.Generator) -> np.ndarray:
+    """K-means++ D^2 seeding. Returns (K, d) initial centroids."""
+    n = X.shape[0]
+    cent = np.empty((K, X.shape[1]), dtype=np.float32)
+    cent[0] = X[rng.integers(n)]
+    d2 = ((X - cent[0]) ** 2).sum(axis=1)
+    for k in range(1, K):
+        total = d2.sum()
+        if total <= 1e-12:
+            cent[k:] = X[rng.integers(n, size=K - k)]
+            break
+        probs = d2 / total
+        idx = rng.choice(n, p=probs)
+        cent[k] = X[idx]
+        d2 = np.minimum(d2, ((X - cent[k]) ** 2).sum(axis=1))
+    return cent
+
+
+def kmeanspp_lloyd(X: np.ndarray, K: int, n_init: int = 1, sweeps: int = 24,
+                   seed: int = 7,
+                   tol: float = 1e-6) -> Tuple[np.ndarray, np.ndarray, float]:
+    """Returns (centroids (K,d) float32, labels (n,), inertia float).
+
+    Theorem T-KM(2) is enforced structurally: we only accept a sweep result
+    if inertia is non-increasing (guards float noise with a tolerance).
+    """
+    assert K >= 1 and X.ndim == 2
+    X = np.ascontiguousarray(X, dtype=np.float32)
+    n = X.shape[0]
+    rng = np.random.default_rng(seed)
+    best = None
+    for init in range(max(1, n_init)):
+        cent = _d2_sample(X, K, rng)
+        prev_J = np.inf
+        labels = np.zeros(n, dtype=np.int64)
+        for _ in range(sweeps):
+            # assignment (chunked to bound RAM)
+            chunk = max(1, (64 * 1024 * 1024) // max(1, 4 * K * X.shape[1]))
+            mins = np.empty(n, dtype=np.float32)
+            for s in range(0, n, chunk):
+                e = min(n, s + chunk)
+                d2 = ((X[s:e, None, :] - cent[None, :, :]) ** 2).sum(axis=2)
+                labels[s:e] = d2.argmin(axis=1)
+                mins[s:e] = d2.min(axis=1)
+            J = float(mins.sum())
+            if prev_J - J < tol * max(1.0, prev_J):
+                prev_J = min(prev_J, J)
+                break
+            prev_J = J
+            # update
+            for k in range(K):
+                m = labels == k
+                if m.any():
+                    cent[k] = X[m].mean(axis=0)
+                else:  # dead centroid -> reseed at farthest point
+                    far = ((X - cent[labels]) ** 2).sum(axis=1)
+                    cent[k] = X[int(far.argmax())]
+        if best is None or prev_J < best[2]:
+            best = (cent.copy(), labels.copy(), prev_J)
+    return best