🤖 machine learning

Scalable Uncertainty Quantification for Extreme Weather Forecasting via Empirical Neural Tangent Kernels

This paper introduces a scalable, inference-time uncertainty quantification framework for deep learning weather models that leverages Neural Tangent Kernels and adaptive feature decomposition (ICA or SVD) to generate sharp, severity-scaling prediction intervals for extreme events, outperforming conformal baselines while requiring no model retraining.

Original authors: Jose Marie Antonio Miñoza, Rex Gregor Laylo, Sebastian C. Ibañez

Published 2026-06-03

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jose Marie Antonio Miñoza, Rex Gregor Laylo, Sebastian C. Ibañez

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Weather Forecasts Without a "Confidence Meter"

Imagine you are checking the weather app before a big trip. The app says, "It will rain tomorrow." That's a point forecast. It's a single, definite answer.

But what if the app could also say, "It will rain tomorrow, and I'm 90% sure of that, but there's a small chance it might be a hurricane"? That second part is Uncertainty Quantification (UQ).

Currently, the newest, fastest AI weather models (like FourCastNetV2 or Pangu-Weather) are incredibly accurate and fast. However, they only give you the "point forecast." They act like a confident driver who never admits they might be wrong. For extreme events like typhoons or floods, knowing how much to trust the forecast is just as important as the forecast itself.

The Solution: The "NTK-UQ" Pipeline

The authors propose a new method called NTK-UQ (Neural Tangent Kernel Uncertainty Quantification). Think of it as a "confidence meter" you can attach to any existing weather AI without having to rebuild the car.

Here is how it works, step-by-step:

1. The "Last-Layer" Snapshot

Imagine the weather AI is a giant library of atmospheric knowledge. When it looks at a new weather pattern, it processes the data through many layers of "thinking."

The Trick: The authors don't look at the whole library. They only look at the very last shelf (the last layer) where the AI stores its final "thoughts" about the weather.
The Analogy: It's like asking a student, "What is the main idea of this story?" You don't need to read the whole book again; you just look at their summary notes. These notes (called features) contain the AI's understanding of the atmosphere.

2. The "Similarity Check" (The Kernel)

Once the AI extracts these summary notes, the system asks: "Does this new weather pattern look like anything we've seen before in our training data?"

The Analogy: Imagine you are a chef. You taste a new soup.
- If the soup tastes exactly like a recipe you've made 1,000 times, you are very confident (low uncertainty).
- If the soup tastes weird and nothing like your past recipes, you are very unsure (high uncertainty).
The math behind this is called a Gaussian Process. It calculates how "strange" the current weather is compared to the AI's memory bank.

3. The "Sorting Hat" (SVD vs. ICA)

This is the paper's biggest discovery. To make the math fast and accurate, the system has to sort the "summary notes" into groups. The authors found that one size does not fit all.

Method A (SVD): Think of this as sorting books by thickness. It groups things that are very common and have a lot of "volume" (variance). This works well for routine weather but misses the weird, extreme stuff.
Method B (ICA): Think of this as sorting books by weirdness or unique flavor. It looks for rare, distinct patterns (like a specific type of storm) that don't show up in the "average" data.
The Finding:
- For some AI models (like those based on "Transformers"), the ICA (Weirdness) method is better. It can spot a coming hurricane and say, "Hey, this is dangerous!"
- For other models (like those based on "Fourier" math), the SVD (Thickness) method is better. If you use the wrong sorter, the system might crash or give you a false sense of security.

The paper provides a simple "rule of thumb" (an algorithm) to tell you which sorter to use based on the math of the specific AI model you are using.

4. The Result: Adaptive Intervals

When you combine these steps, you get a forecast that doesn't just say "Rain." It says:

Routine Day: "Rain, 100% chance." (Narrow interval, low uncertainty).
Extreme Day: "Rain, but the amount could be anywhere from 1 inch to 10 inches." (Wide interval, high uncertainty).

Unlike older methods that give you the same "wide or narrow" answer for every day, this system scales. It gets wider and more cautious exactly when the weather gets dangerous.

Why This Matters (According to the Paper)

