feat(metrics): add WelfordVariance and WelfordCovariance helpers (PR 1 of #3748)

[joemunene-by]

Summary

PR 1 of the plan in #3748: introduces ignite/metrics/_running_stats.py with two numerically stable running-statistics primitives that variance- and covariance-bearing metrics can share, instead of each one rolling its own naive Σx² − (Σx)²/n implementation.

• WelfordVariance: running mean, variance, std for a single variable.
• WelfordCovariance: running variance_x, variance_y, covariance, and Pearson correlation() for a paired (x, y) stream.

Both classes:

• keep internal state in float64 regardless of input dtype, so the classic E[X²] − E[X]² cancellation does not bite at large means (the failure mode from [Bug] Numerical instability in PearsonCorrelation due to naive variance formula #3662);
• update incrementally via Welford's online algorithm;
• expose merge() implementing the Chan / Welford parallel formula, suitable for cross-rank distributed reduction or any other case where two accumulators need to be combined without re-iterating the raw data.

No existing metric is touched in this PR. The helper lands first so the two consumer PRs can each diff against a stable shared API.

Follow-ups (already scoped in #3748)

• PR 2: port R2Score to WelfordVariance, with a regression test mirroring [Bug] Numerical instability in PearsonCorrelation due to naive variance formula #3662's mean=1e6, std=1 failure case.
• PR 3: update fix(metrics): use Welford's algorithm in PearsonCorrelation for numerical stability #3741 to drop its inline Welford state and consume WelfordCovariance instead. Pure refactor on top of the existing review.

API

class WelfordVariance:
    n_samples: int
    mean: torch.Tensor                       # running mean
    sum_sq_dev_from_mean: torch.Tensor       # M2 = Σ (x_i − mean)^2

    def update(self, batch: torch.Tensor) -> None: ...
    def merge(self, other: WelfordVariance) -> None: ...

    @property
    def variance(self) -> torch.Tensor: ...  # M2 / n
    @property
    def std(self) -> torch.Tensor: ...       # sqrt(variance)


class WelfordCovariance:
    n_samples: int
    mean_x: torch.Tensor
    mean_y: torch.Tensor
    sum_sq_dev_x: torch.Tensor
    sum_sq_dev_y: torch.Tensor
    sum_product_of_devs: torch.Tensor        # Σ (x_i − mean_x)(y_i − mean_y)

    def update(self, batch_x, batch_y) -> None: ...
    def merge(self, other: WelfordCovariance) -> None: ...

    @property
    def variance_x(self) -> torch.Tensor: ...
    @property
    def variance_y(self) -> torch.Tensor: ...
    @property
    def covariance(self) -> torch.Tensor: ...
    def correlation(self, eps: float = 1e-8) -> torch.Tensor: ...  # Pearson r

Tests

20 unit tests in tests/ignite/metrics/test_running_stats.py, all passing locally with pytest tests/ignite/metrics/test_running_stats.py:

WelfordVariance

• empty accumulator returns zero variance / std without crashing
• single-batch update matches numpy.mean and numpy.var to 1e-12
• multi-batch update matches single-batch on the same data
• merge of two accumulators matches update on the concatenated data
• merge with an empty accumulator on either side is a no-op or absorbs the other side
• numerical-stability regression: 10,000 samples at mean=1e6, std=1 in float32; Welford in float64 recovers the true variance, and the same data through the naive Σx² − (Σx)²/n formula in float32 is asserted to fail by ≥ 0.1, so the test documents the failure mode it is protecting against
• single-sample case (variance is 0, not undefined)
• empty-batch update is a no-op
• reset clears state
• integer input dtypes upcast to float64

WelfordCovariance

• empty accumulator returns zero everywhere
• single-batch update matches numpy.cov(..., bias=True) and numpy.corrcoef to 1e-10 to 1e-12
• multi-batch update matches single-batch
• merge matches concatenated update
• numerical-stability regression: 10,000 paired samples at mean=1e6, true correlation > 0.99; Welford recovers r within 1e-4 of the float64 ground truth
• shape mismatch raises ValueError
• empty-batch update is a no-op
• constant-y series returns r = 0.0 (not NaN) via the eps clamp
• reset clears state

Cross-class sanity

• WelfordCovariance.variance_x matches WelfordVariance.variance fed the same x. Catches future drift between the two implementations.

Conventions

• ruff format clean, ruff check clean.
• Internal state, return values, and properties are all torch.Tensor on the user-supplied device, mirroring the dtype / device convention used elsewhere in ignite.metrics.
• File is named _running_stats.py (leading underscore) and not re-exported from ignite/metrics/__init__.py, so it is internal to the metrics module. Public access stays through individual metric classes; the helper can graduate to a public API later if there is demand.

Test plan

• pytest tests/ignite/metrics/test_running_stats.py: 20 / 20 passing locally
• ruff format --check, ruff check
• CI on this PR

cc @vfdev-5. Opening PR 2 (R2Score port) and PR 3 (PearsonCorrelation refactor on top of #3741) once this lands.

[vfdev-5]

This class should not handle device, neither dtype

[vfdev-5]

Let's make it as a dataclass?

[vfdev-5]

I do not think we should flatten it and for the mean computation we may need to specify the axis (to confirm)

[joemunene-by]

Thanks for the review. Pushed 1b82442 addressing all three points.

1. No device / dtype in the helper. Dropped the constructor args and all internal .to(device) / .to(torch.float64) calls. State now takes the dtype and device of the first batch passed to update. Callers (R2Score, PearsonCorrelation) are responsible for the float64 upcast, which they were already doing in their own update methods anyway, so this is a no-op for the planned consumers and a cleaner contract for any future caller.

2. Dataclasses. Both classes are now @dataclass with field(default_factory=...) for the tensor fields. Dropped the manual __init__ and the reset() method; "reset" becomes reconstruction (self.welford = WelfordVariance()), which fits cleanly into how Metric.reset() is normally written in the consumer classes. Tests updated accordingly: test_reset is now test_fresh_instance_has_zero_state and just asserts the default-factory state.

3. Flatten. Removed the explicit .flatten() call. batch.mean() and batch.numel() already reduce over the full tensor regardless of shape, so the current scalar-reduction behavior is preserved and the code reads cleaner for any input shape.

On axis-aware reduction: the current consumers (R2Score, PearsonCorrelation) both produce a single scalar variance / covariance, so they only need the full-reduction behavior. If we want axis support for future per-channel use cases (running variance over (N, C, H, W) with reduction over (0, 2, 3)), the natural shape of the change is an optional dim parameter on update:

def update(self, batch: torch.Tensor, dim: Optional[Union[int, tuple[int, ...]]] = None) -> None:
    ...
    n_b = batch.numel() if dim is None else _samples_along(batch.shape, dim)
    mean_b = batch.mean(dim=dim, keepdim=False) if dim is not None else batch.mean()
    ...

The state on first update would take the shape of mean_b, and n_samples would become an int count of samples-per-position (still scalar because the reduction is the same at every position).

Two ways we can take it:

• (a) Add the dim parameter in this PR so the helper ships fully-featured, even though no consumer uses it yet.
• (b) Leave it scalar for now; add dim as a follow-up the first time a metric actually needs it (avoids speculative API surface).

I lean toward (b): the consumers don't need it, axis-aware running stats have a few subtleties (Bessel correction conventions, mixed-shape merge, axis reordering) that are easier to design alongside a concrete user. Happy to go either way; let me know which you prefer.

[aaishwarymishra]

@joemunene-by hi, I was trying to review but I find the code to be quite confusing I think we can try to make more readable :)

[aaishwarymishra]

why we writing a function for a code which is literally 1 line?

[aaishwarymishra]

This part is redundant as merge is already doing it we can just use that.

[aaishwarymishra]

Why do we need merge?

[aaishwarymishra]

same here

[joemunene-by]

@aaishwarymishra thanks, your inline points are all fair. Pushed e153f57 addressing each.

1. _zero() for a one-liner. Removed. Each tensor field now uses field(default_factory=lambda: torch.tensor(0.0)) directly. Six lines of one-purpose helper become six explicit defaults.

2. update's body duplicates merge. Refactored. update now builds a single-batch WelfordVariance from the input and calls self.merge(batch_acc). The Chan / Welford parallel formula now lives in exactly one place. Same refactor applied to WelfordCovariance.update (your "same here" on line 192).

The mental model that fell out of this is worth stating: update is the degenerate case of merge where the right-hand side has just been built from one batch. There's one formula, not two. I lifted that into the module docstring and the method docstrings so the next reader sees it immediately.

3. "Why do we need merge?" Distributed reduction. In multi-rank training, each rank accumulates its own WelfordVariance / WelfordCovariance over its local samples; at eval time the ranks fold their accumulators together rank-by-rank to produce the population statistic. Without an explicit merge, that cross-rank reduction has to re-iterate the raw data, which defeats the whole point of an online algorithm. I made this rationale explicit in merge's docstring (previously just one line, now a paragraph).

The bonus is that the two-path design (update for the local-batch case, merge for the cross-state case) collapses to a single formula at the implementation level, which I think also lands the "more readable" feedback from your top-level comment. Let me know if anything's still murky.

cc @vfdev-5 — points from your earlier review are still addressed in 1b82442; this commit is a follow-up to @aaishwarymishra's readability pass. Ready for another look when you have a moment.

[Copilot]

Pull request overview

Adds internal running-statistics primitives intended to centralize numerically stable variance/covariance accumulation for future metric refactors (per #3748), along with a dedicated unit-test suite to validate update/merge correctness and numerical-stability behavior.

Changes:

• Introduce WelfordVariance and WelfordCovariance helpers implementing Welford/Chan online + merge formulas.
• Add unit tests covering empty/single/multi-batch updates, merges, and large-mean numerical-stability regressions.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
ignite/metrics/_running_stats.py	New internal Welford-based running variance/covariance accumulators with merge support.
tests/ignite/metrics/test_running_stats.py	New test suite validating correctness, merge equivalence, and numerical-stability scenarios.

Comments suppressed due to low confidence (1)

ignite/metrics/_running_stats.py:224

• merge() is not wrapped in torch.no_grad() and does not defensively detach incoming tensors unless self is empty. If callers merge state that still requires grad, this can create an autograd graph. Consider adding @torch.no_grad() on merge() and/or detaching other.mean_x/mean_y and the sum_* tensors at the start of the method.

    def merge(self, other: "WelfordCovariance") -> None:
        """Combine ``other`` into ``self`` using the Chan / Welford parallel formula.

        Same correction term as the univariate version, applied three
        times: once for ``sum_sq_dev_x``, once for ``sum_sq_dev_y``, and
        once for ``sum_product_of_devs`` (using ``delta_x * delta_y``
        instead of ``delta * delta``). See
        :meth:`WelfordVariance.merge` for the derivation.
        """

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[joemunene-by]

Copilot's auto-review surfaced four concrete points — pushed 89d5b36 addressing each:

1. Docstring claimed the helpers were already used by PearsonCorrelation / R2Score. They aren't yet (that's PRs 2 and 3 of Numerical stability across running-moment metrics — unified plan (tracking #3662 / #3672 / #3741) #3748). Reworded to "intended consumers (in follow-up PRs of Numerical stability across running-moment metrics — unified plan (tracking #3662 / #3672 / #3741) #3748)" so the doc matches the state at this commit.
2. Docstring claimed internal state in float64. It doesn't the helper is intentionally dtype/device agnostic per @vfdev-5's earlier review, and the caller supplies the float64 cast when stability matters. Realigned the prose so users don't read a false safety guarantee on float32 input.
3. merge() lacked @torch.no_grad(). Both WelfordVariance.merge and WelfordCovariance.merge now have the decorator, matching update. A caller merging an accumulator whose tensors still carry grad would have built an autograd graph and leaked memory over time.
4. cross_coef = n_a * n_b / n_ab precomputed as a Python float in WelfordCovariance.merge. Dropped the temporary, inlined n_a * n_b / n_ab directly into each parallel-formula line mirrors WelfordVariance.merge's existing style and keeps the arithmetic on the same dtype/device as the tensor operands.

20 / 20 tests still passing locally.