Multidimensional Distributional Neural Network Output Demonstrated in Super-Resolution of Surface Wind Speed

TL;DR

This paper introduces a neural network framework that models multidimensional Gaussian distributions for uncertainty quantification, specifically applied to super-resolution of surface wind speed, capturing spatial correlations efficiently.

Contribution

It presents a novel training approach using a Gaussian loss with Fourier-based covariance stabilization and a regularization strategy for better covariance estimation.

Findings

Effective super-resolution of surface wind speed with uncertainty quantification

Stable training via Fourier representation of covariance matrices

Enhanced covariance estimation through information sharing regularization

Abstract

Accurate quantification of uncertainty in neural network predictions remains a central challenge for scientific applications involving high-dimensional, correlated data. While existing methods capture either aleatoric or epistemic uncertainty, few offer closed-form, multidimensional distributions that preserve spatial correlation while remaining computationally tractable. In this work, we present a framework for training neural networks with a multidimensional Gaussian loss, generating closed-form predictive distributions over outputs with non-identically distributed and heteroscedastic structure. Our approach captures aleatoric uncertainty by iteratively estimating the means and covariance matrices, and is demonstrated on a super-resolution example. We leverage a Fourier representation of the covariance matrix to stabilize network training and preserve spatial correlation. We introduce a novel regularization strategy -- referred to as information sharing -- that interpolates between image-specific and global covariance estimates, enabling convergence of the super-resolution downscaling network trained on image-specific distributional loss functions. This framework allows for efficient sampling, explicit correlation modeling, and extensions to more complex distribution families all without disrupting prediction performance. We demonstrate the method on a surface wind speed downscaling task and discuss its broader applicability to uncertainty-aware prediction in scientific models.