No Re-training Needed: You don't have to teach the AI anything new. You just take the AI as it is, run it through this "confidence meter," and get the results. It's like putting a new dashboard on an old car.
Super Fast: Once the setup is done, checking the confidence takes almost no extra time. It's just a quick math calculation (a matrix-vector product).
Better than the Competition: The authors tested this against a standard method called "Conformal Prediction." Their method produced 31–37% tighter (sharper) intervals while still being accurate.
- Analogy: If the old method said, "The storm will be between 1 and 100 miles wide," the new method says, "It will be between 40 and 60 miles wide." Both are correct, but the new one is much more useful.
Architecture Matters: The paper proves that you can't just pick one math trick for all AI models. You have to match the math trick (SVD or ICA) to the specific "brain" of the AI model, or the uncertainty estimates will fail (a problem they call "Variance Collapse").

Summary

The paper introduces a way to give "confidence scores" to super-fast AI weather models. It works by looking at the AI's final thoughts, checking how strange the current weather is compared to history, and using a smart sorting system (choosing between SVD or ICA) to ensure the confidence score is accurate. This allows forecasters to know exactly when to trust the AI and when to be extra careful during extreme weather events.

Technical Summary: Scalable Uncertainty Quantification for Extreme Weather Forecasting via Empirical Neural Tangent Kernels

1. Problem Statement

Deep learning weather models (e.g., FourCastNetV2, Pangu-Weather, Aurora, AIFS) have achieved accuracy comparable to traditional numerical weather prediction (NWP) systems while operating orders of magnitude faster. However, these models predominantly produce deterministic point forecasts lacking calibrated uncertainty estimates. This gap is critical for high-stakes decision-making during extreme weather events (e.g., tropical cyclones, floods), where understanding the reliability of a forecast is as important as the forecast itself.

Existing uncertainty quantification (UQ) methods face significant barriers when applied to billion-parameter foundation models:

Deep Ensembles: Computationally prohibitive due to the need to train multiple copies of massive models.
Bayesian Methods: Require architectural modifications and incur substantial memory/compute overhead.
Conformal Prediction: Provides distribution-free coverage guarantees but produces uniform interval widths that do not correlate with actual forecast errors or event severity, failing to distinguish extreme events from routine conditions.

The paper addresses the need for a scalable, post-hoc UQ method that requires no retraining, works with pre-trained checkpoints, and produces adaptive uncertainty intervals that scale with event severity.

2. Methodology: NTK-UQ Framework

The proposed framework, Neural Tangent Kernel-based Uncertainty Quantification (NTK-UQ), treats the last-layer features of a frozen pre-trained weather model as an empirical Neural Tangent Kernel (NTK). Under the last-layer NTK–Gaussian Process (GP) correspondence, the model is viewed as a GP where the kernel is defined by the learned feature representations.

The methodology consists of three phases:

2.1 Last-Layer Feature Extraction

For a pre-trained model $f_\theta = g_\psi \circ \phi_\omega$ , the method extracts the last-layer features $\phi(x)$ learned from decades of ERA5 reanalysis data. These features encode physically meaningful atmospheric structure. For spatial feature maps, multi-statistic aggregation (mean, std, min, max, percentiles) is used to create fixed-dimensional feature vectors regardless of spatial resolution.

2.2 Kernel Decomposition and GP Posterior Construction

The empirical kernel $K(x, x') = \phi(x)^\top \phi(x')$ is constructed. To make the GP posterior variance computationally tractable for large calibration sets, the centered feature matrix is decomposed. The paper compares two decomposition strategies:

Singular Value Decomposition (SVD): Captures orthogonal directions of maximum variance (second-order statistics).
Independent Component Analysis (ICA): Exploits higher-order statistics (kurtosis, negentropy) to isolate statistically independent sources.

The predictive variance $\sigma^2(x^*)$ is calculated using the GP posterior formula:
$\sigma^2(x^*) = K(x^*, x^*) + \sigma_n^2 - k_*^\top (K + \sigma_n^2 I)^{-1} k_*$
where the correction term depends on the similarity of the test input to the calibration distribution in feature space.

2.3 Decomposition Selection Rule

The paper proposes a data-driven rule (Algorithm 1) to select between SVD and ICA based on the eigenspectrum concentration ratio ( $\rho = \lambda_1 / \sum \lambda_j$ ):

