# Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

arXiv: https://arxiv.org/abs/2604.19530 · HTML: https://mlsanity.com/paper/2604.19530/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 23 of its 27 references that are in the map.
- First posted on arXiv 21 Apr 2026.

## Builds on (most cited first, 23 of 23)

- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Informer: Beyond Efficient Transformer for Long Sequence Time-Series Forecasting](https://mlsanity.com/paper/2012.07436/) (2020)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [Averaging Weights Leads to Wider Optima and Better Generalization](https://mlsanity.com/paper/1803.05407/) (2018)
- [Revisiting Deep Learning Models for Tabular Data](https://mlsanity.com/paper/2106.11959/) (2021)
- [A decoder-only foundation model for time-series forecasting](https://mlsanity.com/paper/2310.10688/) (2023)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [ClimaX: A foundation model for weather and climate](https://mlsanity.com/paper/2301.10343/) (2023)
- [Benchmarking LLMs via Uncertainty Quantification](https://mlsanity.com/paper/2401.12794/) (2024)
- [Better Uncertainty Calibration via Proper Scores for Classification and Beyond](https://mlsanity.com/paper/2203.07835/) (2022)
- [Variational Learning is Effective for Large Deep Networks](https://mlsanity.com/paper/2402.17641/) (2024)
- [Position: Bayesian Deep Learning is Needed in the Age of Large-Scale AI](https://mlsanity.com/paper/2402.00809/) (2024)
- [Calibrated and Sharp Uncertainties in Deep Learning via Density Estimation](https://mlsanity.com/paper/2112.07184/) (2021)
- [Contextual Dropout: An Efficient Sample-Dependent Dropout Module](https://mlsanity.com/paper/2103.04181/) (2021)
- [Uncertainty-Aware Attention for Reliable Interpretation and Prediction](https://mlsanity.com/paper/1805.09653/) (2018)
- [Transformer Uncertainty Estimation with Hierarchical Stochastic Attention](https://mlsanity.com/paper/2112.13776/) (2021)
- [Gaussian Stochastic Weight Averaging for Bayesian Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2405.03425/) (2024)
- [Variational Low-Rank Adaptation Using IVON](https://mlsanity.com/paper/2411.04421/) (2024)
- [Bayesian Optimization under Uncertainty for Training a Scale Parameter in Stochastic Models](https://mlsanity.com/paper/2510.06439/) (2025)

## Built on by (most cited first, 1 of 1)

- [Decision-Oriented Uncertainty Quantification for Risk Control in Earth System Spatiotemporal Foundation Models](https://mlsanity.com/paper/2609.14821/) (2026)

## Abstract (arXiv)

Transformer-based scientific foundation models are increasingly deployed in high-stakes settings, but current architectures give deterministic outputs and provide limited support for calibrated predictive uncertainty. We propose Stochastic Attention, a sample average lightweight inference-time modification that randomizes attention by replacing softmax weights with normalized multinomial samples controlled by a single concentration parameter, and produces predictive ensembles without retraining. To set this parameter, we introduce a calibration objective that matches the stochastic attention output with the target, yielding an efficient univariate post-hoc tuning problem. We evaluate this mechanism on scientific foundation models for weather and time-series forecasting, as well as several regression tasks. Across benchmarks against uncertainty-aware baselines, we find that Sample Average Stochastic Attention achieves the strongest native calibration and the sharpest prediction intervals at comparable calibration, with adaptation costs nearly three orders of magnitude lower than the next-best baseline.