Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

2026-04-21

AI Summary

Scientific foundation models are often limited in high-stakes scenarios due to poorly calibrated predictive uncertainty. This work proposes a stochastic attention mechanism applied at inference time, which replaces Softmax weights with normalized multinomial sampling to generate prediction ensembles without requiring model retraining. A single-parameter posterior calibration objective is introduced to enable efficient tuning. The method uniquely incorporates controllable randomness in a lightweight manner during inference, significantly improving both calibration efficiency and the sharpness of predictive intervals. Experimental results on weather and time series forecasting tasks demonstrate that the approach achieves state-of-the-art native calibration performance and yields substantially tighter prediction intervals compared to existing uncertainty-aware baselines, with tuning completed in minutes rather than days.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 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 two scientific foundation models for weather and timeseries forecasting along with an additional regression task. Across benchmarks against uncertainty-aware baselines, we find that Stochastic Attention achieves the strongest native calibration and the sharpest prediction intervals at comparable coverage, while requiring only minutes of post-hoc tuning versus days of retraining for competitive baselines.

Problem

• scientific foundation models
• predictive uncertainty
• calibration
• deterministic outputs
• high-stakes settings

Innovation

• Stochastic Attention
• Inference-Time Calibration
• Predictive Uncertainty
• Foundation Models
• Post-hoc Tuning