fix(training): pair only valid anchors in contrastive pretraining

deeptab/training/pretraining.py

@@ -164,6 +164,32 @@ def forward(self, x):
         return x  # Return unpooled sequence embeddings (N, S, D)
 
     def get_knn(self, labels):
+        """Return the positive and negative neighbor indices for a batch.
+
+        For classification, rows without both a same-class and a
+        different-class neighbor are left out, so row ``i`` of the result need
+        not belong to batch row ``i``.
+
+        Parameters
+        ----------
+        labels : torch.Tensor
+            Batch targets: class indices of shape ``(N,)`` for classification,
+            or target values of shape ``(N, 1)`` for regression.
+
+        Returns
+        -------
+        knn_indices, neg_indices : torch.Tensor
+            Positive and negative neighbor indices, each of shape ``(A, k)``.
+        """
+        knn_indices, neg_indices, _ = self._pair_indices(labels)
+        return knn_indices, neg_indices
+
+    def _pair_indices(self, labels):
+        """Return ``(knn_indices, neg_indices, anchor_indices)`` for a batch.
+
+        Row ``i`` of the neighbor indices belongs to the anchor row
+        ``anchor_indices[i]`` of the batch.
+        """
         batch_size = labels.size(0)
         k_neighbors = min(self.k_neighbors, batch_size - 1)
 
@@ -191,22 +217,45 @@ def get_knn(self, labels):
                 )
             knn_indices = torch.stack([knn_indices_list[i] for i in valid])
             neg_indices = torch.stack([neg_indices_list[i] for i in valid])
+            anchor_indices = torch.tensor(valid, dtype=torch.long)
         else:
             with torch.no_grad():
                 target_distances = torch.cdist(labels.float(), labels.float(), p=2).squeeze(-1)
             knn_indices = target_distances.topk(k_neighbors + 1, largest=False).indices[:, 1:]
             neg_indices = target_distances.topk(k_neighbors, largest=True).indices
+            anchor_indices = torch.arange(batch_size)
+
+        return knn_indices.to(self.device), neg_indices.to(self.device), anchor_indices.to(self.device)
+
+    def contrastive_loss(self, embeddings, knn_indices, neg_indices, anchor_indices=None):
+        """Cosine-embedding loss between anchors and their positive/negative neighbors.
 
-        return knn_indices.to(self.device), neg_indices.to(self.device)
+        Parameters
+        ----------
+        embeddings : torch.Tensor
+            Batch embeddings, shape ``(N, D)``, or ``(N, S, D)`` when
+            ``pool_sequence=False``.
+        knn_indices, neg_indices : torch.Tensor
+            Positive and negative neighbor indices, shape ``(A, k)``; row ``i``
+            belongs to the anchor ``anchor_indices[i]``.
+        anchor_indices : torch.Tensor, optional
+            Batch rows used as anchors, shape ``(A,)``. ``None`` uses every row.
 
-    def contrastive_loss(self, embeddings, knn_indices, neg_indices):
+        Returns
+        -------
+        torch.Tensor
+            Scalar loss (summed over sequence positions when
+            ``pool_sequence=False``).
+        """
         if not self.pool_sequence:
             N, S, D = embeddings.shape
             loss = 0.0
             for i in range(S):
                 embs = embeddings[:, i, :]
                 k_neighbors = knn_indices.shape[1]
                 embs = F.normalize(embs, p=2, dim=-1)
+                anchors = embs if anchor_indices is None else embs[anchor_indices]
+                N = anchors.shape[0]
 
                 positive_pairs = embs[knn_indices] if self.use_positive else None
                 negative_pairs = embs[neg_indices] if self.use_negative else None
@@ -229,16 +278,18 @@ def contrastive_loss(self, embeddings, knn_indices, neg_indices):
 
                 # Anchor layout must match cat(pairs): one anchor-major block of
                 # k rows per anchor, tiled once per pair type.
-                embeddings_s = embs.repeat_interleave(k_neighbors, dim=0).repeat(len(pairs), 1)
+                embeddings_s = anchors.repeat_interleave(k_neighbors, dim=0).repeat(len(pairs), 1)
                 _loss = self.loss_fn(embeddings_s, all_pairs, all_pair_labels)
                 loss += _loss
 
             return loss
 
         else:
-            N, D = embeddings.shape
+            D = embeddings.shape[1]
             k_neighbors = knn_indices.shape[1]
             embeddings = F.normalize(embeddings, p=2, dim=-1)
+            anchors = embeddings if anchor_indices is None else embeddings[anchor_indices]
+            N = anchors.shape[0]
 
             positive_pairs = embeddings[knn_indices] if self.use_positive else None
             negative_pairs = embeddings[neg_indices] if self.use_negative else None
@@ -261,7 +312,7 @@ def contrastive_loss(self, embeddings, knn_indices, neg_indices):
 
             # Anchor layout must match cat(pairs): one anchor-major block of
             # k rows per anchor, tiled once per pair type.
-            embeddings_s = embeddings.repeat_interleave(k_neighbors, dim=0).repeat(len(pairs), 1)
+            embeddings_s = anchors.repeat_interleave(k_neighbors, dim=0).repeat(len(pairs), 1)
             loss = self.loss_fn(embeddings_s, all_pairs, all_pair_labels)
             return loss
 
@@ -270,25 +321,25 @@ def training_step(self, batch, batch_idx):
 
         data, labels = batch
         embeddings = self(data)
