Fix under-coverage in CovariateLabel's finite-sample correction - #1223
Open
lehendo wants to merge 4 commits into
Open
Fix under-coverage in CovariateLabel's finite-sample correction#1223lehendo wants to merge 4 commits into
lehendo wants to merge 4 commits into
Conversation
_query_weighted_quantile added the test point's reserved weight only to the normalizing denominator, never inserting it as an actual point in the weighted empirical distribution. Verified against Corollary 1 of Tibshirani, Barber, Candes, and Ramdas, "Conformal Prediction Under Covariate Shift" (NeurIPS 2019, arXiv:1904.06019): before this fix, the "corrected" quantile under-covered relative to target (0.81-0.89 vs a 0.90 target in Monte Carlo simulation) -- worse than no correction at all. Fix: prepend the reserved weight to the cumulative sum before dividing, matching the paper's construction exactly (confirmed via a hand-computed example and against a from-scratch reference implementation of the paper's formula). Also fix a related gap: calibrate() computed a single threshold using the *mean* calibration likelihood ratio as a stand-in for a test point's own w(x), but Corollary 1 defines the threshold per test point using that point's actual weight. forward() now accepts an optional test_embeddings argument to compute the real per-point threshold matching the paper exactly; omitting it keeps the old single-threshold behavior as a documented, explicitly-warned approximation, since not every wrapped model exposes an embedding extraction path. 18 tests pass (4 new), including a deterministic regression test against a hand-computed Corollary 1 example and new coverage of the forward() per-point/fallback paths.
lehendo
changed the title
CI failed on test_forward_with_embeddings_uses_per_point_weight
(AssertionError: 0.3241... == 0.3241...). Root-caused this to a real
design flaw, not just missing randomness control: a test point's
likelihood-ratio weight is drawn from the same KDE-derived distribution
as the calibration weights, so by construction it can never be more than
a small fraction of total calibration weight mass. Whether prepending it
shifts _query_weighted_quantile's selected order-statistic index depends
entirely on where the alpha-quantile boundary happens to fall relative to
the (data-dependent, effectively random given the model's unseeded init)
distribution of calibration weights along the sorted-score axis.
Confirmed this wasn't just 'rare bad luck': reproduced the exact failure
deterministically for multiple fixed seeds, and directly inspected the
internals (total_weight, cum_weights) showing the two test points'
weights (0.05 vs 10.0) simply didn't straddle a boundary for that data --
an inherent fragility in comparing *output values* for two randomly-KDE
-derived weights, not a bug in the underlying Corollary-1 implementation.
Replaced the flaky output-comparison with two more precise checks:
1. Spy on _query_weighted_quantile during a real forward() call and
verify it's invoked once per test point with that exact point's own
weight -- directly verifies the claim ('forward uses per-point
weight') without depending on the resulting threshold *values*
differing.
2. A new, fully hand-computed unit test proving test_weight does change
_query_weighted_quantile's result in general (by forcing the
documented -inf fallback with a deliberately large weight), isolated
from any KDE/model randomness.
Also seeded _build_pointwise_setup()'s model init for reproducibility.
lehendo
changed the title
Contributor
|
The threshold should be calculated separately for each test patient, using that patient's embedding. However, I think put the embeddings in each batch or compute them inside |
Contributor
|
Another boundary issue -- Corollary 1 includes scores equal to the threshold. Since larger scores should be included, this should use >= instead of >.. Tied scores at the threshold are excluded and can under-cover. Switch it to |
Two independent bugs, both flagged in review: 1. forward() included a class only if conformity_score > threshold (strict). _query_weighted_quantile defines the threshold as the smallest score whose cumulative weight reaches alpha, so that score's own mass is part of what clears the target -- excluding it with a strict comparison silently under-covers whenever a real score lands on a calibration tie. scores.py's own all_class_conformity_scores docstring already documents ">=" as the intended convention for this sign convention; forward() just didn't follow it. Fixed to ">=", and added a length-mismatch check for test_embeddings (previously extra rows were silently ignored rather than raising). 2. The per-test-point thresholding path (forward(test_embeddings=...), the actual Corollary 1 construction from Tibshirani et al. 2019) has existed since the previous commit on this branch, but examples/conformal_eeg/tuev_covariate_shift_conformal.py evaluated it via Trainer(model=cov_predictor).inference(test_loader, ...), which calls model(**batch) with no way to also supply a per-batch test_embeddings slice. So the example -- meant to demonstrate and validate this exact PR's fix -- silently only ever exercised the approximate mean-weight fallback, never the real per-point guarantee. Replaced with a manual inference loop that slices test_embeddings in the same fixed order test_loader yields batches in (mirroring Trainer.inference()'s own logic) and passes each batch's slice through. Added a regression test for (1) that directly forces a conformity score equal to the threshold and checks it's included.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
test_forward_with_embeddings_uses_per_point_weight failed on CI (0.3241... == 0.3241...). Root cause: a test point's KDE-derived weight is always a small fraction of total calibration weight, so whether it shifts the selected quantile index is essentially a coin flip depending on random model init