diff --git a/monai/metrics/embedding_collapse.py b/monai/metrics/embedding_collapse.py index d33a3eb0bc..51d8b47c9a 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"}) @@ -159,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"``. @@ -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: 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)