Adaptive Individual Uncertainty under Out-Of-Distribution Shift with Expert-Routed Conformal Prediction

TL;DR

Problem: current UQ methods struggle to provide reliable, per-sample uncertainty under distribution shift in high-stakes drug discovery. Approach: TESSERA combines MoE-based epistemic and aleatoric signals and applies a single split-conformal calibration to yield individualized intervals with finite-sample, distribution-free marginal coverage at $1-\alpha$. Findings: on ChEMBL31 PLI data with scaffold-based OOD and i.i.d. splits, TESSERA achieves near-nominal coverage and strong width adaptivity (CWC) across multiple encoders, with low AUSE and informative SSC behavior. Significance: enables actionable, selective prediction in drug-discovery pipelines and demonstrates that combining MoE diversity with conformal calibration yields trustworthy uncertainty under distribution shift.

Abstract

Reliable, informative, and individual uncertainty quantification (UQ) remains missing in current ML community. This hinders the effective application of AI/ML to risk-sensitive domains. Most methods either fail to provide coverage on new data, inflate intervals so broadly that they are not actionable, or assign uncertainties that do not track actual error, especially under a distribution shift. In high-stakes drug discovery, protein-ligand affinity (PLI) prediction is especially challenging as assay noise is heterogeneous, chemical space is imbalanced and large, and practical evaluations routinely involve distribution shift. In this work, we introduce a novel uncertainty quantification method, Trustworthy Expert Split-conformal with Scaled Estimation for Efficient Reliable Adaptive intervals (TESSERA), that provides per-sample uncertainty with reliable coverage guarantee, informative and adaptive prediction interval widths that track the absolute error. We evaluate on protein-ligand binding affinity prediction under both independent and identically distributed (i.i.d.) and scaffold-based out-of-distribution (OOD) splits, comparing against strong UQ baselines. TESSERA attains near-nominal coverage and the best coverage-width trade-off as measured by the Coverage-Width Criterion (CWC), while maintaining competitive adaptivity (lowest Area Under the Sparsification Error (AUSE)). Size-Stratified Coverage (SSC) further confirms that intervals are right-sized, indicating width increases when data are scarce or noisy, and remain tight when predictions are reliable. By unifying Mixture of Expert (MoE) diversity with conformal calibration, TESSERA delivers trustworthy, tight, and adaptive uncertainties that are well-suited to selective prediction and downstream decision-making in the drug-discovery pipeline and other applications.

Figures (5)

• Figure 1: Scaffold-wise coverage under distribution shift. (A) Prediction-interval coverage probability (PICP) computed per Bemis–Murcko scaffold for the 15 most frequently occuring scaffolds in the test set. (B) The same analysis for the 15 least frequent scaffolds. The dashed line marks the nominal target $1-\alpha$=0.90. Bars are color-coded by methods and the scaffold sketches below the axis correspond to the x-axis groups. The panels highlight how coverage varies across scaffolds and under rarity and shift. TESSERA generally remain closer to the target line, whereas non-calibrated baselines tend to under-cover on rare scaffolds.
• Figure 2: Visualizing predictive intervals and adaptivity across methods. (A–F) Scatter of predictions $\hat{y}$ versus ground truth $y$ on the test set. The semi-transparent band show the method’s prediction intervals (smoothed along sorted $y$ for visual clarity), and the dashed line is $y{=}x$. Panels correspond to: (A) CP–Predicted Variance, (B) CP–Expert Disagreement, (C) Monte Carlo Dropout, (D) RIO-GP, (E) eMOSAIC, and (F) Classical CP.
• Figure 3: Size–Stratified Coverage (SSC). For each method, test points are sorted by prediction–interval (PI) width and split into equal–sized bins (left $\rightarrow$ right: narrow $\rightarrow$ wide). Bars show the empirical coverage within each bin. The three panels report SSC with $J\!\in\!\{3,5,10\}$ bins (top to bottom). The dashed horizontal line indicates the nominal target coverage $1-\alpha$ (here, $0.90$). Ideally, all bars should lie near the dashed line, indicating conditionally well–calibrated PIs across easy (narrow) and hard (wide) regions. We compare our method with various baselines. We omit SSC for Classical CP because its split-residual construction produces constant-width intervals making SSC uninformative, SSC is designed to assess adaptive methods whose widths vary with difficulty.
• Figure 4: Disentangling Uncertainty
• Figure 5: Risk–coverage (sparsification) curves: RMSE on the remaining data after removing the fraction on the x-axis with the highest predicted uncertainty. The dashed black curve is the oracle (removing points by true error). The legend reports AUSE (area between each curve and the oracle, lower is better).