Credits: Gilles Louppe
Poseidon: Three Changes
for a better Prior
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.

Progress

We want to improve the results of the paper, and learn what works best for our other projects:

  • Rewrite: Poseidon rebuilt on azula, on the shelf and with the variables of the paper.
  • Overfit study: days of summer 2010, to check that the prior can produce realistic states.
  • Mutation-style search: many runs, sometimes from scratch, but often from the best one with one change, reusing its weights to save time.

Daily standardization

In the paper, each channel is standardized by a mean and a standard deviation computed over all days. In summer, the eddies of the day ( to ) are then a small fraction of a deviation. Instead, we standardized each state by the statistics of its day of the year:

Surface, 1 July Global statistics Daily statistics
Temperature []
Oxygen []

We took the ViT trained from scratch for k updates on global statistics, and trained it for k more updates on daily statistics.

Figure A. Surface temperature [], each map with its own – quantiles.

Note: the daily ViT also trained longer; a fair test would train both from scratch for as many updates. The filaments appeared within 16k updates.

From a U-Net to a ViT, then a U-ViT

Here we explore a vision transformer (ViT), which turns each patch of values into one token of features. Its prior samples showed the grid of the patches, so we added a convolutional head after the ViT, at full resolution, to smooth these discontinuities.

We then added a convolutional encoder before the ViT, which coarsens the states once, so that the tokens are built from learned features rather than raw pixels, and a decoder after it: the ViT becomes the bottleneck of a U-Net, the U-ViT.

Figure B. The U-ViT; the coloured lines show which parts make the ViT, the ViT with its head, and the U-ViT.

Figure C. Spectrum of surface temperature, mean of generated days; validation loss along each training.

Key result. The U-ViT follows the simulated spectrum, and goes below the plateau the ViT kept for about k updates.

Samplers

We were curious whether the remaining blur comes from the sampler. With the U-ViT (k updates), we generated two memorized days from the same initial noise with the samplers of azula: deterministic (Adams–Bashforth of order , Heun), stochastic (DDIM with , DDPM) and predictor-corrector (one Langevin correction per step), from to steps.

Figure D. Error to the memorized day (temperature) and small-scale power of oxygen (wavenumbers , simulation ).

Key result. The deterministic sampler of order does best: the lowest error and the most small-scale power.

Figure E. Surface temperature [] from the same initial noise, steps, each map with its own scale.

Problem. The same slight blur with every sampler: it does not come from the sampler, but from the denoiser.

Other changes

Also helped

  1. Input projection in float32: at .
  2. Learning rate , and a wider ViT (, blocks).
  3. Fine-tuning from the best weights at each change.

Did not help

  1. Patches of or : slower, worse at –.
  2. Date as input channels: equal or worse.
  3. Float32, , more low-noise training: no change.
  4. Variance exploding schedule: better at the extremes, worse at –.

Summary

Key results

  1. The U-ViT follows the simulated spectrum.
  2. The deterministic sampler of order does best.

Problems to solve

  1. A slight blur remains, from the denoiser: of the small-scale power of oxygen.

What's next

  • Sharper samples: more depth and more channels at full resolution (runs in progress).
  • Full training: the chosen U-ViT on 1998–2017, on A100.
  • Posterior: skill, spread and SSR per depth, and hypoxia, as in the paper.