Skillful Global Ocean Emulation and the Role of Correlation-Aware Loss

TL;DR

This paper presents a global ocean emulator based on GraphCast architecture, trained on NOAA data, that achieves skillful medium-range forecasts up to 15 days, enhanced by a correlation-aware loss function.

Contribution

It introduces a dedicated ocean emulator with a novel Mahalanobis distance loss that improves forecast skill and acts as a regularizer for ocean dynamics.

Findings

The emulator forecasts 10-15 days ahead with skillful accuracy.

Mahalanobis distance loss outperforms MSE loss in forecast skill.

Correlation-aware loss regularizes slow ocean dynamics.

Abstract

Machine learning emulators have shown extraordinary skill in forecasting atmospheric states, and their application to global ocean dynamics offers similar promise. Here, we adapt the GraphCast architecture into a dedicated ocean-only emulator, driven by prescribed atmospheric conditions, for medium-range predictions. The emulator is trained on NOAA's UFS-Replay dataset. Using a 24 hour time step, single initial condition, and without using autoregressive training, we produce an emulator that provides skillful forecasts for 10-15 day lead times. We further demonstrate the use of Mahalanobis distance as loss that improves the forecast skill compared to the Mean Squared Error loss by explicitly accounting for the correlations between tendencies of the target variables. Using spatial correlation analysis of the forecasted fields, we also show that the proposed correlation-aware loss acts as a statistical-dynamical regularizer for the slow, correlated dynamics of the global oceans, offering a better background forecast for downstream tasks like data assimilation.