Fix calibration self-inclusion leak in NeighborhoodLabel (NCP)

docs/api/calib/pyhealth.calib.predictionset.rst

@@ -81,6 +81,19 @@ ClusterLabel (K-means Cluster-based Conformal)
 NeighborhoodLabel (Neighborhood Conformal Prediction)
 -----------------------------------------------------
 
+.. note::
+
+   ``NeighborhoodLabel.calibrate()`` excludes each calibration point from
+   its own k-nearest-neighbor set during the alpha_tilde search.
+   ``sklearn``'s ``NearestNeighbors.kneighbors()`` returns each point as
+   its own nearest neighbor (distance 0) when queried with an explicit
+   ``X`` equal to the fitted set -- unlike the implicit no-argument form,
+   which excludes self-matches. Without excluding this self-match, a
+   calibration point's own score leaks into its own threshold during
+   calibration (something a genuine test point never benefits from),
+   biasing the search toward an overly permissive threshold and causing
+   real under-coverage at test time.
+
 .. autoclass:: pyhealth.calib.predictionset.NeighborhoodLabel
    :members:
    :undoc-members:

examples/conformal_eeg/tuev_ncp_conformal.py

@@ -15,6 +15,11 @@
   python examples/conformal_eeg/tuev_ncp_conformal.py --root /srv/local/data/TUH/tuh_eeg_events/v2.0.0/edf
   python examples/conformal_eeg/tuev_ncp_conformal.py --quick-test --log-file quicktest_ncp.log
   python examples/conformal_eeg/tuev_ncp_conformal.py --alpha 0.1 --n-seeds 5 --split-seed 0 --log-file ncp_seeds5.log
