diff --git a/docs/api/calib/pyhealth.calib.predictionset.rst b/docs/api/calib/pyhealth.calib.predictionset.rst index 740b1c87c..0f0708d83 100644 --- a/docs/api/calib/pyhealth.calib.predictionset.rst +++ b/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: diff --git a/examples/conformal_eeg/tuev_ncp_conformal.py b/examples/conformal_eeg/tuev_ncp_conformal.py index 77cc98475..b4a1b2543 100644 --- a/examples/conformal_eeg/tuev_ncp_conformal.py +++ b/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 diff --git a/pyhealth/calib/predictionset/cluster/neighborhood_label.py b/pyhealth/calib/predictionset/cluster/neighborhood_label.py index 36fdbaa0a..c8426a348 100644 --- a/pyhealth/calib/predictionset/cluster/neighborhood_label.py +++ b/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) diff --git a/tests/core/test_neighborhood_label.py b/tests/core/test_neighborhood_label.py index 0c812d141..2986ed50d 100644 --- a/tests/core/test_neighborhood_label.py +++ b/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()