Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improving Coverage in Combined Prediction Sets with Weighted p-values

Gina Wong, Drew Prinster, Suchi Saria, Rama Chellappa, Anqi Liu

TL;DR

The paper tackles the challenge that aggregating conformal prediction sets from multiple sources typically yields a weaker overall guarantee of $1-2\alpha$ despite per-model validity. It introduces a weighted p-value framework to aggregate conformal p-values, with both data-independent and data-dependent weights, and proves finite-sample validity via a scaling correction $m^*$ that preserves weight proportions. This framework enables adaptive, input-specific coverage, demonstrated in mixture-of-experts (MoE) settings where routing weights modulate each expert’s influence on the final prediction set. Empirical results show improved worst-slice coverage and more uniform validity across data regions, highlighting practical benefits for local reliability in complex, multi-source prediction tasks.

Abstract

Conformal prediction quantifies the uncertainty of machine learning models by augmenting point predictions with valid prediction sets. For complex scenarios involving multiple trials, models, or data sources, conformal prediction sets can be aggregated to create a prediction set that captures the overall uncertainty, often improving precision. However, aggregating multiple prediction sets with individual $1-α$ coverage inevitably weakens the overall guarantee, typically resulting in $1-2α$ worst-case coverage. In this work, we propose a framework for the weighted aggregation of prediction sets, where weights are assigned to each prediction set based on their contribution. Our framework offers flexible control over how the sets are aggregated, achieving tighter coverage bounds that interpolate between the $1-2α$ guarantee of the combined models and the $1-α$ guarantee of an individual model depending on the distribution of weights. Importantly, our framework generalizes to data-dependent weights, as we derive a procedure for weighted aggregation that maintains finite-sample validity even when the weights depend on the data. This extension makes our framework broadly applicable to settings where weights are learned, such as mixture-of-experts (MoE), and we demonstrate through experiments in the MoE setting that our methods achieve adaptive coverage.

Improving Coverage in Combined Prediction Sets with Weighted p-values

TL;DR

The paper tackles the challenge that aggregating conformal prediction sets from multiple sources typically yields a weaker overall guarantee of despite per-model validity. It introduces a weighted p-value framework to aggregate conformal p-values, with both data-independent and data-dependent weights, and proves finite-sample validity via a scaling correction that preserves weight proportions. This framework enables adaptive, input-specific coverage, demonstrated in mixture-of-experts (MoE) settings where routing weights modulate each expert’s influence on the final prediction set. Empirical results show improved worst-slice coverage and more uniform validity across data regions, highlighting practical benefits for local reliability in complex, multi-source prediction tasks.

Abstract

Conformal prediction quantifies the uncertainty of machine learning models by augmenting point predictions with valid prediction sets. For complex scenarios involving multiple trials, models, or data sources, conformal prediction sets can be aggregated to create a prediction set that captures the overall uncertainty, often improving precision. However, aggregating multiple prediction sets with individual coverage inevitably weakens the overall guarantee, typically resulting in worst-case coverage. In this work, we propose a framework for the weighted aggregation of prediction sets, where weights are assigned to each prediction set based on their contribution. Our framework offers flexible control over how the sets are aggregated, achieving tighter coverage bounds that interpolate between the guarantee of the combined models and the guarantee of an individual model depending on the distribution of weights. Importantly, our framework generalizes to data-dependent weights, as we derive a procedure for weighted aggregation that maintains finite-sample validity even when the weights depend on the data. This extension makes our framework broadly applicable to settings where weights are learned, such as mixture-of-experts (MoE), and we demonstrate through experiments in the MoE setting that our methods achieve adaptive coverage.
Paper Structure (72 sections, 3 theorems, 81 equations, 14 figures, 6 tables, 1 algorithm)

This paper contains 72 sections, 3 theorems, 81 equations, 14 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Quantile and the empirical quantile function
  4. p-values
  5. Conformal prediction
  6. Related work
  7. Combining conformal prediction sets by weighted p-value
  8. Weighted aggregation provides more flexible coverage guarantees
  9. From p-values to prediction sets
  10. Extending to data-dependent weights
  11. Computing $m^*$ with a merging set
  12. Data-dependent weights give a form of conditional coverage
  13. Achieving tighter guarantees
  14. Application: mixture-of-experts
  15. Baselines
  16. ...and 57 more sections

Key Result

Proposition 4.0

Let $\widehat{C}_{\alpha,1}^{\text{split}}(X_{n+1}), \dots, \widehat{C}_{\alpha,K}^{\text{split}}(X_{n+1})$ be $K$ prediction sets defined by p-value functions $\widehat{p}_1, \dots, \widehat{p}_K$eq:pk on $X_{n+1}$, where $1-\alpha$ coverage eq:split_coverage holds for each set $k \in [K]$. Then, t where $v = \max\{v_1, v_2, \dots, v_K\}$ is the largest weight assigned to any of the p-values.

Figures (14)

  • Figure 1: Left: Storm forecasting example with different models tracking humidity, sea temperature, rainfall, and wind shear over time. Below, an abstract representation of how models vary in predictive strength across the input space (colored regions). At the given test point (black), the red and green models dominate, so their prediction sets matter most. Right: Model prediction sets are combined with weighted aggregation. (a) Data-independent weights reflect expert priors (e.g. the blue model covers the largest portion of the input space and so is considered more generally reliable, and up-weighted accordingly). (b) Data-dependent weights adapt to context, yielding forecasts better aligned with current conditions (e.g. at the test point, red and green dominate).
  • Figure 2: Left: Network diagram for MoE. For traditional MoE, the aggregation module takes a weighted sum of the outputs from each expert. To learn weight-dependent prediction sets, we instead propose to combine the prediction sets of each expert by weighted p-value. Right, top row: Comparison of split conformal prediction sets with those learned from weighted aggregation. Weighted aggregation allows overall coverage to follow the coverage of the dominant expert, rather than remain purely marginal. Right, bottom row: Another comparison of split conformal with weighted aggregation, with the latter showing local coverage with smooth transitions.
  • Figure 3: Local validity experiments comparing split conformal to weighted aggregation using absolute residual scores and CQR scores. Each row corresponds to a dataset, with plots for marginal coverage, WS coverage, and prediction set size from left to right. The weighted aggregation methods consistently improves WS coverage across datasets, in most cases meeting target coverage on the WS slab. Split conformal is less conservative marginally, but undercovers on the WS slab; using CQR scores is not enough to cover the gap.
  • Figure 4: Mean coverage compared to the size of the merging set $S_{\text{merge}}$, with different feature assignments (left) and different weighted aggregation variants (right). Overall, we find that coverage improves as the merging set gets larger for all methods. We note that the merging set does not have to be prohibitively large to produce decent results: for example, most feature assignment methods overcover by only 2% with fewer than 200 samples.
  • Figure 5: Coverage (left) and prediction set size (right) for different feature assignment methods and weighted aggregation methods. Our results indicate that sharing fewer features leads to tighter coverage, but sharing more features leads to more efficient (smaller) prediction sets. We also see that WA precise tends to have coverage that is closest to nominal and the most efficient prediction sets, albeit with a much looser guarantee.
  • ...and 9 more figures

Theorems & Definitions (3)

  • Proposition 4.0
  • Proposition 5.0: Infinite-sample guarantee
  • Proposition 5.0: Finite-sample guarantee