Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

TL;DR

This paper introduces Stochastic Attention, a lightweight inference-time method for calibrating uncertainty in transformer-based scientific models, improving calibration and prediction interval sharpness without retraining.

Contribution

It proposes a novel stochastic attention mechanism with a calibration objective, enabling effective uncertainty calibration at inference time for scientific models.

Findings

Achieves the strongest native calibration among tested methods.

Produces sharper prediction intervals at comparable calibration levels.

Requires nearly three orders of magnitude less adaptation cost than baselines.

Abstract

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.