-        knn_indices, neg_indices = self.get_knn(labels)
-        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices)
+        knn_indices, neg_indices, anchor_indices = self._pair_indices(labels)
+        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices, anchor_indices)
 
         self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
         return loss
 
     def test_step(self, batch, batch_idx):
         data, labels = batch
         embeddings = self(data)
-        knn_indices, neg_indices = self.get_knn(labels)
-        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices)
+        knn_indices, neg_indices, anchor_indices = self._pair_indices(labels)
+        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices, anchor_indices)
         self.log("test_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True)
         return loss
 
     def validation_step(self, batch, batch_idx):
         data, labels = batch
         embeddings = self(data)
-        knn_indices, neg_indices = self.get_knn(labels)
-        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices)
+        knn_indices, neg_indices, anchor_indices = self._pair_indices(labels)
+        loss = self.contrastive_loss(embeddings, knn_indices, neg_indices, anchor_indices)
         self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True, logger=True)
         return loss
 

docs/api/training/index.rst

@@ -166,9 +166,10 @@ pooling), use :class:`ContrastivePretrainer` directly.
 
 .. note::
 
-   Pairs are formed **within each batch**. For classification, a batch must
-   contain at least two classes; pretraining raises a ``ValueError`` if any
-   sample has no same-class or no different-class neighbor. Use a batch size
+   Pairs are formed **within each batch**. For classification, a sample acts as
+   an anchor only when its batch holds at least one same-class and one
+   different-class neighbor; other samples are skipped, and pretraining raises a
+   ``ValueError`` only when no sample in a batch qualifies. Use a batch size
    large enough to cover the classes, or a stratified sampler.
 
 See Also

tests/training/test_training_pretraining.py

@@ -238,6 +238,25 @@ def test_get_knn_classification_all_same_class_raises():
         pt.get_knn(labels)
 
 
+def test_class_singleton_is_dropped_from_anchors_not_from_pairs():
+    """A class with a single row in the batch has no positive neighbor: that row
+    is not used as an anchor, but it still serves as a negative for the other rows."""
+    pt = _make_pretrainer(regression=False, k_neighbors=4)
+    labels = torch.tensor([0, 0, 0, 1, 1, 1, 2])
+    knn_indices, neg_indices, anchor_indices = pt._pair_indices(labels)
+    assert anchor_indices.tolist() == [0, 1, 2, 3, 4, 5]
+    assert knn_indices.shape == neg_indices.shape == (6, 4)
+    # Each anchor has 4 different-class rows (3 of the other class and the singleton), so all are drawn.
+    assert all(6 in row for row in neg_indices.tolist())
+
+    loss = pt.contrastive_loss(torch.randn(7, 8), knn_indices, neg_indices, anchor_indices)
+    assert torch.isfinite(torch.as_tensor(loss))
+
+    unpooled = _make_pretrainer(regression=False, k_neighbors=2, pool_sequence=False)
+    sequence_loss = unpooled.contrastive_loss(torch.randn(7, 3, 8), knn_indices, neg_indices, anchor_indices)
+    assert torch.isfinite(torch.as_tensor(sequence_loss))
+
+
 # ---------------------------------------------------------------------------
 # ContrastivePretrainer init
 # ---------------------------------------------------------------------------
@@ -284,3 +303,36 @@ def test_contrastive_loss_anchor_layout_matches_pair_layout():
     anchors = embeddings.repeat_interleave(k, dim=0).repeat(n_pair_types, 1).view(-1)
     expected = torch.arange(n_anchors * k * n_pair_types) % (n_anchors * k) // k
     assert torch.equal(anchors, expected.float())
+
+
+def test_contrastive_loss_uses_the_given_anchor_rows():
+    """Self-positives and orthogonal negatives give zero loss only when the
+    neighbor rows are paired with the anchor rows they were selected for."""
+    pt = _make_pretrainer(k_neighbors=1, pool_sequence=True)
+    pt.loss_fn = nn.CosineEmbeddingLoss(margin=0.0)
+    anchor_indices = torch.tensor([1, 2, 4])
+    knn_indices = anchor_indices.view(3, 1)
+    neg_indices = torch.zeros(3, 1, dtype=torch.long)
+    loss = pt.contrastive_loss(torch.eye(5), knn_indices, neg_indices, anchor_indices)
+    assert float(loss) == pytest.approx(0.0, abs=1e-6)
+
+
+def test_classifier_pretraining_completes_with_a_class_singleton(tmp_path, monkeypatch):
+    import numpy as np
+    import pandas as pd
+
+    from deeptab.configs import FTTransformerConfig
+    from deeptab.models import FTTransformerClassifier
+
+    monkeypatch.chdir(tmp_path)
+    rng = np.random.default_rng(0)
+    features = pd.DataFrame({"a": rng.normal(size=48), "b": rng.normal(size=48)})
+    # Class 2 occurs once among the 32 training rows, so the batch holding it has a class singleton.
+    targets = np.array([0, 1] * 15 + [2, 0] + [0, 1] * 8)
+    model = FTTransformerClassifier(
+        model_config=FTTransformerConfig(d_model=8, n_layers=1, n_heads=2),
+        random_state=0,
+    )
+    model.build_model(features[:32], targets[:32], X_val=features[32:], y_val=targets[32:], batch_size=16)
+    model.pretrain(pretrain_epochs=1, k_neighbors=2, save_path=str(tmp_path / "pretrained.pth"))
+    assert (tmp_path / "pretrained.pth").is_file()