Njord: A Probabilistic Graph Neural Network for Ensemble Ocean Forecasting

arXiv:2605.15470 · cs.LG, physics.ao-ph · Submitted 2026-05-14 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: Today's paper: "Njord: A Probabilistic Graph Neural Network for Ensemble Ocean Forecasting".

Jane: The gist Njord introduces a probabilistic data-driven model for ocean forecasting that combines deep latent variable frameworks with graph neural networks to provide calibrated ensemble forecasts and uncertainty estimates for…

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So we're looking at Njord today, which is this new probabilistic data-driven model for ocean forecasting. It’s designed to give us calibrated ensemble forecasts and uncertainty estimates for both global and regional ocean domains, which is pretty significant because most existing models just give you one answer.

Jane: That sounds like a big deal when you're dealing with chaotic systems like the ocean, Tom. What does this model actually combine? What's the core idea behind Njord?

Tom: Well, it combines a deep latent variable framework with a graph neural network architecture. The main mechanical thing is that it lets us sample each forecast step in a single forward pass, which is pretty efficient for these kinds of complex models.

Lu: Exactly, and what makes the GNN part interesting is that they use the hierarchical Graph Neural Networks from Graph-EFM to capture those spatial relationships over the entire forecasted area. It’s about mapping gridded inputs to latent representations on this mesh graph, where stochasticity from a learned latent variable Zt affects all of the outputs >

Meng: I gotta ask, how does it handle the sheer size of these ocean grids? If we're talking global models at zero point two five degree resolution, that sounds computationally heavy <ref:2605.15470#pg1>.

Tom: That’s where they introduced something specific for scaling up: a clustering-based graph layout constructed via spherical K-means clustering of ocean grid point three dee Cartesian coordinates with latitude-based area weights <ref:2605.15470#pg2>. It’s designed to place nodes based on the density of those grid points to better conform to the irregular geometry of the ocean surface >

Jane: So, it's not just using a standard fixed mesh, but dynamically building the graph layout based on where the ocean data actually is, which makes sense for irregular coastlines. What about training? How do they teach this thing to be reliable?

Tom: They use a combination of the Evidence Lower Bound loss and the Continuous Ranked Probability Score loss. The full loss function involves an expectation term, a KL divergence term to regularize the latent variable Zt, and then they add the CRPS term for calibration >

Lu: And they trained it using a curriculum approach. They gradually activate things, starting with pure autoencoder training before introducing the KL divergence term and finally enabling the CRPS loss. For Njord-Baltic specifically, they trained it for three hundred fifty epochs following a staged curriculum that begins with pure autoencoder training >

Meng: Three hundred fifty epochs sounds like a lot of time on expensive hardware. And how do we know it’s actually learning physics, not just memorizing patterns?

