cli99 · October 17, 2024 03:41 · Oct 17, 2024
diff --git a/one_pass_variance.py b/one_pass_variance.py
@@ -0,0 +1,41 @@
+# https://jonisalonen.com/2013/deriving-welfords-method-for-computing-variance/
+
+import torch
+
+
+def two_pass_variance(data):
+    n = len(data)
+    mean = sum(data) / n
+    var = sum([(x - mean) ** 2 for x in data]) / (n - 1)
+    return var
+
+
+def one_pass_variance(data):
+    # numerically unstable
+    n = len(data)
+    sum = 0.0
+    sumq = 0.0
+    for x in data:
+        sum += x
+        sumq += x**2
+    mean = sum / n
+    return (sumq - n * mean**2) / (n - 1)
+
+
+def welford_variance(data):
+    m = 0.0
+    s = 0.0
+    for idx, x in enumerate(data):
+        old_m = m
+        m = m + (x - m) / (idx + 1)
+        s = s + (x - m) * (x - old_m)
+    return s / (len(data) - 1)
+
+
+x = torch.randn(10)
+y1 = two_pass_variance(x)
+y2 = one_pass_variance(x)
+y3 = welford_variance(x)
+print(y1, y2, y3)
+assert torch.allclose(y1, y2), "two_pass_variance and one_pass_variance are not equal"
+assert torch.allclose(y1, y3), "two_pass_variance and welford_variance are not equal"