Improving Conformal Prediction Sets Through Semantic Neighborhood Diffusion

In safety-critical applications such as robotics, healthcare, and weather prediction, the reliability of uncertainty estimates is of crucial importance, as decisions guided by misplaced confidence risk costly and hazardous outcomes. Current deep learning models routinely produce unreliable and overconfident uncertainty scores, severely limiting their use in such applications. One promising direction is conformal prediction which transforms heuristic uncertainty scores into prediction sets with rigorous statistical coverage guarantees. While conformal prediction guarantees coverage, the efficiency of prediction sets is limited by how well-aligned the model's predictive distribution is to the true data distribution; misaligned or noisy predictive distributions force current methods to produce excessively large, uninformative prediction sets in order to maintain guarantees. In this work, we propose a simple method that smooths over noisy and overconfident nonconformity scores by diffusing over neighboring data points on the data manifold, leveraging the growing capabilities of off-the-shelf foundation models to produce this embedding space. We characterize conditions for when diffusion reduces prediction set size, and empirically demonstrate that the proposed method significantly reduces the average prediction set size, also under distribution shifts.