Towards fully differentiable neural ocean model with Veros

TL;DR

This paper introduces a fully differentiable extension of the VEROS ocean model, enabling gradient-based optimization and parameter calibration, thus advancing end-to-end learning in ocean modeling using JAX autodifferentiation.

Contribution

The authors develop and evaluate a differentiable version of the VEROS ocean model compatible with JAX, facilitating automatic differentiation for ocean state correction and parameter calibration.

Findings

Successful implementation of a differentiable ocean model

Demonstrated gradient-based correction of initial ocean states

Enabled direct calibration of physical parameters from observations

Abstract

We present a differentiable extension of the VEROS ocean model, enabling automatic differentiation through its dynamical core. We describe the key modifications required to make the model fully compatible with JAX autodifferentiation framework and evaluate the numerical consistency of the resulting implementation. Two illustrative applications are then demonstrated: (i) the correction of an initial ocean state through gradient-based optimization, and (ii) the calibration of unknown physical parameters directly from model observations. These examples highlight how differentiable programming can facilitate end-to-end learning and parameter tuning in ocean modeling. Our implementation is available online.