Research map / Time series / Causal and generative time series

Calibrating Scientific Foundation Models with Inference-Time Stochastic Attention

NeurIPS 2026 · Time series › Causal and generative time series · approach: Probabilistic and conformal forecasting · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Stochastic Attention, an inference-time method that replaces softmax attention weights with normalized multinomial samples controlled by one concentration parameter nu, selected by a calibration objective matching attention-induced deviation to held-out residuals.

Problem

Scientific foundation models give deterministic outputs with limited calibrated predictive uncertainty, and existing UQ methods (ensembles, SWAG, IVON, MC Dropout) are costly or poorly calibrated at scale.

Approach

Interpret softmax attention weights as a categorical distribution and replace the expectation with a normalized multinomial sample average of size nu (mean-preserving, variance scales as 1/nu); choose nu by a one-dimensional post-hoc objective matching stochastic deviation magnitude to held-out residual magnitude, using Bayesian optimization.

Key results

• On TimesFM ETTh1 (H=96), SA reaches 61.5% empirical coverage versus 15.6%-38.5% for baselines against a 95% target.
• On ClimaX 72-hour Z500, SA achieves native PIT W1=0.056 while all baselines are U-shaped (W1 0.10-0.22); SA gives the narrowest PI-95 intervals with MultiSWAG closest at ~1.02x SA width.
• On UCI regression, SA has the best native PIT W1 on 5 of 8 datasets and the narrowest calibrated intervals on 7 of 8 after baseline temperature scaling.

Surprising findings

• On ClimaX, HSA's training-stage stochasticity degrades anomaly correlation from 0.964 to 0.7476.
• Score-optimal nu=25 gives better CRPS and Energy Score than calibration-first nu=4, at the cost of PIT uniformity.

Limitations

• nu is a single global concentration shared across all heads and layers, limiting expressivity.
• Calibration objective is stated for continuous regression-style residuals; classification and structured targets are not directly covered.
• On Wine and Protein (highest RMSE), SA is competitive but not the leader in native calibration.

Experimental setup

Datasets: ETTh1 ETTh2 ETTm1 ETTm2 ERA5 UCI Concrete UCI Energy UCI Naval UCI Yacht UCI Protein UCI Wine UCI Kin8nm

Benchmarks: ETT-small ClimaX 72-hour Z500 forecasting (ERA5) UCI regression (8 datasets)

Metrics: MAE latitude-weighted ACC normalized RMSE PIT Wasserstein-1 from uniform empirical coverage PI-95 width CRPS Energy Score

Baselines: SWAG MultiSWAG IVON MC Dropout Contextual Dropout HSA

Builds on: LoRA temperature scaling SWAG MC Dropout Bayesian optimization under uncertainty

Closest papers by meaning

• Introducing Spectral Attention for Long-Range Dependency in Time Series Forecasting (NeurIPS 2024 · similarity 0.86)
• In-Context Fine-Tuning for Time-Series Foundation Models (ICML 2025 · similarity 0.85)
• Sundial: A Family of Highly Capable Time Series Foundation Models (ICML 2025 · similarity 0.84)
• Common Task Framework For a Critical Evaluation of Scientific Machine Learning Algorithms (NeurIPS 2025 · similarity 0.83)
• Revisiting Attention for Multivariate Time Series Forecasting (AAAI 2025 · similarity 0.83)
• SEMPO: Lightweight Foundation Models for Time Series Forecasting (NeurIPS 2025 · similarity 0.83)