[Paper] AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning

[jacobbieker]

Arxiv/Blog/Paper Link

Paper: https://arxiv.org/abs/2308.13280
Code: https://github.com/clessig/atmorep

Detailed Description

Self-supervised representation learning that uses historical record and observations to constrain the representation of the atmosphere. The point of the model is to feed into downstream networks that take this representation as input (I.e. instead of analysis fields like most weather models, or direct observations, like in #64

Some interesting thoughts:

1. They also struggle with getting areas from multi-dimensional NWP data efficiently, and try to mask the loading through data-parallel loading and making batches in parallel in memory.
2. MultiFormer is their name for having one essentially vision transformer per input field (i.e. temperature, u wind, v wind, etc.) that is combined into one large model. The rationale for the multiple transformers is that each one can better learn the processes for how that field varies in time and space, and can have individually smaller transformers that can be swapped in and out for the final model predictions. So if you don't need the temperature, for example, then you don't need to load the temperature transformer, and the rest of MultiFormer would still run. They use cross-attention heads to have the different transformers communicate to each other.
3. U-Net like connections between the encoder and decoder are used to keep multi-resolution features.
4. Trained with a masked-reconstruction loss, has an internal ensemble prediction that gives probabilities for reconstructing the input masked tokens.

[jacobbieker]

@JackKelly another group of people who've had to work around getting random data in high-dimensional datasets. I think this is what we tried in nowcasting-dataset?

[JackKelly]

Oooh, very nice, thank you!