Heimdall, Neptune and Poseidon

Credits: Gilles Louppe
Understanding the Ocean
from Observations Using
Deep Learning
PRESENTED BY Victor Mangeleer1–4
PHD SUPERVISORS Marilaure Grégoire1,3 and Gilles Louppe2,3 HOST SUPERVISOR Laure Zanna4 IN AFFILIATION WITH Modelling for Aquatic Systems1, Science with AI Laboratory2,
University of Liège3, New York University4.

What can we predict from space?

Progress

Since our last meeting, here is what we worked on:

  • Encoders: the encoders, for GLORYS and for the satellite products, are trained.
  • Latents: every dataset is encoded, with its statistics computed.
  • Experiment: a study to assess whether the observations are helpful, now training*.

* Note: we assume it is already worth showing in the latent space whether the observations help, even though a decoder would also let us see it on the fields.

Scientific question

In the latent space, does a model given the surface observations predict the states of a block of days better than the same model without them, both knowing the states on each side of the block?

Experimental setup

We learn , with and *:

  • Observations: each model is trained with and without the observations .
  • Models: seeds each, updates on the latents of 2020–2021.
  • Baselines: the linear interpolation between the sides, close days being close states, and a climatology.
  • Evaluation: skill, spread, SSR and CRPS on days of 2022–2023, members, the same noise.

* : the number of states of the block to predict; : the number of known states on each side of the block, used for its prediction.

Figure A. CRPS on the central state of the block, on our validation set.

The linear interpolation gives a single sample: its CRPS reduces to the mean absolute error.

Figure B. An example of the two figures we will draw once all the experiments end, with made-up values: left, the gain from the observations; right, how the error of every state changes with .

Choice. Is this comparison, with and without observations, enough to tell whether they help?

Attention

To run this experiment, we use a transformer, whose attention we had to define. Its attention is local (Neptune and Poseidon showed that a local attention works well), following STRATA: each block attends within windows of GLORYS tokens, holding the states of every day and every observation token whose nearest GLORYS token, on the sphere, lies in the window. The windows are shifted by half a window every other block.

Figure C. The processor in the latent space of the encoders: every source becomes tokens of channels.

* Latent grids: SST , SSS , wind , chlorophyll ; SSS and wind have an ascending and a descending pass.

Figure D. What a state to predict has read, block after block: the windows alternate with shifted windows, and its view grows.

Choice. Windows over every day and source, so that sides and observations are read in one block.

Summary

Choices to make

  1. Is the comparison with and without observations enough?
  2. Windows over every day and source.

What's next

  • Observations: do they help locally, and up to which block size?
  • Decoder: a GLORYS decoder, to score in physical space and tell the surface from the depth.
  • GiBBS: sample the global state with the learned local posterior.

Neural emulator of Black Sea dynamics

Progress

Since our last meeting, here is what we worked on:

  • Encoders: the three encoders, surface, ocean and ERA5, are trained.
  • Latents: the whole dataset, 1998–2022, is encoded, and its statistics are computed.
  • Decoders: the surface decoder is trained (MSE), the ocean decoder waits*.
  • Processor: experimenting to find a good baseline for the prototype.
* Our fair-share on Lucia makes jobs wait for hours…

Experimental setup

Each design of the processor is compared in the same setup:

  • Training: one year of latents, 2008, for min on one A100, times per design*.
  • Evaluation: on the whole validation set, 2018–2020, with members.
  • Metric: the CRPS over the standardized latents, against persistence.
* Models should normally be compared at an equal number of updates rather than an equal training time.

Processor

The design of the processor will shape every forecast of the emulator. Before scaling it up, we studied two of its choices: the attention mechanism, which tokens read each other, and the target, what the network outputs.

Figure E. The emulator: encoders, processor and decoder, with the two choices we studied numbered.

Attention mechanism

In each block, every token updates itself by reading other tokens. In one day, a current only travels to km, two to four latent pixels. Our hypothesis is that a token only needs to attend to what happens around it, without information from far away: a local attention should be enough.

Figure F. Visualization of the three different attention mechanisms.

Credits: Hu et al., Scaling Storm-Resolving Atmospheric AI Simulation to the Entire Planet, 2026. Windows are fixed tiles of the grid, shifted by half a window every other block.

Figure G. CRPS of the three attention mechanisms on fixed validation days, during training.