Lalam: The paper states they are tackling the forecasting task of mapping a sequence of initial ocean states X−p+one:zero = (X−p+one <ref:2605.15470#pg2>..., Xzero), where p is the number of past steps to a sequence of future states Xone:T = (Xone..., XT) >

Title and authors: Jane: So we’re mapping past ocean states to future ones. And the input forcing F−p+one:T includes known dynamic factors like the day of year, static features like depth, and atmospheric forcing at the surface from a weather model > <ref:2605.15470#pg2>

Tom: Right, and that forces it to account for things like wind, pressure, and temperature driving surface dynamics. That's how you get those initial conditions into the model >

Lu: And on page two of this work they mention that in a regional setting Njord also conditions on boundary data from an independent global ocean model, which is interesting because previous emulators either lack boundary forcing or depend on data from the very system they aim to replace at the boundary > <ref:2605.15470#pg1>

Meng: That means it’s not just predicting the future in isolation; it’s using external information, which is good for realism but adds complexity to deployment. What about sea ice? I saw a mention of how they handle that.

Tom: They model sea ice by incorporating physical constraints instead of post-processing them. They use a density channel mechanism where they construct a binary density channel d from Sea Ice Concentration SIC as d equals I

SIC > zero: > <ref:2605.15470#pg1>

Jane: So, instead of just guessing the ice, the model predicts this density logit, and then it passes that through a sigmoid and thresholded at zero point five to set the density channel and ice variables to zero in the next state if it falls below that threshold >

Lu: That’s an enhancement over just predicting a raw value because it enforces physical reality for those specific variables, ensuring locations predicted to be ice-free get clean zero-ice inputs rather than small residual values that might accumulate over rollouts >

Meng: That sounds like a solid fix for the sea ice representation. Now, what about the results? How did it compare to those deterministic machine learning baselines they used?

Tom: In their experiments, they demonstrate strong performance on both global and regional domains compared to deterministic machine learning baselines > On the global OceanBench benchmark, Njord achieves the lowest errors on average across upper-ocean variables when evaluated against real-world observations >

Jane: That’s a key finding. Lower errors on real data is one thing, but what about uncertainty? The paper focuses heavily on providing those uncertainty estimates from the sampled ensemble forecasts >

Tom: They also show that the ensemble spread exhibits good calibration, as reflected by the Spread-Skill Ratio SSR which increases steadily toward one > On Njord-Baltic at two km resolution, they achieved errors comparable to a deterministic baseline while still providing probabilistic forecasts > <ref:2605.15470#pg2>

Title and authors: Lu: That’s a solid comparison. The paper also shows that increasing the ensemble size to M = twenty yields a systematic RMSE reduction across all evaluated variables and depth levels >

Meng: Increasing the ensemble size is always good for uncertainty estimation, but running twenty members at once sounds like it would push the computational cost up significantly in a production environment. What about scaling this to even higher resolutions?

Tom: They address scaling by modifying the Interaction and Propagation Networks to use separate dimensionalities for grid embeddings, edge representations, and graph node representations > And for global zero point two five degree resolution, they use a two-stage training schedule: pretraining on one degree resolution data followed by finetuning at zero point two five° > <ref:2605.15470#pg2>

Jane: So they’ve built a framework that adapts to both the spatial structure of the ocean and the physical constraints of sea ice, all while providing calibrated uncertainty estimates, which is what we need for operational use. So, where does this leave us?

Tom: Well, it positions Njord as a probabilistic generative machine learning model for global and regional ocean forecasting that successfully scales generative ensemble modeling to high-resolution ocean grids with irregular coastal geometries >

Jane: It really sounds like a step forward in moving ocean AI from just deterministic answers to providing the uncertainty needed for better decision making about things like currents or sea ice > So, we've seen how Njord works and what it achieves on its own. Now let’s see what the authors are suggesting they should do next >

Lu: The paper suggests scaling to large ocean grids by modifying those Interaction and Propagation Networks for separate dimensionalities, which is a good direction for future work > They also mention applying the framework regionally to construct Njord-Baltic using quadrilateral graphs, which is useful for regional models >

Meng: I wonder if that scaling mechanism will be easy to implement on existing cloud infrastructure. If we have to completely restructure how the embeddings are handled across different dimensionalities, that’s a big engineering lift >

Lalam: The authors also showed how increasing the ensemble size to twenty reduces RMSE systematically, so perhaps pushing that ensemble size further is something worth exploring for better uncertainty quantification >

Tom: Yeah, it’s about getting more reliable estimates. So this paper is Njord: A Probabilistic Graph Neural Network for Ensemble Ocean Forecasting, and it shows how probabilistic graph neural networks can handle the complexity of ocean forecasting > We've covered the architecture, the training methods, and the results on performance versus deterministic models.

Jane: It’s a solid piece of work that moves us closer to having AI models that give us not just a prediction, but a probability distribution about what might happen in complex environments like the ocean > We’ll see what else is coming from this research next >

The paper's summary: Tom: So to recap, Njord is this new model for ocean forecasting that uses deep latent variables and graph neural networks to give us ensemble forecasts and uncertainty estimates for both global and regional ocean areas >

Jane: That’s right, Tom, it’s moving away from just giving you one single answer because the ocean is so complex >

Tom: Exactly. The core idea here is that they let the model sample each step of the forecast in just a single forward pass which makes it pretty efficient for these huge calculations >

Jane: And what really stands out, I think is how they handle those massive ocean grids, especially when dealing with irregular coastlines, because most models use fixed shapes >

Tom: They built a unique graph layout by doing a spatial K-means clustering of the ocean grid point three dee Cartesian coordinates using latitude-based area weights >

Jane: That means the model’s structure actually conforms to where the ocean data is located, which is a big improvement over just sticking to a standard mesh >

Tom: Right. And for training, they use a combination of losses—the Evidence Lower Bound loss and the Continuous Ranked Probability Score loss—to make sure it's both accurate and calibrated >

Jane: I’m really curious about how this translates into real-world predictions, because that’s where the uncertainty estimates matter most for folks who actually need to make decisions about currents or sea ice >

Tom: The experiments show strong performance against deterministic machine learning baselines on both global and regional scales > On the global OceanBench benchmark, Njord achieves lower errors on average across upper-ocean variables compared to real-world observations >

Jane: That’s a solid number to have, but what about the reliability of those predictions? Does it actually give you a good sense of how wrong it might be?

Tom: They show that the ensemble spread has good calibration, which they measure using the Spread-Skill Ratio or SSR, and that number increases steadily toward one > This means the uncertainty estimates are actually trustworthy >

Jane: So, even when things get complicated in places like the Baltic Sea at a two kilometer resolution, it still delivers probabilistic forecasts comparable to traditional deterministic models >

Tom: Precisely. And they also found that if you increase the ensemble size to twenty members, you see a systematic reduction in the root mean square error across all variables and depth levels >

Jane: It sounds like this model is getting really good at capturing the messy physics of ocean dynamics while giving us reliable probabilistic outputs >

Tom: So Njord is positioned as a probabilistic generative model for ocean forecasting that scales to high-resolution grids with irregular shapes, providing calibrated uncertainty estimates for operational use >

Jane: And it sets up some interesting questions about how we can use this kind of AI to help us make better decisions about our oceans and the weather >

The paper's improvements: Tom: So, we’re talking about how they plan to make Njord even better now that they have this initial version >

Jane: Right, Tom, what are the specific tweaks they suggest for future work? They can't just stop at a good result when you’re dealing with something as vast as the ocean >

Tom: The authors say the framework is scalable to large ocean grids by modifying those Interaction and Propagation Networks to use separate dimensionalities for grid embeddings, edge representations, and graph node representations >

Jane: So they are looking at how to handle scaling up the model structure itself so it can manage bigger datasets without breaking >

Tom: And for going even larger, say to global zero point two five degree resolution, they propose a two-stage training schedule: pretraining on one degree resolution data followed by finetuning at that finer resolution >

Jane: That means you start with a big picture view and then refine it down to very fine detail without starting completely from scratch >

Tom: They also mention applying the regional Njord model to the Baltic Sea using quadrilateral graphs instead of the general mesh graphs used for global models >

Jane: That’s interesting because it shows they are tailoring the graph structure specifically to regional geometries, which we already talked about earlier >

Tom: And on ensemble size, they showed that increasing the ensemble size up to twenty members leads to a systematic reduction in the root mean square error across all variables and depth levels >

Jane: That’s concrete evidence that more members in your forecast actually gives you a better prediction, not just more noise >

Tom: So, beyond those specific scaling points, the main implication is that Njord proves probabilistic graph neural networks can handle the complexity of ocean forecasting on a massive scale >

Jane: It moves us past just deterministic answers and into a realm where we can quantify risk and uncertainty in our predictions about ocean currents or sea ice >

Tom: Exactly. The paper shows they’ve built something that matches or beats deterministic machine learning models while providing well-calibrated uncertainty estimates which is crucial for operational use >

Jane: So, the next thing we need to figure out is how this kind of AI can actually get integrated into a real ocean modeling system for decision making >

Conclusion: Tom: So to wrap up, Njord is this model for ocean forecasting that uses probabilistic graph neural networks to give us ensemble forecasts and uncertainty estimates for both global and regional domains >

Jane: That’s right, Tom, it really moves the field from just giving you one answer to showing you the range of possibilities when dealing with chaotic systems like the ocean >

Tom: They showed strong performance against deterministic machine learning baselines on both global and regional scales using a combination of latent variables and graph neural networks >

Jane: And they did it in a way that gives us reliable uncertainty estimates, which is what operational users need to understand risk in things like sea ice or currents >

Tom: Yeah, the calibration of that ensemble spread using the Spread-Skill Ratio shows the uncertainty estimates are actually trustworthy, increasing toward one as expected >

Jane: It’s a really solid piece of work because it shows AI can handle the irregular geometry of the ocean while giving us measurable confidence in its predictions >

Lu: The ability to map gridded inputs to latent representations on a graph structure based on spatial K-means clustering is a powerful way to handle those complex, non-uniform shapes in ocean modeling >

Meng: From an engineering standpoint, getting that level of performance across both global and regional domains is impressive, but I'm still thinking about the computational cost of running twenty ensemble members at once >

Lalam: I’m pretty sure this approach to conditioning the latent variable Zt via KL divergence regularization is what makes the stochasticity in the forecast step actually useful rather than just random noise >

Tom: It seems like Njord successfully scales generative ensemble modeling to high-resolution ocean grids with irregular coastal geometries >

Jane: It certainly sets a high bar for how we can use AI to predict complex, dynamic natural systems and provide those necessary probability distributions >

Lu: The next step in this research direction would be exploring how these graph structures could be dynamically adapted during the forecast itself based on real-time observations >

Meng: If you could optimize the way the embeddings are handled across different dimensionalities, that would make it much easier to deploy this kind of model in a production environment >

Lalam: I’m confident that Njord provides a very reliable foundation for using generative models in forecasting complex physical processes like ocean dynamics >

Tom: We've seen how Njord works and what it achieves on its own, so we'll keep an eye on these scaling improvements as the research continues >

Daniel Holmberg, Joel Oskarsson, Erik Wikingsson, Fredrik Lindsten, Teemu Roos

University of Helsinki · ETH AI Center

cs.LG, physics.ao-ph

Submitted: 2026-05-14

Updated: 2026-10-14

Comments: NeurIPS 2026. Code is available at https://github.com/deinal/njord

Code: https://github.com/deinal/njord

Project page: https://oceanbench.lab

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 92/100

The gist: The gist Njord introduces a probabilistic data-driven model for ocean forecasting that combines deep latent variable frameworks with graph neural networks to provide calibrated ensemble forecasts and

Key concepts

Deep Latent Variable Framework: This is a deep learning structure where the model compresses complex ocean data into a lower-dimensional, hidden representation called the latent variable (Zt). This learned latent state captures the underlying, unobserved dynamics of the ocean system. By sampling from this learned distribution, Njord generates multiple possible future states for forecasting.
Graph Neural Networks (GNNs): GNNs are neural networks specifically designed to operate on graph-structured data. In Njord, the ocean grid points are treated as nodes in a graph, and the connections between them represent spatial relationships. The GNN architecture allows the model to effectively learn how changes at one point influence neighboring points across the entire forecasted area.
Continuous Ranked Probability Score (CRPS): CRPS is a metric used to evaluate probabilistic forecasts. It measures both the accuracy of the predicted mean value and the sharpness or spread of the predicted probability distribution. By optimizing against CRPS, Njord ensures that its ensemble forecasts are not only accurate but also provide well-calibrated uncertainty estimates.
K-means Cluster Meshes: These are specialized graph layouts used to handle irregular ocean geometries. Instead of using a uniform grid, K-means clusters the 3D coordinates of ocean points to create nodes that conform better to the actual shape of the sea surface. This adaptation allows Njord to accurately model regional areas like the Baltic Sea with complex boundaries.

Terminology

Benchmark — 1× in this paper · explained in 2 episodes

Summary

The gist Njord introduces a probabilistic data-driven model for ocean forecasting that combines deep latent variable frameworks with graph neural networks to provide calibrated ensemble forecasts and uncertainty estimates for both global and regional ocean domains.

Model Architecture

Njord combines a deep latent variable framework with a graph neural network architecture, enabling sampling each forecast step in a single forward pass. It employs the hierarchical Graph Neural Networks (GNNs) from Graph-EFM [9] to capture spatial relationships over the forecasted area. The model operates by mapping gridded inputs to latent representations on the mesh graph, where stochasticity from a learned latent variable Zt affects all outputs of the model. For global models, Njord utilizes a clustering-based graph layout constructed via spherical K-means clustering of ocean grid point 3D Cartesian coordinates with latitude-based area weights. This approach places nodes based on the density of ocean grid points to better conform to the irregular geometry of the ocean surface.

Training and Objectives

Njord is trained using a combination of the Evidence Lower Bound (ELBO) and Continuous Ranked Probability Score (CRPS) loss. The full loss function for predictions at time t is defined as L = − Eqϕ log pθ(Xt Xt−2:t−1, Ft−2:t, Zt) + λKL DKLqϕ(Zt Xt−2:t, Ft−2:t)∥ pθ(Zt Xt−2:t-1, Ft-2:t) + λCRPS. Training follows a curriculum in multiple phases, gradually activating the KL term (λKL), introducing multi-step rollout training, and finally enabling the CRPS loss (λCRPS). For Njord-Baltic, the model is trained for 350 epochs using a staged curriculum that begins with pure autoencoder training before introducing the KL divergence term and finally incorporating the CRPS loss.

Handling Physical Constraints

Njord models sea ice by incorporating physical constraints to ensure realistic fields. Instead of post-processing constraints, it uses a density channel mechanism where a binary density channel d ∈ [0, 1] is constructed from Sea Ice Concentration (SIC) as d = I[SIC > 0]. The model predicts this density logit, which is then passed through a sigmoid and thresholded at 0.5 to set the density channel and ice variables to zero in the next state if the predicted density falls below this threshold.

Performance Evaluation

Experiments demonstrate strong performance on both global and regional domains compared to deterministic machine learning baselines. On the global OceanBench benchmark, Njord achieves the lowest errors on average across upper-ocean variables when evaluated against real-world observations, with the largest improvements in surface temperature prediction. For regional modeling of the Baltic Sea at 2 km resolution (Njord-Baltic), it achieves errors comparable to a deterministic baseline while also providing probabilistic forecasts. The model's ensemble spread exhibits good calibration, as reflected by the Spread-Skill Ratio (SSR) which increases steadily toward 1.

Scaling and Extensions

The framework is scalable to large ocean grids by modifying the Interaction and Propagation Networks to use separate dimensionalities for grid embeddings, edge representations, and graph node representations. To scale to global 0.25° resolution, a two-stage training schedule is used: pretraining on 1° resolution data followed by finetuning at 0.25°. The framework can be applied regionally to construct Njord-Baltic using quadrilateral graphs, which are used for regional models. Furthermore, the model's ensemble size comparison shows that increasing the ensemble size to M = 20 yields a systematic RMSE reduction across all evaluated variables and depth levels.

Conclusion

Njord is the first probabilistic generative machine learning model for global and regional ocean forecasting, successfully scaling generative ensemble modeling to high-resolution ocean grids with irregular coastal geometries. The results show that Njord not only matches or exceeds the accuracy of deterministic machine learning models and traditional numerical systems but also provides well-calibrated uncertainty estimates which are important for operational use. This is important for operational use, where understanding forecast uncertainties in the ocean currents or advancing sea ice can help improve decision making.

--- Page 52 ---

The individual ensemble members appear sharp and exhibit noticeable variability, whereas the ensemble mean is smoother due to averaging. Sea-ice fields display well-defined edges and are exactly zero in ice-free regions, reflecting the use of clamping and a dedicated density channel. The ensemble standard deviation is elevated near the ice edge, where its position varies between members. For SIT, uncertainty also extends into the interior of the ice pack, reflecting ensemble disagreements on dynamic thickening processes such as ridging. The physical drivers of uncertainty for salinity and potential temperature are distinctly visible. Salinity variance is overwhelmingly concentrated in the Skagerrak and Kattegat straits. Potential temperature uncertainty is also elevated in these straits but extends further into the Baltic Proper, reflecting complex thermal mixing fronts and internal mesoscale eddies. Surface currents exhibit widespread uncertainty across the entire basin because the Baltic Sea’s surface circulation is heavily wind-driven. Because the Baltic Sea’s surface circulation is heavily wind-driven, ensemble spread in these velocities reflects the chaotic, rapid response of surface waters to varying meteorological forcing across the members. This variance naturally peaks in the narrow, high-flow bottlenecks of the Danish straits. Finally, SLA shows elevated uncertainty not just in the straits but also at the northern and eastern extremities of the basin. In a shallow, enclosed sea, water levels are highly sensitive to wind stress piling water up against the coasts (storm surges and seiches). Furthermore, some visual artifacts remain apparent in the SLA standard deviation. This could potentially arise because SLA fields are derived from interpolating sparse along-track satellite observations.

--- Page 38 ---

Njord samples a single next-step ensemble member in 1 s on one AMD MI250X GPU. In these experiments, we generate a 20-member global ensemble. Lead-time performance is evaluated in Figure 4 for selected variables.

--- Page 39 ---

Figure 29 demonstrates the Zonal current at six different depths. Columns from left to right show RMSE, CRPS, and SSR <ref:2605.

Improvements for AI systems

Bold header: Scalable Graph Construction for Irregular Grids

The improvement involves replacing fixed geometric meshes (like icosahedral) with a spatial K-means clustering of the ocean grid point 3D Cartesian coordinates, with latitude-based area weights to construct graphs, which is noted as better conforming to the geometry of the ocean surface. This allows Njord to operate over the actual ocean grid points and avoids wasting computation on land.

Bold header: Enhanced Physical Constraint Enforcement for Sea Ice

The system can be improved by using a density channel mechanism where a binary density channel d is constructed from SIC, and if the predicted density falls below a threshold, the model is set to zero-ice inputs, ensuring that locations predicted to be ice-free receive clean zero-ice inputs rather than small residual values that may accumulate over rollouts.

Bold header: Improved Temporal Context Utilization

The ensemble forecasting capability can be refined by using a richer temporal context; specifically, the default configuration should utilize two input states (t−2, t−1) and three forcing steps (t−2, t-1, t) instead of simpler inputs to achieve better performance on variables like sea ice.

Abstract

Ocean dynamics are inherently chaotic, yet existing machine learning ocean models produce only deterministic forecasts. We introduce Njord, a probabilistic data-driven model for ocean forecasting, applicable to both global and regional domains. Njord combines a deep latent variable framework with a graph neural network architecture, enabling sampling each forecast step in a single forward pass. We apply Njord globally at 0.25° resolution and regionally to the Baltic Sea at 2 km resolution. To scale to these large ocean grids we introduce K-means cluster meshes that adapt to irregular sea surface geometry. Experiments demonstrate strong performance on both domains compared to deterministic machine learning baselines, while also providing uncertainty estimates from the sampled ensemble forecasts. On the global OceanBench benchmark, Njord achieves the lowest errors on average across upper-ocean variables when evaluated against real-world observations, with the largest improvements in surface temperature prediction.