fix(metrics): use the unbiased HSIC estimator in linear CKA

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"})
 
 
@@ -159,7 +163,8 @@ class centroids. ``None`` if fewer than 2 classes.
           effective rank. Always present.
         - ``per_class_rank_<cls>``: 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"``.
@@ -336,10 +341,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 +360,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 +403,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:

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)