Key result. At equal training time and capacity ( M parameters), local attention reaches a lower CRPS, and each update is faster.

Which target?

The processor can output the next latent state or its increment . Which one is easier?

Figure H. CRPS on the whole validation set after min of training: one point per run, the mean as a line, and the relative gap between the means.

Even if we keep the state here, we will also train the increment for longer and compare decoded trajectories.

Decoders

We train decoders, from the latents back to the physical fields; so far, only the surface decoder is trained. It mirrors the encoder, with wider channels in its two deepest stages, and is trained on 1998–2017 with a mean squared error averaged per variable and level over the sea.

Figure I. The surface fields of 30 consecutive days, decoded from their encoded latents.

Summary

Key results

  1. A local attention reaches a lower CRPS than the global one.

What's next

  • Position: the absolute embedding against RoPE, running now.
  • Full training: a larger processor on the 20 training years.
  • Decoders: the ocean decoder, then a fine-tuning on latents noised like the processor's error.
  • Rollout: a multi-step fine-tuning, to keep 10-day trajectories stable.

Poseidon

Introduction

We redid the experiments of the paper* with a vision transformer (instead of a U-Net), to verify our results, learn how to design transformers, and get ideas for the assimilation module of Neptune. Three of the most interesting observations:

  • Compression of the tokens: how wide the transformer must be for its data.
  • Loss weighting: which voxels of water the loss should favour.
  • Retraining the prior: the chosen model, on years of data.

Experimental setup

Each model approximates the prior and trains on the days of summer , from scratch with the same budget (k updates of states on one A100). We compare the models with two diagnostics:

  • Denoising loss: at a fixed noise level , we noise of these days, denoise them and measure the error.
  • Samples: for a training date , the model should reproduce the same state ; we measure the RMSE between the two.

Figure J. Our U-ViT: convolutions compress the state, then a transformer processes its tokens; the dashed box holds our question.

Compression of the tokens

Each token carries values: the projection maps them to features (1), and the blocks work with (2). We compare three cases:

  • Baseline, : the tokens are compressed in every block.
  • Entry compression, , : compressed at the entry only.
  • No compression, : the tokens keep all their values.
Note: the fourth case, W1 = 2048 and W2 = 1024, was not tested.

Figure K. RMSE of the samples per variable for the three cases: the green model is the best everywhere.

Key result. No compression, , is the best everywhere: widening the tokens again after a compression creates no information.

Note: RMSE against the simulation on the 90 training dates, in standardized units, one run per case: it measures the fit, not the generalization.

Loss weighting

We test three ways of weighting the loss between the voxels of water: None, where each voxel counts the same, Level, where each depth level counts the same, and Column, where each water column counts the same.

Figure L. The three weightings on a small shelf, coast on the left: the colour is the weight of each voxel in the loss.

Figure M. RMSE of the samples with depth and per variable: the weighting per level of the paper does worse than none or per water column.

Note: RMSE against the simulation on the 90 training dates, in standardized units, one run per case: it measures the fit, not the generalization.

Retraining the prior on 20 years

  • Data: all the training days, –.
  • Model: , heads, Column weighting, from to .
  • Size: M parameters ( M for the U-Net of the paper).
  • Training: k updates of states on A100, about h (the paper: days on A100).

Figure N. Training and validation losses, then validation loss per noise level: over the whole training (131k updates).

Note: the validation loss is the denoising loss on 256 fixed days of 2018–2020, at 8 fixed noise levels σ = 10−5 to 102 (the left panel averages them), always with the same noise; the training draws σ at random.

Figure O. A prior sample at the end of the training for 15 August 2019 (validation), under the simulation of that day.

Summary

Key results

  1. No compression of the tokens is the best; widening them again after a compression creates no information.
  2. The weighting per level of the paper does worse than none or per water column.

What's next

  • Prior samples: judge the retrained prior on its generations once trained, maps first.
  • Posterior: MMPS on validation then test, with skill, spread and SSR per depth.
  • Hypoxia: check the gain of the Column weighting at the bottom of the shelf, where hypoxia forms.

Sources

Appendix

A
B

Figure P. The latent processor with windows and columns, with the shape of every tensor.

C

Figure Q. The full prior, with the shape of the tensor after each stage (channels × latitude × longitude).