diff --git a/monai/metrics/embedding_collapse.py b/monai/metrics/embedding_collapse.py index b427660b90..d33a3eb0bc 100644 --- a/monai/metrics/embedding_collapse.py +++ b/monai/metrics/embedding_collapse.py @@ -295,7 +295,10 @@ def _effective_rank_score(emb: torch.Tensor) -> torch.Tensor: probs = sv / sv_sum safe_probs = probs.clamp_min(torch.finfo(probs.dtype).tiny) eff_rank: torch.Tensor = torch.exp(-(probs * safe_probs.log()).sum()) - max_rank: torch.Tensor = emb.new_tensor(float(min(emb.shape[0], emb.shape[1]))) + # Mean-centering above forces the rows to sum to zero, a linear dependency, so + # rank(centered) <= min(N - 1, D). Normalizing by min(N, D) would impose a 1/N + # floor on the score whenever N <= D. No-op when N > D: both expressions equal D. + max_rank: torch.Tensor = emb.new_tensor(float(min(emb.shape[0] - 1, emb.shape[1]))) return (emb.new_tensor(1.0) - eff_rank / max_rank).clamp(0.0, 1.0) diff --git a/tests/metrics/test_embedding_collapse.py b/tests/metrics/test_embedding_collapse.py index b406f92215..bdf079c0f1 100644 --- a/tests/metrics/test_embedding_collapse.py +++ b/tests/metrics/test_embedding_collapse.py @@ -224,10 +224,31 @@ def test_formula_matches_manual(self): probs = sv / sv_sum safe_probs = probs.clamp_min(torch.finfo(probs.dtype).tiny) eff_rank = (-(probs * safe_probs.log()).sum()).exp() - expected = (1.0 - eff_rank / min(20, 16)).clamp(0.0, 1.0) + expected = (1.0 - eff_rank / min(20 - 1, 16)).clamp(0.0, 1.0) score = _effective_rank_score(emb) self.assertAlmostEqual(float(score), float(expected), places=5) + def test_effective_rank_no_collapse_when_n_leq_d(self): + # Isotropic embeddings have NO dimensional collapse -> score must be ~0. + # RED on dev: returns ~0.25 (the 1/N floor). GREEN after fix: ~0.0004. + torch.manual_seed(0) + score = _effective_rank_score(torch.randn(4, 768)) + self.assertLess(float(score), 0.05) + + def test_effective_rank_two_distinct_samples_report_no_collapse(self): + # Two maximally-spread points span the only subspace 2 samples CAN span. + # RED on dev: returns exactly 0.5 ("50% collapsed"). GREEN after fix: 0.0. + emb = torch.zeros(2, 768) + emb[0, 0] = 1.0 + emb[1, 0] = -1.0 + self.assertAlmostEqual(float(_effective_rank_score(emb)), 0.0, places=5) + + def test_effective_rank_unchanged_when_n_gt_d(self): + # Regression guard: the fix MUST be a no-op where the code is already correct. + torch.manual_seed(0) + score = _effective_rank_score(torch.randn(1024, 768)) + self.assertAlmostEqual(float(score), 0.1234, places=3) + class TestPerClassRank(unittest.TestCase): def test_keys_present_for_each_class(self):