+
+Note: NeighborhoodLabel.calibrate() excludes each calibration point from
+its own k-nearest-neighbor set during the alpha_tilde search, to avoid a
+calibration point's own score leaking into its own threshold; see
+docs/api/calib/pyhealth.calib.predictionset.rst for details.
 """
 
 from __future__ import annotations

pyhealth/calib/predictionset/cluster/neighborhood_label.py

@@ -29,7 +29,35 @@ class NeighborhoodLabel(SetPredictor):
     Reference:
         Ghosh, S., Belkhouja, T., Yan, Y., & Doppa, J. R. (2023).
         Improving Uncertainty Quantification of Deep Classifiers via
-        Neighborhood Conformal Prediction.
+        Neighborhood Conformal Prediction. https://arxiv.org/abs/2303.10694
+
+    Note:
+        This implements Algorithm 2 of the paper (marginal coverage only --
+        the paper does not provide a class-conditional variant). Two
+        deviations from a literal reading of the paper, both deliberate:
+
+        1. The paper's k-NN weighting (Eq. 5) does not state whether a
+           calibration point's own score should be excluded from its own
+           neighbor set when searching for alpha_tilde (Eq. 2). This
+           implementation excludes it (leave-one-out): including it would
+           let a point's own score dominate its own threshold during
+           calibration -- something a genuine test point, never part of
+           the calibration set, cannot benefit from -- which biases the
+           search toward an overly permissive threshold and causes real
+           under-coverage at test time.
+        2. If a prediction set would come out empty (no class's score
+           reaches the threshold), the model's top-1 predicted class is
+           force-included. This is not specified in the paper, but is
+           coverage-safe: adding a class to a set can only raise coverage,
+           never lower it.
+
+        Like standard split-conformal prediction, the coverage guarantee
+        here assumes exchangeability between the calibration and test
+        distributions. It does not correct for covariate shift; see
+        :class:`~pyhealth.calib.predictionset.CovariateLabel` for a method
+        that does. Under real distribution shift (common in clinical data),
+        expect under-coverage even with a fully correct implementation --
+        that is an assumption violation, not an implementation bug.
 
     Args:
         model: A trained base model that supports embedding extraction
@@ -184,10 +212,28 @@ def calibrate(
         self.cal_embeddings_ = np.atleast_2d(cal_embeddings)
         self.cal_conformity_scores_ = np.asarray(conformity_scores, dtype=np.float64)
 
-       # this is the ncp calibration step
-        distances_cal, indices_cal = self._nn.kneighbors(
-            self.cal_embeddings_, n_neighbors=k
+        # This is the NCP calibration step. Querying kneighbors() with an
+        # explicit X argument equal to the fitted set makes each point its
+        # own nearest neighbor at distance 0 (unlike the implicit no-argument
+        # form, which sklearn special-cases to exclude self-matches) -- so we
+        # request one extra neighbor and drop each point's own self-match
+        # (leave-one-out). Without this, a calibration point's own score
+        # leaks into its own threshold computation during the alpha_tilde
+        # search below, which a genuine test point -- never part of the
+        # calibration set -- can't benefit from; that asymmetry biases the
+        # search toward an overly permissive threshold and causes
+        # under-coverage at test time.
+        k_query = min(k + 1, N)
+        distances_all, indices_all = self._nn.kneighbors(
+            self.cal_embeddings_, n_neighbors=k_query
         )
+        n_neighbors_loo = k_query - 1
+        distances_cal = np.zeros((N, n_neighbors_loo), dtype=distances_all.dtype)
+        indices_cal = np.zeros((N, n_neighbors_loo), dtype=indices_all.dtype)
+        for i in range(N):
+            mask = indices_all[i] != i
+            distances_cal[i] = distances_all[i][mask][:n_neighbors_loo]
+            indices_cal[i] = indices_all[i][mask][:n_neighbors_loo]
         cal_weights = np.exp(-distances_cal / self.lambda_L)
         cal_weights = cal_weights / cal_weights.sum(axis=1, keepdims=True)
 

tests/core/test_neighborhood_label.py

@@ -162,13 +162,97 @@ def test_calibrate_without_embeddings_extracts(self):
         self.assertIsNotNone(ncp.cal_conformity_scores_)
 
     def test_calibration_empirical_coverage_at_least_1_minus_alpha(self):
-        """After calibrate(), empirical coverage on calibration set >= 1-alpha."""
+        """After calibrate(), empirical coverage on calibration set >= 1-alpha,
+        recomputed the same leave-one-out way calibrate() itself does (a
+        calibration point's own score must never appear in its own
+        neighbor set -- see test_calibrate_excludes_self_from_neighbors for
+        why: querying kneighbors() with an explicit X equal to the fitted
+        set makes each point its own nearest neighbor at distance 0, which
+        would leak a point's own score into its own threshold and trivially
+        inflate this exact coverage check if not excluded).
+
+        Uses its own seeded, larger (N=30), *trained* calibration set rather
+        than the shared 6-sample untrained fixture. Two independent issues
+        made the original version flaky/wrong, not just one:
+
+        1. With only 6 points, achievable coverage values are exactly
+           {0, 1/6, ..., 1}, so the target 1-alpha=0.8 sits squarely between
+           4/6=0.667 and 5/6=0.833 -- a single point tipping either way
+           (driven only by the model's unseeded random init, since setUp()
+           never seeds torch) flips the assertion.
+        2. More fundamentally: the shared fixture's model is never trained,
+           so its predicted probabilities -- and therefore the conformity
+           scores NCP calibrates on -- are just noise from a random init,
+           uncorrelated with embedding-space locality. NCP's per-point
+           threshold comes from each point's k-nearest *neighbors only*
+           (itself excluded), not a global quantile over all N points the
+           way plain split conformal's is, so it has no automatic "any
+           alpha_tilde >= 0 must reach 1-alpha coverage" guarantee the way
+           a global quantile would -- that guarantee is a property of the
+           *underlying scores actually correlating with locality*, which an
+           untrained model doesn't provide. Confirmed directly: with the
+           untrained model, alpha_tilde_ converges to its floor of 0.0 (the
+           search's most permissive setting) and still only covers 23/30,
+           because roughly 1/(k+1) of points have a lower score than all k
+           of their neighbors purely by chance when scores are pure noise.
+           Training the model so scores genuinely reflect how "easy" each
+           point is relative to its neighborhood removes this floor issue.
+
+        N=30 gives a 3.3%-wide step size (fine enough not to sit on a knife
+        edge), a fixed seed makes the outcome reproducible, and training
+        makes the target coverage actually achievable by the method's own
+        theory rather than accidentally testing an unmet precondition.
+        """
         from pyhealth.calib.predictionset.base_conformal import _query_weighted_quantile
 
-        ncp = NeighborhoodLabel(model=self.model, alpha=0.2, k_neighbors=3, lambda_L=50.0)
-        cal_indices = [0, 1, 2, 3, 4, 5]
-        cal_dataset = self.dataset.subset(cal_indices)
-        cal_emb = self._get_embeddings(cal_dataset)
+        torch.manual_seed(0)
+        np.random.seed(0)
+
+        n_per_class = 10
+        samples = []
+        for label in range(3):
+            for j in range(n_per_class):
+                idx = label * n_per_class + j
+                samples.append(
+                    {
+                        "patient_id": f"p{idx}",
+                        "visit_id": f"v{idx}",
+                        "conditions": [f"c{idx}"],
+                        "procedures": [float(label) + 0.01 * j],
+                        "label": label,
+                    }
+                )
+        dataset = create_sample_dataset(
+            samples=samples,
+            input_schema=self.input_schema,
+            output_schema=self.output_schema,
+            dataset_name="test_coverage",
+        )
+        model = MLP(
+            dataset=dataset,
+            feature_keys=["conditions", "procedures"],
+            label_key="label",
+            mode="multiclass",
+        )
+
+        # Brief training: this trivially class-separable fixture (procedures
+        # cleanly bucketed by label) only needs a handful of epochs for
+        # predicted probabilities -- and thus conformity scores -- to
+        # reflect real structure instead of random-init noise.
+        model.train()
+        optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
+        train_loader = get_dataloader(dataset, batch_size=len(samples), shuffle=True)
+        for _ in range(50):
+            for batch in train_loader:
+                optimizer.zero_grad()
+                ret = model(**batch)
+                ret["loss"].backward()
+                optimizer.step()
+        model.eval()
+
+        ncp = NeighborhoodLabel(model=model, alpha=0.2, k_neighbors=3, lambda_L=50.0)
+        cal_dataset = dataset.subset(list(range(len(samples))))
+        cal_emb = extract_embeddings(model, cal_dataset, batch_size=32, device="cpu")
         ncp.calibrate(cal_dataset=cal_dataset, cal_embeddings=cal_emb)
 
         self.assertIsNotNone(ncp.alpha_tilde_)
@@ -178,9 +262,17 @@ def test_calibration_empirical_coverage_at_least_1_minus_alpha(self):
         # Recompute per-sample thresholds using alpha_tilde (Q^NCP definition: alpha_tilde-quantile of conformity)
         N = ncp.cal_conformity_scores_.shape[0]
         k = min(ncp.k_neighbors, N)
-        distances_cal, indices_cal = ncp._nn.kneighbors(
-            ncp.cal_embeddings_, n_neighbors=k
+        k_query = min(k + 1, N)
+        distances_all, indices_all = ncp._nn.kneighbors(
+            ncp.cal_embeddings_, n_neighbors=k_query
         )
+        n_loo = k_query - 1
+        distances_cal = np.zeros((N, n_loo))
+        indices_cal = np.zeros((N, n_loo), dtype=int)
+        for i in range(N):
+            mask = indices_all[i] != i
+            distances_cal[i] = distances_all[i][mask][:n_loo]
+            indices_cal[i] = indices_all[i][mask][:n_loo]
         cal_weights = np.exp(-distances_cal / ncp.lambda_L)
         cal_weights = cal_weights / cal_weights.sum(axis=1, keepdims=True)
 
@@ -201,6 +293,42 @@ def test_calibration_empirical_coverage_at_least_1_minus_alpha(self):
             msg=f"Calibration empirical coverage {empirical_coverage:.4f} should be >= 1-alpha={1 - ncp.alpha}",
         )
 
+    def test_calibrate_excludes_self_from_neighbors(self):
+        """Regression test: calibrate()'s alpha_tilde search must not let a
+        calibration point be its own nearest neighbor.
+
+        Querying sklearn's NearestNeighbors.kneighbors() with an explicit X
+        argument equal to the fitted set returns each point as its own
+        nearest neighbor at distance 0 (unlike the implicit no-argument
+        form, which sklearn special-cases to exclude self-matches). Without
+        excluding this self-match, a calibration point's own score leaks
+        into its own threshold computation during calibration -- something
+        a genuine test point (never part of the calibration set) can't
+        benefit from -- which biases the alpha_tilde search toward an
+        overly permissive threshold and causes real under-coverage at test
+        time (empirically ~0.82-0.83 actual vs. 0.90 target in isolated
+        simulation, before this fix).
+        """
+        ncp = NeighborhoodLabel(model=self.model, alpha=0.2, k_neighbors=3, lambda_L=50.0)
+        cal_dataset = self.dataset.subset([0, 1, 2, 3, 4, 5])
+        cal_emb = self._get_embeddings(cal_dataset)
+        ncp.calibrate(cal_dataset=cal_dataset, cal_embeddings=cal_emb)
+
+        N = ncp.cal_conformity_scores_.shape[0]
+        k = min(ncp.k_neighbors, N)
+        k_query = min(k + 1, N)
+        _, indices_all = ncp._nn.kneighbors(ncp.cal_embeddings_, n_neighbors=k_query)
+
+        for i in range(N):
+            with self.subTest(point=i):
+                mask = indices_all[i] != i
+                kept = indices_all[i][mask][: k_query - 1]
+                self.assertNotIn(
+                    i,
+                    kept,
+                    f"calibration point {i} leaked into its own neighbor set",
+                )
+
 
 if __name__ == "__main__":
     unittest.main()