Concentrated Spectrum ( $\rho > 0.8$ ): Typical of spectral operators (e.g., SFNO). SVD is preferred with aggressive rank truncation ( $k \le 10$ ) to avoid variance collapse.
Distributed Spectrum ( $\rho < 0.5$ ): Typical of attention-based or graph-based models. ICA is preferred to exploit non-Gaussian, heavy-tailed structures inherent in extreme weather.
Intermediate: Both methods are evaluated on a validation split.

2.4 Post-Hoc Calibration

Raw NTK uncertainties are scaled via a per-variable, per-lead-time factor $\alpha$ (learned via binary search) to achieve a target coverage level (e.g., 90%). This ensures the prediction intervals are calibrated to the ground truth.

3. Theoretical Insights

The paper provides theoretical analysis explaining why UQ quality is architecture-dependent:

Variance Collapse (Proposition 1): If the truncation rank $k$ approaches the effective rank of the feature space, the GP correction term consumes nearly all prior variance. This destroys the discrimination between extreme events and routine conditions. The paper establishes a diagnostic criterion $R_k = C_k/P < 0.9$ to prevent this collapse, noting that architectures with concentrated spectra require aggressive truncation.
Non-Gaussian Discrimination (Proposition 2): Extreme weather events exhibit heavy-tailed, non-Gaussian distributions. While SVD captures only second-order variance, ICA exploits higher-order statistics (kurtosis) to isolate rare extreme-event signatures. This allows ICA to produce adaptive uncertainty estimates that scale with event severity, a capability SVD lacks for heavy-tailed features.

4. Experimental Results

The framework was evaluated on four diverse AI weather models (FourCastNetV2, Pangu-Weather, Aurora, AIFS) using 100 extreme weather events from the EM-DAT database (2021), including tropical cyclones, floods, and heatwaves.

Coverage: All models achieved target 90% coverage (89–91%) after post-hoc scaling.
Decomposition Performance:
- ICA achieved proper coverage and lower CRPS (Continuous Ranked Probability Score) for three models (AIFS, Aurora, FourCastNetV2). It produced adaptive intervals with high coefficient of variation (CV = 0.07–1.81), successfully distinguishing tropical cyclones from routine weather.
- SVD failed to achieve coverage for Aurora (58.1%) and produced more uniform intervals (CV = 0.01–0.49) that failed to discriminate event severity.
- Pangu-Weather was an exception where SVD satisfied coverage while ICA failed, though both exhibited numerical instabilities.
Comparison with Conformal Prediction: NTK-UQ (using the optimal decomposition per model) achieved 31–37% sharper prediction intervals at 90% coverage compared to split conformal prediction. Crucially, unlike conformal prediction which produces uniform intervals by construction, NTK-UQ produces adaptive intervals that scale with extreme event severity.
Spatial Uncertainty: For models with spatially-resolved features (e.g., AIFS), ICA enabled spatially-varying uncertainty maps that correlated with forecast error concentrations, allowing for targeted warnings.

5. Significance and Claims

The paper claims the following significance:

Scalability: NTK-UQ requires no model retraining or architectural changes, making it applicable to any pre-trained deterministic checkpoint. Inference adds minimal overhead (a single matrix-vector product).
Adaptive Discrimination: Unlike conformal prediction, NTK-UQ produces uncertainty intervals that adapt to the severity of the weather event, providing actionable discrimination between routine and extreme conditions.
Architecture-Aware Design: The work establishes that UQ quality is not universal but depends on the interaction between neural architecture (eigenspectrum properties) and the decomposition method. It provides a predictive diagnostic (eigenspectrum concentration) to select the appropriate UQ strategy without exhaustive experimentation.
Operational Utility: The framework enables resource-constrained meteorological agencies to add calibrated uncertainty estimates to existing high-performance models, lowering the barrier to trustworthy forecasting for disaster preparedness.

The authors note limitations, including the reliance on approximations (last-layer NTK-GP correspondence holds rigorously only at infinite width) and the evaluation being bounded to a single out-of-distribution year (2021). However, empirical results consistently validate the theoretical predictions across diverse architectures.