Robust Bayes-Assisted Conformal Prediction

Abstract: Bayes-assisted conformal prediction combines the strengths of Bayesian modelling with exact, distribution-free frequentist coverage guarantees. Although conformal validity is preserved even when the Bayesian working model (BWM) is misspecified, the size of the resulting prediction sets can degrade substantially when the prior is poorly aligned with the observed data. We address this limitation by introducing RoBAS (Robust Bayes-Assisted Shrinkage): a Bayes-assisted framework for constructing robust nonconformity scores, with two instantiations: one induced by a heavy-tailed BWM, and a closed-form empirical Bayes shrinkage score. The resulting scores adapt to the quality of the working information encoded in the prior: when this information is reliable, they exploit it to produce efficient prediction sets; when it is weak or inaccurate, they revert to the Distance-To-Average (DTA) score, a robust non-informative baseline. We evaluate the proposed scores on tabular and image regression tasks where the training distribution may differ from the calibration and test distributions, while the calibration and test data themselves remain exchangeable. We find that they are competitive with widely used scores in the absence of such shift, while substantially reducing interval widths in shifted settings.

• The paper introduces RoBAS, a framework that dynamically balances Bayes-optimal and data-driven regimes to maintain efficient prediction intervals under prior-data conflict.
• It leverages heavy-tailed priors and residual-level modeling to enable closed-form predictive calculations and adaptive shrinkage.
• Empirical evaluations show that RoBAS stabilizes interval widths across small calibration sizes and distribution shifts, outperforming fixed prior methods.

Robust Bayes-Assisted Conformal Prediction: A Technical Overview

Introduction and Motivation

Bayes-Assisted Conformal Prediction (BACP) integrates Bayesian working models (BWMs) with the frequentist coverage guarantee of conformal prediction (CP), using the posterior predictive as a nonconformity score to generate adaptive prediction sets. While BACP achieves optimal width under model-alignment [Hoff2023], its interval widths degrade when prior information is misaligned with calibration/test data, as the nonconformity score amplifies prior-data conflict into excessive conservativeness. This limitation is especially acute in applications involving distribution shift—ubiquitous in contemporary machine learning pipelines—where the informativeness of the BWM can vary dramatically across subpopulations or over time.

The paper introduces RoBAS (Robust Bayes-Assisted Shrinkage), a framework for residual-based Bayes-assisted scores that dynamically interpolates between a Bayes-optimal regime (when prior information is accurate) and a robust, data-driven regime (when the BWM is misaligned), with computationally tractable instantiations. RoBAS provably contracts to Distance-To-Average (DTA) scoring under extreme prior-data conflict, ensuring that predictive efficiency does not collapse. This mechanism subsumes standard residual-based conformal prediction scores as limiting cases, while maintaining closed-form computations typical of conjugate Bayesian models.

Technical Framework

Conformal Prediction and Score Efficiency

Conformal prediction [Vovk2005] delivers finite-sample, marginal frequentist coverage for $Y_{n+1}$ under exchangeability by defining a prediction set based on a nonconformity score. The choice of score function is critical for efficiency: simple choices such as DTO (absolute residuals) yield tight intervals for unbiased predictors but suffer under systematic bias, while DTA (centered residuals) is robust to mean-shift at the cost of increased variance for small $n$.

Recent work [Hoff2023] shows that using the negative posterior predictive (Bayes-optimality in the conformal framework) yields minimax expected width under the BWM prior, but only when the prior is representative. Standard residual-based CP methods do not correct for prior misspecification, and efficiency can degrade uncontrollably.

Residual-Level Bayes-Assisted Scores

RoBAS reverses the conventional BACP paradigm by specifying the BWM at the residual level of a fixed predictor $f$, rather than at the full generative model $Y|X, \theta$. This decoupling yields several advantages:

• Priors are placed on low-dimensional statistics (mean/variance of residuals) rather than intractable parameter spaces (e.g., neural network weights).
• Computational complexity is reduced: closed-form or efficient numerical scoring is achievable.
• The method generalizes across arbitrary predictors $f$ (including non-Bayesian and ensemble methods).

The nonconformity score is constructed via the conditional posterior predictive for $r_{n+1}|r_{1:n}$, where $r_i = y_i - f(x_i)$. Using heavy-tailed priors on the mean of the residuals, RoBAS adaptively shrinks toward DTA under mean shift, while retaining Bayes-optimal contraction near zero residual mean.

Shrinkage Mechanism and Asymptotic Robustness

The key innovation is the theoretical analysis of the heavy-tailed (regularly varying) prior on the residual mean. The hierarchical model is:

• $R|\theta \sim \mathcal{N}(\theta, \sigma^2)$
• $\theta | \tau^2 \sim \mathcal{N}(0, \gamma \tau^2)$
• $\tau^2 \sim g_{\tau^2}(\cdot)$, with $n$0 regularly varying

