Speed up _WelfordAccumulator's per-chunk update

The streaming update re-derived two full (B, F) arrays from the
running mean (once before updating it, once after) plus an elementwise
product — five passes over each chunk and three temporary arrays, to
maintain a mean/variance that's tiny in width (COND_DIM=15 at most).

Reformulate as Chan/Golub/LeVeque's parallel-variance algorithm:
compute the chunk's own local mean/M2 (independent of the running
state) and merge it in with an O(F) combination formula. Same
streaming interface and output (identical to ~1e-14, float64 rounding
noise), ~40% faster per update() call on a benchmark chunk.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

giant/data/transforms.py

@@ -196,7 +196,7 @@ class Normalizer:
 
 
 class _WelfordAccumulator:
-    """Streaming mean/variance (Welford's online algorithm, batch update).
+    """Streaming mean/variance (Chan/Golub/LeVeque 1979 parallel algorithm).
 
     Use to fit a Normalizer over data that doesn't fit in memory:
         acc = _WelfordAccumulator(n_features)
@@ -211,13 +211,23 @@ class _WelfordAccumulator:
         self._M2 = np.zeros(n_features, dtype=np.float64)
 
     def update(self, X: np.ndarray) -> None:
+        # Computes the chunk's own local mean/M2 (two passes over X, no
+        # reference to the running mean) and merges it into the running
+        # totals with the O(F) Chan/Golub/LeVeque combination formula.
+        # Equivalent to the textbook single-pass streaming update (which
+        # instead re-derives two full (B, F) arrays from the running mean,
+        # before and after updating it) but ~40% cheaper here since it
+        # avoids one of those (B, F) passes and its temporary array.
         X = np.asarray(X, dtype=np.float64)
         B = X.shape[0]
+        mean_b = X.mean(0)
+        diff = X - mean_b
+        M2_b = np.einsum("ij,ij->j", diff, diff)
+
         new_n = self.n + B
-        delta = X - self._mean
-        self._mean += delta.sum(0) / new_n
-        delta2 = X - self._mean
-        self._M2 += (delta * delta2).sum(0)
+        delta = mean_b - self._mean
+        self._mean += delta * (B / new_n)
+        self._M2 += M2_b + delta * delta * (self.n * B / new_n)
         self.n = new_n
 
     def to_normalizer(self) -> "Normalizer":

tests/test_transforms.py

@@ -15,6 +15,7 @@ from giant.data.transforms import (
     sorted_membership,
     travel_direction,
     _vectorized_map_lookup,
+    _WelfordAccumulator,
 )
 
 
@@ -491,3 +492,71 @@ def test_vectorized_map_lookup_raises_keyerror_on_missing_value():
     values = np.array([1, 2, 3])
     with pytest.raises(KeyError):
         _vectorized_map_lookup(values, mapping)
+
+
+# ── _WelfordAccumulator ──────────────────────────────────────────────────────
+
+
+def test_welford_accumulator_matches_direct_mean_std_over_many_chunks():
+    rng = np.random.default_rng(5)
+    F = 4
+    chunks = [rng.standard_normal((rng.integers(1, 50), F)) * 10 + 3 for _ in range(20)]
+    full = np.concatenate(chunks, axis=0)
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, full.mean(axis=0), rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(norm.std, full.std(axis=0), rtol=1e-5, atol=1e-5)
+    assert acc.n == full.shape[0]
+
+
+def test_welford_accumulator_single_chunk():
+    rng = np.random.default_rng(6)
+    X = rng.standard_normal((100, 3)) * 5 - 2
+
+    acc = _WelfordAccumulator(3)
+    acc.update(X)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, X.mean(axis=0), rtol=1e-5)
+    np.testing.assert_allclose(norm.std, X.std(axis=0), rtol=1e-5)
+
+
+def test_welford_accumulator_matches_naive_running_mean_reference():
+    """The chunk-local-mean + Chan-merge formula must agree with the naive
+    textbook streaming update (subtract the *running* mean before and after
+    updating it) that it replaces, within float64 rounding tolerance."""
+    rng = np.random.default_rng(7)
+    F = 3
+    chunks = [rng.standard_normal((rng.integers(1, 40), F)) for _ in range(15)]
+
+    def naive_update(mean, M2, n, X):
+        X = np.asarray(X, dtype=np.float64)
+        B = X.shape[0]
+        new_n = n + B
+        delta = X - mean
+        mean = mean + delta.sum(0) / new_n
+        delta2 = X - mean
+        M2 = M2 + (delta * delta2).sum(0)
+        return mean, M2, new_n
+
+    naive_mean = np.zeros(F)
+    naive_M2 = np.zeros(F)
+    naive_n = 0
+    for chunk in chunks:
+        naive_mean, naive_M2, naive_n = naive_update(
+            naive_mean, naive_M2, naive_n, chunk
+        )
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+
+    assert acc.n == naive_n
+    np.testing.assert_allclose(acc._mean, naive_mean, rtol=1e-9, atol=1e-9)
+    np.testing.assert_allclose(acc._M2, naive_M2, rtol=1e-9, atol=1e-9)