[Fixes #9116] Include per-class rank in EmbeddingCollapseMetric aggregate score

monai/metrics/embedding_collapse.py

@@ -51,8 +51,10 @@ class EmbeddingCollapseMetric(Metric):
     Args:
         reduction: how to aggregate individual scores into ``aggregate``.
             ``"max"`` returns the worst-case score (recommended for
-            safety-critical use). ``"mean"`` returns the average of
-            available scores. ``"none"`` omits the ``aggregate`` key.
+            safety-critical use), including the most collapsed class.
+            ``"mean"`` returns the average of available scores, where all
+            per-class scores contribute a single term. ``"none"`` omits the
+            ``aggregate`` key.
         include_indicators: optional list of indicator names to compute.
             If ``None``, all applicable indicators are computed.
             Valid names: ``"centroid_similarity"``, ``"effective_rank"``,
@@ -162,7 +164,8 @@ class centroids. ``None`` if fewer than 2 classes.
           ``None`` if ``target_embeddings`` not provided.
         - ``separation``: silhouette-based inter-class separation score.
           ``None`` if sklearn unavailable or fewer than 2 classes.
-        - ``aggregate``: reduced score. Omitted when ``reduction="none"``.
+        - ``aggregate``: reduced over the global indicators plus the worst
+          ``per_class_rank_<cls>``. Omitted when ``reduction="none"``.
 
     Raises:
         ValueError: if inputs are invalid (shape, reduction, indicators).
@@ -180,6 +183,7 @@ class centroids. ``None`` if fewer than 2 classes.
 
     inc = set(include_indicators) if include_indicators is not None else None
     scores: dict[str, torch.Tensor | None] = {}
+    per_class_scores: dict[str, torch.Tensor | None] = {}
 
     # -- Label-dependent indicators
     if labels is not None:
@@ -191,7 +195,8 @@ class centroids. ``None`` if fewer than 2 classes.
             scores["centroid_similarity"] = _centroid_similarity(emb, lbl)
 
         if inc is None or "per_class_rank" in inc:
-            scores.update(_per_class_rank(emb, lbl))
+            per_class_scores = _per_class_rank(emb, lbl)
+            scores.update(per_class_scores)
 
         if inc is None or "separation" in inc:
             scores["separation"] = _separation(emb, lbl)
@@ -218,6 +223,11 @@ class centroids. ``None`` if fewer than 2 classes.
     if reduction != "none":
         primary = {"centroid_similarity", "effective_rank_score", "domain_shift", "separation"}
         available = [v.to(device=emb.device) for k, v in scores.items() if k in primary and v is not None]
+        # Per-class collapse is reduced to its worst class before joining the pool, so that
+        # `mean` keeps one vote per indicator regardless of how many classes are present.
+        by_class = [v.to(device=emb.device) for v in per_class_scores.values() if v is not None]
+        if by_class:
+            available.append(torch.stack(by_class).max())
         if not available:
             scores["aggregate"] = None
         elif reduction == "max":

tests/metrics/test_embedding_collapse.py

@@ -462,5 +462,69 @@ def test_raises_without_sklearn(self):
                 )
 
 
+class TestAggregateIncludesPerClassRank(unittest.TestCase):
+    """Asymmetric collapse must reach the summary score, not only the per-class keys."""
+
+    @staticmethod
+    def _majority_healthy_minority_collapsed(n_major=200, n_minor=8, d=64, seed=17):
+        """Build a two-class embedding set: a healthy majority class and a minority
+        class collapsed to a single point.
+
+        Args:
+            n_major: number of samples in the healthy majority class.
+            n_minor: number of samples in the collapsed minority class.
+            d: embedding dimensionality.
+            seed: seed for the majority class's random draw.
+
+        Returns:
+            Tuple of ``(embeddings, labels)`` with shapes ``[n_major + n_minor, d]``
+            and ``[n_major + n_minor]``, where ``labels`` has dtype ``torch.long``.
+        """
+        torch.manual_seed(seed)
+        major = torch.randn(n_major, d)
+        major[:, 0] -= 4.0  # centroid ~ -e1
+        minor = torch.zeros(n_minor, d)
+        minor[:, 0] = 8.0  # identical rows, centroid = +e1
+        emb = torch.cat([major, minor])
+        lbl = torch.tensor([0] * n_major + [1] * n_minor, dtype=torch.long)
+        return emb, lbl
+
+    def test_collapsed_minority_class_reaches_aggregate(self):
+        """A fully collapsed minority class should drive ``aggregate`` to 1.0 under ``max``."""
+        emb, lbl = self._majority_healthy_minority_collapsed()
+        scores = compute_embedding_collapse(emb, lbl, reduction="max")
+        self.assertAlmostEqual(float(scores["per_class_rank_1"]), 1.0, places=5)
+        self.assertAlmostEqual(float(scores["aggregate"]), 1.0, places=5)
+
+    def test_max_aggregate_at_least_worst_per_class(self):
+        """``aggregate`` under ``max`` must never fall below the worst per-class score."""
+        emb, lbl = self._majority_healthy_minority_collapsed()
+        scores = compute_embedding_collapse(emb, lbl, reduction="max")
+        per_class = [float(v) for k, v in scores.items() if k.startswith("per_class_rank_") and v is not None]
+        self.assertGreaterEqual(float(scores["aggregate"]), max(per_class))
+
+    def test_global_indicators_miss_the_collapse(self):
+        """Guard: the fixture is only meaningful while the global view stays healthy."""
+        emb, lbl = self._majority_healthy_minority_collapsed()
+        scores = compute_embedding_collapse(emb, lbl, reduction="none")
+        self.assertLess(float(scores["effective_rank_score"]), 0.2)
+        self.assertLess(float(scores["centroid_similarity"]), 0.2)
+
+    def test_per_class_contributes_one_term_to_mean(self):
+        """Three classes must not outvote the global indicators 3:4 under `mean`."""
+        torch.manual_seed(23)
+        emb = torch.cat([torch.randn(40, 32) + k * 6.0 for k in range(3)])
+        lbl = torch.tensor([0] * 40 + [1] * 40 + [2] * 40, dtype=torch.long)
+        scores = compute_embedding_collapse(emb, lbl, reduction="mean")
+        globals_ = [
+            float(scores[k])
+            for k in ("centroid_similarity", "effective_rank_score", "separation")
+            if scores.get(k) is not None
+        ]
+        per_class = [float(v) for k, v in scores.items() if k.startswith("per_class_rank_") and v is not None]
+        expected = (sum(globals_) + max(per_class)) / (len(globals_) + 1)
+        self.assertAlmostEqual(float(scores["aggregate"]), expected, places=5)
+
+
 if __name__ == "__main__":
     unittest.main()