Under strong prior-data conflict (large residual mean magnitude), the marginal posterior predictive for a new residual is shown to be monotonically equivalent to $n$1 (DTA score). Thus, interval widths stabilize at the DTA baseline rather than inflating without bound, regardless of the prior's original informativeness. This is a robustification property specific to the shrinkage structure and tail index of $n$2.

The equivalence between DTA and non-informative Bayes is formalized: DTA is the Bayes-assisted score for the improper uniform prior on $n$3.

Two concrete instantiations are presented:

• RoBAS-Full: Hierarchical model with heavy-tailed prior (e.g., horseshoe-induced, using a scale-mixture of normals structure), implemented with closed-form marginal likelihoods involving special functions.
• RoBAS-EB: Empirical Bayes (EB) surrogate using parametric point estimates of prior variance. The resulting plug-in score is $n$4 with shrinkage $n$5 computed adaptively; it reverts to DTA as calibration residual mean increases.

Both methods maintain the computational tractability necessary for application to large-scale or high-throughput prediction settings.

Grid-Free Interval Computation

Standard CP prediction sets may not be intervals even for unimodal nonconformity scores. The paper introduces a grid-free numerical root-finding approach to determine the acceptance region for a candidate $n$6, avoiding typical hyperparameter sensitivity present in fixed grids. For RoBAS-EB, interval monotonicity is proven—the prediction region is always a contiguous interval (up to numerical tolerance), simplifying downstream confidence set usage.

The computational complexity is $n$7 in calibration set size and target precision, significantly more efficient than grid-based alternatives, as the test score only needs to be compared to $n$8 shifted calibration scores per candidate.

Empirical Evaluation

Experiments cover synthetic and real-world settings, targeting small calibration size ($n$9 in $f$0) and substantial training/calibration distribution shift (covariate or domain shifts in tabular and image regression). The following numerical effects are robustly highlighted:

• Prior-Data Alignment: Under no (or weak) shift ($f$1), RoBAS intervals are as tight as optimal Bayes-assisted and DTO scores; DTA lags due to mean variance inflation.
• Prior-Data Conflict: Under strong mean shift, RoBAS interval widths are stable, closely tracking DTA—variance is inherited from observed data and not from prior variance overestimation.
• Comparisons: RoBAS outperforms all Bayes-assisted alternatives under shift, as other methods (with fixed or plug-in prior) suffer severe width inflation. RoBAS remains computationally competitive despite its adaptivity.

This adaptive behavior is most pronounced for small $f$2 and high noise; as $f$3 increases, DTA and RoBAS merge due to reduced mean estimation error, and DTO performance degrades sharply under persistent mean shift.

Across UCI and specialized image regression datasets (UTKFaces, VentricularVolume), RoBAS yields best-in-class or competitive interval widths in IID settings and robust performance under domain adaptation scenarios, with no loss of validity.

Implications and Future Developments

RoBAS addresses an important gap in conformal prediction: score calibration under prior misspecification. Unlike approaches targeting robustness to non-exchangeability (covariate shift between calibration and test) [tibshirani2019conformal, gibbs2021adaptive], RoBAS targets efficiency robustness to the quality of encoded information in the BWM, which is crucial given the unreliability of prior knowledge in modern predictive workflows. Its key implication is that deployment of Bayesian or penalized conformal techniques need not risk loss of actionable informativeness under model misspecification—the method provably reverts to safe data-driven baselines.

From a practical perspective, RoBAS eliminates the need for case-by-case prior elicitation and enables the use of powerful, possibly misaligned predictors (e.g., black-box neural nets trained on disjoint datasets). Closed-form and EB instantiations enable use with batch pipelines and real-time systems.

Theoretical directions include extending robustness guarantees to scale- and higher-moment misspecification (current results target mean misspecification; robustification under heavy-tailed or heteroscedastic residuals is open), and further developing grid-free inference with monotonicity certificates beyond specific score structures.

In the AI pipeline, RoBAS provides a principled means to quantify and penalize model bias in the context of downstream uncertainty quantification, opening directions for robustified predictions in high-stakes domains (healthcare, operations, scientific discovery).

Conclusion

RoBAS provides a theoretically-grounded, computationally-efficient extension of Bayes-assisted conformal prediction that is robust to prior-data conflict, interpolating between Bayes-optimal and DTA efficiency regimes as residual mean shift increases. Its design yields tractable, interval-valued prediction sets, and empirical results demonstrate consistent gains or non-inferiority versus extant methods for both exchangeable and distribution-shifted calibration/test regimes. RoBAS is a significant technical advance in reliable, adaptive uncertainty quantification for deployed predictive systems.

---

[Hoff2023] Bayes-optimal prediction with frequentist coverage control. Bernoulli 29(2):901-928, 2023. (2607.04236) Robust Bayes-Assisted Conformal Prediction. [Vovk2005] Algorithmic Learning in a Random World. Springer, 2005. [romano2019conformalized] Conformalized quantile regression. NeurIPS 32, 2019. [gibbs2021adaptive] Adaptive conformal inference under distribution shift. NeurIPS 34, 2021. [tibshirani2019conformal] Conformal prediction under covariate shift. NeurIPS 32, 2019.