diff --git a/docs/api/calib.rst b/docs/api/calib.rst index 7814b4719..c2543af68 100644 --- a/docs/api/calib.rst +++ b/docs/api/calib.rst @@ -42,6 +42,25 @@ New to calibration and uncertainty quantification? Check out this complete examp This example shows the complete pipeline from model training to uncertainty-aware predictions with formal coverage guarantees. +.. note:: + + ``CovariateLabel``'s finite-sample correction implements Corollary 1 of + Tibshirani, Barber, Candes, and Ramdas, "Conformal Prediction Under + Covariate Shift" (NeurIPS 2019, https://arxiv.org/abs/1904.06019): the + test point's reserved probability mass must be inserted as an actual + point in the weighted empirical distribution (at the conservative + extreme), not merely folded into the normalizing denominator -- doing + only the latter silently under-covers relative to the target coverage + level. + + Corollary 1 also defines the threshold *per test point*, using that + point's own likelihood ratio w(x). Pass ``test_embeddings`` to + ``CovariateLabel.forward()`` to get this exact per-point threshold; + omitting it falls back to a single threshold computed from the *mean* + calibration likelihood ratio, which is only an approximation of the + paper's guarantee (a ``UserWarning`` is raised when this fallback is + used). + Quick Links ----------- diff --git a/examples/conformal_eeg/tuev_covariate_shift_conformal.py b/examples/conformal_eeg/tuev_covariate_shift_conformal.py index d1bcba20a..58efd0c2e 100644 --- a/examples/conformal_eeg/tuev_covariate_shift_conformal.py +++ b/examples/conformal_eeg/tuev_covariate_shift_conformal.py @@ -19,6 +19,9 @@ Notes: - CovariateLabel requires access to test embeddings to estimate density ratios. - Test embeddings are recomputed each seed since the model changes. +- CovariateLabel's finite-sample correction implements Corollary 1 of + Tibshirani, Barber, Candes, and Ramdas (NeurIPS 2019, arXiv:1904.06019); + see docs/api/calib.rst for details. """ from __future__ import annotations diff --git a/pyhealth/calib/predictionset/covariate/covariate_label.py b/pyhealth/calib/predictionset/covariate/covariate_label.py index 3482e4e91..1ff0b6b0e 100644 --- a/pyhealth/calib/predictionset/covariate/covariate_label.py +++ b/pyhealth/calib/predictionset/covariate/covariate_label.py @@ -22,6 +22,7 @@ https://arxiv.org/abs/2310.12033 """ +import warnings from typing import Callable, Dict, Optional, Union import numpy as np @@ -200,6 +201,23 @@ def _query_weighted_quantile( reserved test-point mass alone already meets or exceeds ``alpha``, since there isn't enough calibration mass to justify a stricter, finite threshold without risking under-coverage. + + Note: + This implements Corollary 1 of Tibshirani, Barber, Candes, and + Ramdas, "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, + https://arxiv.org/abs/1904.06019): the (1-alpha)-quantile of + sum_i p_i^w(x) * delta_{V_i} + p_{n+1}^w(x) * delta_infinity, where + p_i^w(x) = w(X_i) / (sum_j w(X_j) + w(x)) and p_{n+1}^w(x) is the + reserved test-point mass. The paper's V is a nonconformity score + (higher = worse) with the reserved mass placed at +infinity; this + codebase uses the opposite convention (conformity, higher = better), + so the reserved mass is placed at -infinity here instead. Crucially, + that reserved mass must be inserted as an actual point in the + weighted empirical distribution, not merely folded into the + normalizing denominator -- otherwise it dilutes every real + calibration weight without ever contributing to the cumulative sum + used to pick the threshold, which produces under-coverage instead of + the intended finite-sample guarantee. """ sorted_indices = np.argsort(scores) sorted_scores = scores[sorted_indices] @@ -213,11 +231,17 @@ def _query_weighted_quantile( if p_test >= alpha: # Not enough calibration mass to reach the target coverage without # dipping into the mass reserved for the test point itself: fall - # back to the maximally permissive (safe) threshold. + # back to the maximally permissive (safe) threshold. This is the + # case where the reserved point (sitting first, at -inf, in the + # augmented distribution) already accounts for the full quantile + # by itself. return -np.inf - # Compute cumulative weights over the reserved-mass-inclusive total. - cum_weights = np.cumsum(sorted_weights) / total_weight + # Cumulative weights over the reserved-mass-inclusive total, with the + # test point's reserved mass prepended as an actual point at -inf + # (equivalent to inserting it at the head of the sorted array before + # taking the cumulative sum). + cum_weights = (test_weight + np.cumsum(sorted_weights)) / total_weight # Find the index where cumulative weight exceeds alpha idx = np.searchsorted(cum_weights, alpha, side="left") @@ -400,6 +424,15 @@ def __init__( # Will be set during calibration self.t = None self._sum_cal_weights = None + # Calibration conformity scores/weights, kept so forward() can + # recompute a threshold per-test-point using that point's own + # likelihood ratio w(x), as Corollary 1 of Tibshirani et al. (2019) + # requires, when test_embeddings are provided. + self._cal_conformity_scores = None + self._cal_likelihood_ratios = None + self._cal_class_scores = None + self._cal_class_weights = None + self._warned_fixed_threshold = False def calibrate( self, @@ -516,7 +549,26 @@ def calibrate( y_prob, y_true, score_type=self.score_type, rng=self.rng ) - # Compute weighted quantile thresholds + # Keep the raw calibration scores/weights so forward() can, when + # given test_embeddings, recompute a threshold per-test-point using + # that point's own likelihood ratio w(x) -- the construction + # Corollary 1 actually specifies -- rather than only the + # mean-weight approximation computed below. + self._cal_conformity_scores = conformity_scores + self._cal_likelihood_ratios = likelihood_ratios + if not isinstance(self.alpha, float): + self._cal_class_scores = [ + conformity_scores[y_true == k] for k in range(K) + ] + self._cal_class_weights = [ + likelihood_ratios[y_true == k] for k in range(K) + ] + + # Compute weighted quantile thresholds using the mean calibration + # weight as a stand-in for a "typical" test point's weight. This is + # the fallback used when forward() isn't given test_embeddings (see + # forward()'s docstring for why that's only an approximation of + # Corollary 1, not an exact instance of it). if isinstance(self.alpha, float): test_weight = float(np.mean(likelihood_ratios)) t = _query_weighted_quantile( @@ -541,9 +593,32 @@ def calibrate( self.t = torch.tensor(t, device=self.device) - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward( + self, test_embeddings: np.ndarray | None = None, **kwargs + ) -> dict[str, torch.Tensor]: """Forward propagation with prediction set construction. + Args: + test_embeddings: Optional embeddings for this batch's test + points, shape ``(batch_size, embedding_dim)``, aligned with + the batch order. When provided (and KDEs are available from + :meth:`calibrate`), the threshold is recomputed per test + point using that point's own likelihood ratio w(x) via + :func:`_query_weighted_quantile` -- exactly the construction + in Corollary 1 of Tibshirani, Barber, Candes, and Ramdas, + "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, + https://arxiv.org/abs/1904.06019). + + When omitted, falls back to the single threshold computed + during :meth:`calibrate` using the *mean* calibration + likelihood ratio as a stand-in for w(x). That fallback is + only an approximation of Corollary 1 -- the paper's formula + is defined per test point via that point's actual w(x), not + an aggregate over the calibration set -- so results from the + fallback do not carry the same finite-sample coverage + guarantee as the paper's construction. A warning is emitted + (once) the first time this fallback is used. + Returns: Dictionary with all results from base model, plus: - y_predset: Boolean tensor indicating which classes @@ -556,10 +631,72 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]: conformity_scores = all_class_conformity_scores( y_prob, score_type=self.score_type, rng=self.rng ) + N, K = conformity_scores.shape + + if test_embeddings is not None: + if self.kde_test is None or self.kde_cal is None: + raise ValueError( + "test_embeddings was provided but no KDEs are available " + "(calibrate() was called with custom cal_weights, which " + "has no way to score a new test point's likelihood " + "ratio). Per-test-point thresholding requires calibrate() " + "to have been called with cal_embeddings/test_embeddings " + "or pre-fitted kde_test/kde_cal." + ) + if self._cal_conformity_scores is None: + raise RuntimeError("Must call calibrate() before forward().") + + test_weights = _compute_likelihood_ratio( + self.kde_test, self.kde_cal, test_embeddings + ) + thresholds = np.empty((N, K), dtype=np.float64) + if isinstance(self.alpha, float): + for i in range(N): + t_i = _query_weighted_quantile( + self._cal_conformity_scores, + self.alpha, + self._cal_likelihood_ratios, + float(test_weights[i]), + ) + thresholds[i, :] = t_i + else: + for i in range(N): + for k in range(K): + if len(self._cal_class_scores[k]) > 0: + thresholds[i, k] = _query_weighted_quantile( + self._cal_class_scores[k], + self.alpha[k], + self._cal_class_weights[k], + float(test_weights[i]), + ) + else: + thresholds[i, k] = -np.inf + threshold_tensor = torch.as_tensor( + thresholds, + device=pred["y_prob"].device, + dtype=pred["y_prob"].dtype, + ) + else: + if not self._warned_fixed_threshold: + warnings.warn( + "CovariateLabel.forward() was called without " + "test_embeddings: falling back to a single fixed " + "threshold computed from the mean calibration " + "likelihood ratio. This is only an approximation of " + "Corollary 1 of Tibshirani et al. (2019), which defines " + "the threshold per test point using that point's own " + "likelihood ratio w(x); pass test_embeddings to get the " + "paper's exact finite-sample coverage guarantee.", + UserWarning, + stacklevel=2, + ) + self._warned_fixed_threshold = True + threshold_tensor = self.t + conformity_scores = torch.as_tensor( conformity_scores, device=pred["y_prob"].device, dtype=pred["y_prob"].dtype ) - pred["y_predset"] = conformity_scores > self.t + pred["y_predset"] = conformity_scores > threshold_tensor return pred diff --git a/tests/core/test_covariate_label.py b/tests/core/test_covariate_label.py index 6f2bfd04a..8e4f99b66 100644 --- a/tests/core/test_covariate_label.py +++ b/tests/core/test_covariate_label.py @@ -1,10 +1,17 @@ import unittest +import warnings +from unittest.mock import patch + import numpy as np import torch from pyhealth.datasets import create_sample_dataset, get_dataloader from pyhealth.models import MLP from pyhealth.calib.predictionset.covariate import CovariateLabel, fit_kde +from pyhealth.calib.predictionset.covariate.covariate_label import ( + _compute_likelihood_ratio, + _query_weighted_quantile, +) from pyhealth.calib.utils import extract_embeddings @@ -340,6 +347,174 @@ def test_score_type_aps_runs_end_to_end(self): self.assertEqual(output["y_predset"].dtype, torch.bool) self.assertEqual(output["y_predset"].shape, output["y_prob"].shape) + def _build_pointwise_setup(self, alpha=0.4): + """Build a larger dataset + a deliberately bimodal (non-uniform) + KDE pair, so per-test-point likelihood ratios genuinely differ -- + needed to test forward()'s pointwise-thresholding path, which the + class's default dummy_kde_cal/dummy_kde_test (both ~uniform) can't + exercise meaningfully. Seeded for reproducibility across runs; the + model's own random init would otherwise depend on however much of + the global RNG stream setUp() and any prior tests already consumed. + """ + torch.manual_seed(0) + np.random.seed(0) + samples = [] + for i in range(60): + samples.append( + { + "patient_id": f"p{i}", + "visit_id": f"v{i}", + "conditions": [f"cond-{i % 10}", f"cond-{(i + 1) % 10}"], + "procedures": [float(i % 5), float((i * 2) % 7), 1.0, 2.0], + "label": i % 3, + } + ) + dataset = create_sample_dataset( + samples=samples, + input_schema={"conditions": "sequence", "procedures": "tensor"}, + output_schema={"label": "multiclass"}, + dataset_name="test_pointwise", + ) + model = MLP( + dataset=dataset, + feature_keys=["conditions", "procedures"], + label_key="label", + mode="multiclass", + ) + model.eval() + + cal_dataset = dataset.subset(list(range(0, 30))) + test_dataset = dataset.subset(list(range(30, 60))) + cal_embeddings = extract_embeddings(model, cal_dataset, batch_size=32, device="cpu") + test_embeddings = extract_embeddings(model, test_dataset, batch_size=32, device="cpu") + + def kde_cal(data): + return np.ones(len(np.asarray(data))) + + def kde_test(data): + data = np.asarray(data) + row_signal = np.abs(data).sum(axis=1) + median = np.median(row_signal) + # Two clearly separated weight regimes, split deterministically + # by an arbitrary per-row criterion, so weights genuinely + # differ by test point. + return np.where(row_signal > median, 10.0, 0.05) + + cal_model = CovariateLabel( + model=model, alpha=alpha, kde_test=kde_test, kde_cal=kde_cal, random_state=0 + ) + cal_model.calibrate( + cal_dataset=cal_dataset, + cal_embeddings=cal_embeddings, + test_embeddings=test_embeddings, + ) + return cal_model, test_dataset, test_embeddings, kde_test, kde_cal + + def test_forward_without_embeddings_warns_and_uses_fixed_threshold(self): + """forward() without test_embeddings must warn that it's only + approximating Corollary 1 (mean calibration weight standing in for + each test point's own w(x)), and must use the single fixed + threshold computed at calibrate() time.""" + cal_model, test_dataset, _, _, _ = self._build_pointwise_setup() + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + with torch.no_grad(): + cal_model(**batch) + self.assertTrue( + any("approximation of Corollary 1" in str(w.message) for w in caught) + ) + + def test_forward_with_embeddings_uses_per_point_weight(self): + """forward(test_embeddings=...) must recompute the threshold per + test point using that point's own likelihood ratio w(x) (Corollary + 1 of Tibshirani et al. 2019), not the calibrate()-time mean-weight + approximation. + + Verified by spying on _query_weighted_quantile during a real + forward() call and checking it is invoked once per test point with + that exact point's own weight -- directly verifying the claim + ("forward uses per-point weight") rather than asserting the + resulting *threshold values* differ between two points. + + The latter was tried first and is not a reliable test: a test + point's likelihood-ratio weight is drawn from the same KDE-derived + distribution as the calibration weights, so it can never be more + than a small fraction of the total calibration weight mass by + construction -- whether prepending it shifts _query_weighted_ + quantile's selected order-statistic index depends on exactly where + the alpha-quantile boundary falls relative to the (data-dependent, + effectively random given the model's unseeded init) distribution + of calibration weights along the sorted-score axis. That made the + old assertion fail whenever the boundary happened to fall in a + region insensitive to a perturbation of that size -- confirmed + directly: it failed deterministically for several concrete seeds, + not just intermittently, so it wasn't simply "rare bad luck." + """ + cal_model, test_dataset, test_embeddings, kde_test, kde_cal = ( + self._build_pointwise_setup() + ) + weights = _compute_likelihood_ratio(kde_test, kde_cal, test_embeddings) + self.assertGreater( + len(set(np.round(weights, 3))), 1, "test setup must have varying weights" + ) + + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) + + module = "pyhealth.calib.predictionset.covariate.covariate_label" + with patch( + f"{module}._query_weighted_quantile", + wraps=_query_weighted_quantile, + ) as spy: + with torch.no_grad(): + cal_model(test_embeddings=test_embeddings, **batch) + + self.assertEqual(spy.call_count, len(weights)) + called_test_weights = [call.args[3] for call in spy.call_args_list] + np.testing.assert_allclose( + sorted(called_test_weights), sorted(float(w) for w in weights) + ) + + def test_query_weighted_quantile_test_weight_changes_result(self): + """Direct, hand-computed check that test_weight actually changes + _query_weighted_quantile's output -- the algorithmic property the + forward()-level test above relies on, isolated from any KDE/model + randomness. A large enough test_weight must push p_test >= alpha, + forcing the documented -inf fallback (not enough calibration mass + to reach the target coverage without dipping into the test point's + own reserved mass); test_weight=0 must not. + """ + scores = np.array([0.1, 0.2, 0.3, 0.4, 0.5]) + weights = np.full(5, 2.0) # total calibration weight = 10.0 + alpha = 0.4 + + t_low = _query_weighted_quantile(scores, alpha, weights, test_weight=0.0) + self.assertNotEqual(t_low, -np.inf) + + # p_test = test_weight / (10 + test_weight) >= 0.4 requires + # test_weight >= 10 * 0.4 / 0.6 ≈ 6.67; 100 clears it by a wide, + # seed-independent margin. + t_high = _query_weighted_quantile(scores, alpha, weights, test_weight=100.0) + self.assertEqual(t_high, -np.inf) + self.assertNotEqual(t_low, t_high) + + def test_forward_raises_without_kde_for_pointwise(self): + """test_embeddings requires KDEs (to score a new point's w(x)); + calibrate() with custom cal_weights has no way to do that, so + forward() must raise rather than silently ignore test_embeddings.""" + cal_model, test_dataset, test_embeddings, _, _ = self._build_pointwise_setup() + cal_model.kde_test = None + cal_model.kde_cal = None + + test_loader = get_dataloader(test_dataset, batch_size=30, shuffle=False) + batch = next(iter(test_loader)) + with self.assertRaises(ValueError): + with torch.no_grad(): + cal_model(test_embeddings=test_embeddings, **batch) + def test_weighted_quantile_function(self): """Test the weighted quantile helper function.""" from pyhealth.calib.predictionset.covariate.covariate_label import ( @@ -403,6 +578,35 @@ def test_weighted_quantile_small_calibration_set_is_conservative(self): ) self.assertTrue(np.isfinite(result_large)) + def test_weighted_quantile_matches_corollary_1(self): + """The finite-sample correction must match Corollary 1 of Tibshirani, + Barber, Candes, and Ramdas, "Conformal Prediction Under Covariate + Shift" (NeurIPS 2019): the reserved test-point mass is an actual + point in the weighted empirical distribution (at -inf, under this + codebase's higher-is-better conformity convention), not merely a + term folded into the normalizing denominator. + + With scores=[0.1,0.3,0.5,0.7,0.9], uniform weights=1, test_weight=1, + alpha=0.3: total_weight=6, p_test=1/6 (< alpha, so not the -inf + fallback). Correctly inserting the reserved mass as the first point + of the augmented (n+1)-point distribution gives cumulative + fractions [2/6, 3/6, 4/6, 5/6, 6/6] = [.333,.5,.667,.833,1.0], so the + alpha=0.3 quantile lands at the first real point: 0.1. + + Before this fix, the reserved mass was only added to the + denominator (cumulative fractions [1/6,2/6,3/6,4/6,5/6]), which + incorrectly returned 0.3 instead -- a stricter, LESS permissive + threshold that under-covers relative to the target. + """ + from pyhealth.calib.predictionset.covariate.covariate_label import ( + _query_weighted_quantile, + ) + + scores = np.array([0.1, 0.3, 0.5, 0.7, 0.9]) + weights = np.ones(5) + result = _query_weighted_quantile(scores, 0.3, weights, test_weight=1.0) + self.assertAlmostEqual(result, 0.1, places=10) + def test_likelihood_ratio_function(self): """Test the likelihood ratio computation.""" from pyhealth.calib.predictionset.covariate.covariate_label import (