Multilabel

stable_pretraining/callbacks/knn.py

@@ -47,6 +47,8 @@ class OnlineKNN(Callback):
             to compute all distances at once. Default is -1.
         distance_metric: Distance metric for finding nearest neighbors. Options are
             'euclidean', 'squared_euclidean', 'cosine', 'manhattan'. Default is 'euclidean'.
+        multi_label: Whether the dataset uses multi-label classification. Default is False.
+        num_classes: Number of total classes. Required if multi_label is True.
 
     Raises:
         ValueError: If k <= 0, temperature <= 0, or chunk_size is invalid.
@@ -73,6 +75,8 @@ def __init__(
         distance_metric: Literal[
             "euclidean", "squared_euclidean", "cosine", "manhattan"
         ] = "euclidean",
+        multi_label: bool = False,
+        num_classes: Optional[int] = None,
     ) -> None:
         super().__init__()
 
@@ -83,6 +87,9 @@ def __init__(
         if chunk_size == 0 or chunk_size < -1:
             raise ValueError(f"chunk_size must be positive or -1, got {chunk_size}")
 
+        if multi_label and num_classes is None:
+            raise ValueError("You must provide `num_classes` when `multi_label=True`.")
+
         if input_dim is not None and isinstance(input_dim, (list, tuple)):
             input_dim = int(np.prod(input_dim))
 
@@ -97,10 +104,20 @@ def __init__(
         self.chunk_size = chunk_size
         self.distance_metric = distance_metric
         self.metrics = metrics
+        self.multi_label = multi_label
+        self.num_classes = num_classes
 
         self._input_queue = None
         self._target_queue = None
 
+        for metric in self.metrics.values():
+            if metric.__class__.__name__ in ["MultilabelAUROC", "MulticlassAUROC"]:
+                logging.warning(
+                    f"[{self.name}] AUROC metric detected. torchmetrics expects target labels to be Long integers. "
+                    "If you are passing float targets, you can cast them to long in your dataset/forward pass. "
+                )
+                break
+
     @property
     def state_key(self) -> str:
         """Unique identifier for this callback's state during checkpointing."""
@@ -195,11 +212,6 @@ def _compute_knn_predictions(
     ) -> Optional[Tensor]:
         """Compute KNN predictions."""
         batch_size = features.size(0)
-        num_classes = int(cached_labels.max().item()) + 1
-
-        predictions = torch.zeros(
-            batch_size, num_classes, device=features.device, dtype=torch.float32
-        )
 
         if cached_features.device != features.device:
             cached_features = cached_features.to(features.device)
@@ -223,13 +235,21 @@ def _compute_knn_predictions(
 
         dist_weight = 1 / dist_weight.add_(self.temperature)
 
-        labels_1d = (
-            cached_labels.squeeze(-1) if cached_labels.dim() > 1 else cached_labels
-        )
-        selected_labels = labels_1d[sim_indices].long()
-        one_hot_labels = F.one_hot(selected_labels, num_classes=num_classes)
+        if self.multi_label:
+            selected_labels = cached_labels[sim_indices].float()
+            weighted_votes = dist_weight.unsqueeze(-1) * selected_labels
+            total_votes = weighted_votes.sum(dim=0)
+            sum_of_weights = dist_weight.sum(dim=0).unsqueeze(-1)
+            predictions = total_votes / sum_of_weights
+        else:
+            num_classes = int(cached_labels.max().item()) + 1
+            labels_1d = (
+                cached_labels.squeeze(-1) if cached_labels.dim() > 1 else cached_labels
+            )
+            selected_labels = labels_1d[sim_indices].long()
+            one_hot_labels = F.one_hot(selected_labels, num_classes=num_classes)
+            predictions = (dist_weight.unsqueeze(-1) * one_hot_labels).sum(0)
 
-        predictions = (dist_weight.unsqueeze(-1) * one_hot_labels).sum(0)
         return predictions
 
     def _log_metrics(

stable_pretraining/callbacks/probe.py

@@ -108,6 +108,15 @@ def __init__(
         self.metrics = metrics
         logging.info(f"{self.name}: We are wrapping up your `forward`!")
         self.wrap_forward(pl_module=module)
+        for metric in self.metrics.values():
+            if metric.__class__.__name__ in ["MultilabelAUROC", "MulticlassAUROC"]:
+                logging.warning(
+                    f"[{self.name}] AUROC metric detected. torchmetrics expects target labels to be Long integers. "
+                    "If you are passing float targets, you can cast them to long in your dataset/forward pass. "
+                    "If your loss function requires floats, you can handle it there (e.g., "
+                    "`loss=lambda p, t: torch.nn.functional.binary_cross_entropy_with_logits(p, t.float())`)."
+                )
+                break
 
     def configure_model(self, pl_module: LightningModule) -> torch.nn.Module:
         """Initialize the probe module from configuration."""

stable_pretraining/tests/unit/test_probing.py

@@ -249,3 +249,77 @@ def test_online_probe_lifecycle_methods(self):
         assert len(module.callbacks_metrics) == 1
         assert len(module.callbacks_modules) == 1
         module.configure_optimizers()
+
+    @patch("stable_pretraining.callbacks.knn.logging.warning")
+    def test_online_knn_multilabel_initialization(self, mock_warning):
+        """Test multi-label initialization, required args, and AUROC warnings."""
+        from stable_pretraining.callbacks import OnlineKNN
+        import torchmetrics
+
+        # Test missing num_classes raises ValueError
+        with pytest.raises(ValueError, match="You must provide `num_classes`"):
+            OnlineKNN(
+                name="knn_probe",
+                input="embedding",
+                target="label",
+                queue_length=100,
+                metrics={},
+                multi_label=True,
+            )
+
+        # Test AUROC warning is triggered properly
+        metrics = {"auroc": torchmetrics.classification.MultilabelAUROC(num_labels=3)}
+        knn = OnlineKNN(
+            name="knn_probe",
+            input="embedding",
+            target="label",
+            queue_length=100,
+            metrics=metrics,
+            multi_label=True,
+            num_classes=3,
+        )
+
+        assert knn.multi_label is True
+        assert knn.num_classes == 3
+        mock_warning.assert_called_once()
+        assert "AUROC metric detected" in mock_warning.call_args[0][0]
+
+    def test_online_knn_multilabel_math(self):
+        """Test the multi-label KNN distance weighting and voting math."""
+        from stable_pretraining.callbacks import OnlineKNN
+
+        knn = OnlineKNN(
+            name="knn_probe",
+            input="embedding",
+            target="label",
+            queue_length=10,
+            metrics={},
+            k=2,
+            temperature=1.0,  # dist_weight = 1 / (dist + 1)
+            multi_label=True,
+            num_classes=3,
+        )
+
+        # Setup mock data
+        query = torch.tensor([[0.0, 0.0]])
+        cached_features = torch.tensor(
+            [
+                [0.0, 0.0],  # dist=0 -> weight=1.0
+                [2.0, 0.0],  # dist=2 -> weight=0.3333
+                [10.0, 10.0],  # Ignored because k=2
+            ]
+        )
+
+        cached_labels = torch.tensor(
+            [[1.0, 0.0, 1.0], [1.0, 1.0, 0.0], [0.0, 0.0, 0.0]]
+        )
+
+        preds = knn._compute_knn_predictions(query, cached_features, cached_labels)
+
+        # Expected votes:
+        # Class 0: (1.0 * 1) + (0.333 * 1) = 1.333 / 1.333 = 1.0
+        # Class 1: (1.0 * 0) + (0.333 * 1) = 0.333 / 1.333 = 0.25
+        # Class 2: (1.0 * 1) + (0.333 * 0) = 1.000 / 1.333 = 0.75
+        expected = torch.tensor([[1.00, 0.25, 0.75]])
+
+        assert torch.allclose(preds, expected, atol=1e-3)