From 3e05e4e378e4390d179721260a6f10547c1178a8 Mon Sep 17 00:00:00 2001 From: rubenuni1009 <183279777+rubenG1009@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:14:21 +0200 Subject: [PATCH 1/2] fix(metrics): use the unbiased HSIC estimator in linear CKA `_hsic` was documented as the unbiased estimator but implemented the biased plug-in one: it kept the Gram diagonals and normalised by `(n - 1) ** 2`. That carries an `O(1/n)` term, so on independent representations the `domain_shift` indicator reported similarity that came from the sample count rather than from the data: n=32 n=64 n=128 n=256 before 0.198 0.114 0.056 0.029 after 0.017 0.009 0.003 0.001 Measured over 20 seeds per point with independent standard-normal inputs, where linear CKA is 0. Testing `_hsic` directly on independent inputs, the old estimator's mean sat 43-72 standard errors above zero at every sample count tested; the new one stays within 1.1. `_domain_shift` subsamples to the smaller of its two inputs, so this also made scores incomparable between domains of different sizes. The `1 / (n - 1) ** 2` factor itself cancels in the CKA ratio, so it was not the cause; zeroing the Gram diagonals is what removes the bias. Two consequences worth calling out: - The unbiased estimator has an `n - 3` factor, so linear CKA is undefined below four samples per domain. `_domain_shift` now returns `None` for n < 4 where it previously returned a number. That number carried the largest bias of all, and `None` is the existing idiom in this module for an indicator that cannot be computed. - An unbiased estimate of a non-negative quantity can be negative, so `_domain_shift` guards against a non-positive self-HSIC that would make the normaliser NaN. Reference: Song et al. (2007), "Supervised Feature Selection via Dependence Estimation", Eq. 5. Kornblith et al. (2019) use this estimator for CKA so that values are comparable across sample counts. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 Signed-off-by: rubenuni1009 <183279777+rubenG1009@users.noreply.github.com> --- monai/metrics/embedding_collapse.py | 44 ++++++++++++++++---- tests/metrics/test_embedding_collapse.py | 53 ++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 8 deletions(-) diff --git a/monai/metrics/embedding_collapse.py b/monai/metrics/embedding_collapse.py index d33a3eb0bc..2059da1190 100644 --- a/monai/metrics/embedding_collapse.py +++ b/monai/metrics/embedding_collapse.py @@ -25,6 +25,10 @@ __all__ = ["EmbeddingCollapseMetric", "compute_embedding_collapse"] _VALID_REDUCTIONS = ("max", "mean", "none") +# The unbiased HSIC estimator carries an ``n - 3`` factor, so linear CKA needs at least +# four samples per domain to be defined at all. +_MIN_HSIC_SAMPLES = 4 + _VALID_INDICATORS = frozenset({"centroid_similarity", "effective_rank", "per_class_rank", "domain_shift", "separation"}) @@ -336,10 +340,11 @@ def _domain_shift(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor | target: ``[M, D]`` float tensor. Returns: - Scalar tensor in ``[0, 1]``, or ``None`` if either set has < 2 samples. + Scalar tensor in ``[0, 1]``, or ``None`` if either set has fewer than + ``_MIN_HSIC_SAMPLES`` samples. 1.0 = representations identical. 0.0 = representations orthogonal. """ - if source.shape[0] < 2 or target.shape[0] < 2: + if source.shape[0] < _MIN_HSIC_SAMPLES or target.shape[0] < _MIN_HSIC_SAMPLES: return None if source.shape[0] != target.shape[0]: @@ -354,10 +359,11 @@ def _domain_shift(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor | hsic_xy = _hsic(source, target) hsic_xx = _hsic(source, source) hsic_yy = _hsic(target, target) - denom = (hsic_xx * hsic_yy).sqrt() - if denom == 0.0: + # The unbiased estimator can return a non-positive self-HSIC on degenerate + # or very small samples, which would make the normaliser NaN. + if hsic_xx <= 0.0 or hsic_yy <= 0.0: return source.new_tensor(0.0) - return (hsic_xy / denom).clamp(0.0, 1.0) + return (hsic_xy / (hsic_xx * hsic_yy).sqrt()).clamp(0.0, 1.0) def _separation(emb: torch.Tensor, labels: torch.Tensor) -> torch.Tensor | None: @@ -396,12 +402,34 @@ def _separation(emb: torch.Tensor, labels: torch.Tensor) -> torch.Tensor | None: def _hsic(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - """Unbiased linear HSIC estimator.""" + """Unbiased linear HSIC estimator of Song et al. (2007), Eq. 5. + + The Gram matrices have their diagonals zeroed before the trace terms are + formed, which is what removes the ``O(1/n)`` bias of the plug-in estimator. + + Args: + x: ``[N, D]`` float tensor. + y: ``[N, D]`` float tensor with the same ``N`` as ``x``. + + Returns: + Scalar tensor. Unlike the biased estimator this may be negative, since + an unbiased estimate of a non-negative quantity is not itself bounded + below by zero; callers must handle that. + + Note: + Requires ``N >= 4``; the ``n - 3`` factor makes the estimator undefined + below that. ``_domain_shift`` enforces this before calling. + """ n = x.shape[0] gram_x = x @ x.T gram_y = y @ y.T - centering = torch.eye(n, dtype=x.dtype, device=x.device) - torch.ones(n, n, dtype=x.dtype, device=x.device) / n - return torch.sum((centering @ gram_x @ centering) * (centering @ gram_y @ centering)) / ((n - 1) ** 2) + gram_x.fill_diagonal_(0) + gram_y.fill_diagonal_(0) + sum_x, sum_y = gram_x.sum(), gram_y.sum() + term_trace = (gram_x * gram_y).sum() + term_product = sum_x * sum_y / ((n - 1) * (n - 2)) + term_cross = 2.0 * (gram_x.sum(dim=0) @ gram_y.sum(dim=0)) / (n - 2) + return (term_trace + term_product - term_cross) / (n * (n - 3)) def _validate_embeddings(embeddings: torch.Tensor) -> None: diff --git a/tests/metrics/test_embedding_collapse.py b/tests/metrics/test_embedding_collapse.py index bdf079c0f1..3ce58c7714 100644 --- a/tests/metrics/test_embedding_collapse.py +++ b/tests/metrics/test_embedding_collapse.py @@ -16,10 +16,12 @@ import torch from monai.metrics.embedding_collapse import ( + _MIN_HSIC_SAMPLES, EmbeddingCollapseMetric, _centroid_similarity, _domain_shift, _effective_rank_score, + _hsic, _per_class_rank, compute_embedding_collapse, linear_probe_accuracy, @@ -302,6 +304,57 @@ def test_score_in_unit_interval(self): self.assertGreaterEqual(float(score), 0.0) self.assertLessEqual(float(score), 1.0) + def test_hsic_is_unbiased_for_independent_inputs(self): + """HSIC of independent representations must average to zero at every n. + + The biased plug-in estimator carries an ``O(1/n)`` term, so on + independent inputs its mean sat tens of standard errors above zero + (t > 40 for every sample count tested) instead of at zero. The + unbiased estimator of Song et al. (2007) removes that term. + """ + for n in (8, 32, 128): + with self.subTest(n=n): + values = [] + for seed in range(100): + torch.manual_seed(seed) + values.append(float(_hsic(torch.randn(n, 8), torch.randn(n, 8)))) + mean = sum(values) / len(values) + variance = sum((v - mean) ** 2 for v in values) / (len(values) - 1) + std_error = (variance / len(values)) ** 0.5 + self.assertLess( + abs(mean / std_error), 3.0, f"HSIC mean {mean:.4f} is {mean / std_error:.1f} SEs from 0 at n={n}" + ) + + def test_score_does_not_drift_with_sample_count(self): + """Unrelated representations must not look more similar at smaller n. + + With the biased estimator the score fell as roughly ``1/n`` -- 0.20 at + n=32, 0.11 at n=64, 0.06 at n=128, 0.03 at n=256 -- so a score was only + comparable against another computed at the same sample count. + ``_domain_shift`` subsamples to the smaller of its two inputs, which + made that easy to hit. + """ + means = [] + for n in (32, 64, 128, 256): + scores = [] + for seed in range(20): + torch.manual_seed(seed) + score = _domain_shift(torch.randn(n, 8), torch.randn(n, 8)) + self.assertIsNotNone(score) + scores.append(float(score)) + means.append(sum(scores) / len(scores)) + self.assertLess(max(means), 0.1, f"independent inputs scored {means} at n=32,64,128,256") + self.assertLess(max(means) - min(means), 0.1, f"score drifts with sample count: {means}") + + def test_below_minimum_samples_returns_none(self): + """Linear CKA is undefined below ``_MIN_HSIC_SAMPLES`` samples.""" + torch.manual_seed(14) + target = torch.randn(10, 8) + for n in range(1, _MIN_HSIC_SAMPLES): + with self.subTest(n=n): + self.assertIsNone(_domain_shift(torch.randn(n, 8), target)) + self.assertIsNone(_domain_shift(target, torch.randn(n, 8))) + def test_single_sample_returns_none(self): src = torch.randn(1, 8) tgt = torch.randn(10, 8) From d991c6fac66f62b0451a78da0ec4e59aab96debd Mon Sep 17 00:00:00 2001 From: rubenuni1009 <183279777+rubenG1009@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:19:52 +0200 Subject: [PATCH 2/2] docs(metrics): note the new domain_shift minimum in the public contract `compute_embedding_collapse` documented `domain_shift` as `None` only when `target_embeddings` is absent. Raising the linear CKA minimum to four samples per domain added a second case that the return contract did not mention. Assisted-by: Claude Code Co-Authored-By: Claude Opus 5 Signed-off-by: rubenuni1009 <183279777+rubenG1009@users.noreply.github.com> --- monai/metrics/embedding_collapse.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/monai/metrics/embedding_collapse.py b/monai/metrics/embedding_collapse.py index 2059da1190..51d8b47c9a 100644 --- a/monai/metrics/embedding_collapse.py +++ b/monai/metrics/embedding_collapse.py @@ -163,7 +163,8 @@ class centroids. ``None`` if fewer than 2 classes. effective rank. Always present. - ``per_class_rank_``: per-class effective rank score. - ``domain_shift``: linear CKA between source and target domains. - ``None`` if ``target_embeddings`` not provided. + ``None`` if ``target_embeddings`` is not provided, or if either + domain has fewer than ``_MIN_HSIC_SAMPLES`` samples. - ``separation``: silhouette-based inter-class separation score. ``None`` if sklearn unavailable or fewer than 2 classes. - ``aggregate``: reduced score. Omitted when ``reduction="none"``.