Weather Generation
Weather Generation
Code: https://github.com/puct9/weather
Ever since implementing the DreamerV3 world model, I had always wanted to do the same to the rain radar precipitation images out of the Bureau of Meteorology. Recently, I had also been sinking my teeth into diffusion and flow matching models, and when I saw the MIRA World Model project get released less than a couple months ago, I figured that there was no better time to try this idea out1.
Data procurement
The National Computational Infrastructure is an Australian public research infrastructure organisation. It publicly hosts an extensive collection of historical data, sourced from various governmental agencies such as the Bureau of Meteorology. Of the historical data, the Operational Weather Radar Network Archive contains radar returns detailing the rate of precipitation at an interval of every 5 minutes from a host of radar stations around the country.
With my hometown being Melbourne, I naturally downloaded the last few years’ worth of radar images for the Melbourne weather station. All in all, this resulted in just over 600K usable images. With those images, some data post-processing was applied to improve the distributional properties of pixel brightnesses by applying a non-linear transform which worked out to be empirically satisfactory on the dataset. The images were then resized from (512, 512) to (128, 128) pixels to reduce the memory and storage footprint and stored with u8 values.
Model architecture
Most of the model elements and architectures borrowed heavily from the implementation choices in MIRA World Model.
The weather generation model is a type of latent diffusion model. Latent diffusion models are the state of the art in terms of computational budget and generation quality. The former property is especially handy for me, as I don’t have infinite money to spend on compute and did most of the training on my personal computer.
Typically, latent diffusion models are trained in 2 stages:
- Train a variational autoencoder (VAE) on the images in a large dataset
- Train a diffusion or flow matching model on the latent space
More recently, the preferred way of training these models has been to move from the VAE to the representational autoencoder (RAE). A representational autoencoder utilises a more modern form of image encoder, trained with self-supervised methods that do not involve image reconstruction like DINO and JEPA. Hence, we train latent diffusion models utilising RAEs in 3 stages:
- Train or download the latest and greatest DINO or JEPA encoder
- Freeze the encoder and train a decoder using their representations by first further compressing the representations into a new latent space
- Train a diffusion or flow matching model on the new latent space
While MIRA used a pretrained version of DINOv3, I found that decoder reconstructions using pretrained DINOv3 embeddings were, for some reason, either not much better than or inferior to using an untrained (randomly initialised) model. Even more interestingly, the latent space produced by using the randomly initialised model admits substantially better generations than a model that directly maps the input image to the latent space through a linear projection.
The decoder trained to reconstruct the images from the latent space is an encoder-only transformer, as is the flow matching model. The patch size used was (8, 8), yielding $16 \times 16 = 256$ patch tokens for the decoder. In keeping with MIRA-WM, the latent space halves the number of tokens for each axis, resulting in $8 \times 8 = 64$ tokens. The spatial attention used a factorised 2d RoPE by rotating one half of the key and query dimensions for the height and width axes.
Training infrastructure
Due to being run on WSL on a machine with only 16GB of memory (so I really only had 8GB of memory to play with), memory constraint was a key consideration for all parts of the training infrastructure, which is famously almost as hard of a problem as the NN modelling challenge itself.
The size of the dataset consisting of ~600K images uncompressed was over 9GB. While this would have been a non-issue when I had 4TB of memory and a H200 node all to myself at my previous workplace, I unfortunately no longer have this luxury. This meant that live dataloading from the disk was almost inevitable. Fortunately, I was the happy owner of a fairly quick NVMe drive which should have had enough random read performance to sustain the training workloads. However, this would be the first major obstacle.
While I had initially planned on using NumPy’s np.memmap facility, it became apparent that, out of the box, it was only capable of delivering a respectable read speed for ordered reads. With some research, the approach of using a madvise was initially trialed.
import mmap
import numpy as np
data = np.memmap(file, dtype=np.uint8).reshape(-1, 128, 128)
data._mmap.madvise(mmap.MADV_RANDOM)
img = data[idx]
While this did improve the random read performance that is typically required when training a model, it was still ultimately unsatisfactory. Instead, it turned out that calling os.pread on an opened file was substantially faster for random reads. This approach, and prefetching the next batch were used for all dataloaders in the project were sufficient to keeping the GPU busy at all times. These approaches were also used for dataloading sequences from the latent cache for training the flow matching model.
import os
import numpy as np
bytes_per_image = 128 * 128
buf = os.pread(fileno, bytes_per_image, idx * bytes_per_image)
img = np.frombuffer(buf, dtype=np.uint8).reshape(128, 128)
Video generation
Video frames are iteratively decoded in a process of diffusion forcing whereby a flow model is conditioned on previous frames’ latents to generate the next latent and hence frame.
A form of classifier-free guidance to improve generations was attempted but to limited effect. This was done by computing the model with and without past-frame conditioning at various points by substantially re-noising the previously generated frames. Auto-guidance was also not implemented for the final model to reduce generation time.
The video shown on the title of this post consists of $6 \times 10 = 60$ individual generations of 200 frames at 32 denoising steps for each frame. $60 \times 200$ frames were generated in about 2:30 excluding compile time, giving a generation speed of about 80 frames per second on a single RTX 2060 Super (TU106).
Deviations
A number of deviations from the MIRA-WM training recipe were employed for reasons primarily relating to compute capacity.
- A DINO encoder was not used (although experiments showed that it may achieve similar results) to reduce the compute required for training the decoder. This was also in part because the DINO caches require a large amount of storage—if we use $256\times 256$ pixel images2, we get $16\times 16 = 256$ patches, where each patch has a dimension of $1024$, yielding $256\times 1024 = 262144$ embedding elements for each image, which is 16 times larger than the number of pixels in the original $128 \times 128$ image to begin with, and that’s not counting for needing to store the embeddings in probably more than 1 byte per element too!
- The LPIPS and LDINO perceptual losses in decoder training were omitted to save training time by 80%.
- The 2x temporal downscale at latent and upscale at decode were omitted because it seemed that MIRA-WM only used it to achieve higher FPS so that the world model could be run in real time, rather than actually needing it for model quality.
- For the same reasoning, the few-step flow maps were omitted.
- No action conditioning was required.
Deficiencies
Because I am bad at machine learning, the model would sometimes do bad things. Here’s an example of a failed, “collapsed” generation. The model would occasionally find itself in this state, and would not seem to be able to snap out of it.
Previously, it was mentioned that using the features out of a randomly initialised (untrained) model and then linearly projecting it to the latent space was superior to directly linearly projecting the pixels to the latent space in a patch-by-patch fashion. This is especially curious since the untrained model could not have been a profound feature extractor, and had a model dimension higher than the number of pixels per patch, so it wasn’t acting as an informational bottleneck either. The difference between approaches becomes pronounced in this stage, where generations are more likely to collapse, on top of collapsing in fewer frames when they do.