Improving Coverage in Combined Prediction Sets with Weighted p-values

TL;DR

This paper introduces a flexible weighted aggregation framework for conformal prediction sets that maintains finite-sample validity and achieves adaptive coverage, especially useful in mixture-of-experts models.

Contribution

It proposes a novel weighted aggregation method for conformal prediction sets that controls coverage bounds and extends to data-dependent weights, broadening applicability.

Findings

Achieves tighter coverage bounds through weighted aggregation.

Maintains finite-sample validity with data-dependent weights.

Demonstrates adaptive coverage in mixture-of-experts models.

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-\alpha$ coverage inevitably weakens the overall guarantee, typically resulting in $1-2\alpha$ 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\alpha$ guarantee of the combined models and the $1-\alpha$ 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.