TabCUQ improves uncertainty estimates for tabular prediction models
TabCUQ: Context-Uncertainty Quantification for Tabular Foundation Models
Machine Learning
Summary
When using computer models to predict outcomes for groups of people based on data tables, it’s important to understand how uncertain these predictions are. The authors study two sources of uncertainty: not having enough context or examples, and variability due to small data samples. They introduce TabCUQ, a way to add noise and use a special technique to better measure and separate these uncertainties. Their method gives more reliable confidence intervals, making predictions safer to trust when making decisions.
What this means in practice
- •For medical data analysts: Provide more reliable uncertainty intervals when estimating treatment effects for patient subgroups using tabular models.
- •For financial risk teams: Improve risk estimates by better accounting for uncertainty in predictions made from small or limited financial datasets.
Authors
Yuebei Shi, Molei Liu, Zijian Guo
Abstract
Tabular foundation models use in-context learning to estimate conditional expected outcomes and average treatment effects for specific groups of individuals. Using these estimates to guide scientific decisions requires reliable uncertainty quantification. The challenge is to account simultaneously for two sources of uncertainty: limited-context bias and context variation arising from the small labeled sample supplied at inference. Predictive spread does not directly quantify both sources of uncertainty, and we observe that the ordinary bootstrap may substantially overestimate context variation, whereas subsampling may underestimate it. We introduce TabCUQ, a perturbation-based framework that injects noise into observed outcomes and applies the foundation model to the perturbed contexts. TabCUQ estimates context variation using wild-bootstrap dispersion and accounts for limited-context bias using the maximum deviation between perturbed and original estimates. For a kernel-regression analogue whose estimation error admits an exact decomposition into these two components, we show that our framework provides asymptotically valid pointwise coverage. In simulations, TabCUQ substantially improves coverage, with average interval widths typically about twice those of an oracle bias-aware benchmark. Across seven real-world tasks, it also improves coverage of reference functions learned from much larger samples.