Variational Autoencoder for Hydrological Stochastic Generation
Latent Space Explorer
The VAE encodes each daily rainfall field into a 64-dimensional latent vector. The scatter below shows a PCA projection of 500 encoded days. Hover a point to see the spatial rainfall pattern it decodes to - demonstrating how distinct weather regimes cluster in latent space.
Background & Motivation
Traditional stochastic weather generators can't produce spatially co-variant rainfall across multiple sites while preserving temporal dependencies — a problem for water supply systems like Auckland's, which span several catchments. This project explores convolutional VAEs as a unified alternative.
Methodology
A convolutional encoder-decoder learns a compressed representation of daily rainfall fields pooled across seven CMIP6 GCMs. Two architectures are compared: a continuous β-VAE with a GRU temporal transition model, and a two-level VQ-VAE-2 with discrete codebooks — see Model Formulation below for the full objective and data details.
Results
The continuous VAE (900 epochs) reconstructs plausible aggregate statistics, but its latent space is heavily degenerate — PCA shows ~99% of variance on a single axis. The VQ-VAE-2 (20 of a planned 500 epochs) shows a far healthier 60/34% split. See Latent Space Diagnostics below.
Lessons Learned
The original Gaussian VAE was effectively learning to predict the average rainfall pattern rather than distinct weather regimes — a common failure mode on skewed, mostly-zero precipitation data. Switching to a discrete codebook forces the model away from this degenerate solution — full explanation below.
Model Formulation
The continuous VAE is trained to maximise a conditional evidence lower bound (ELBO), where the prior over the latent state is itself learned by a GRU-based transition model rather than fixed to $\mathcal{N}(0, I)$, so that it can capture day-to-day persistence:
$$\mathcal{L}_{VAE} = \mathbb{E}_{q_\phi(z_t|x_t,c_t)}\big[\log p_\theta(x_t \mid z_t, c_t)\big] - \beta \, D_{KL}\big(q_\phi(z_t|x_t,c_t)\,\|\,p_\theta(z_t \mid z_{t-1}, c_t)\big)$$
Reconstruction uses a per-pixel Gaussian log-likelihood over the encoded spatial field, and $\beta$ is annealed linearly from 0 to 1 over the first 30 epochs to avoid posterior collapse. Conditioning $c_t$ is a day-of-year cyclic encoding, and the encoder/decoder are convolutional over a 60×60 grid covering the Auckland water supply catchments, at 5km resolution, built from bias-corrected CCAM output for ACCESS-CM2, AWI-CM-1-1-MR, CNRM-CM6-1, EC-Earth3, GFDL-ESM4, NorESM2-MM and NZESM under the historical scenario.
The VQ-VAE-2 instead compresses each field into two levels of discrete codes (a coarse "top" level and a finer "bottom" level conditioned on it). Each encoder output is snapped to its nearest neighbour in a learned codebook of $K=512$ embeddings, with gradients passed through the non-differentiable lookup via the straight-through estimator:
$$\mathcal{L}_{VQ} = \|x - \hat{x}\|_1 \;+\; \|\text{sg}[z_e] - e\|_2^2 \;+\; \beta_c\,\|z_e - \text{sg}[e]\|_2^2$$
where $\text{sg}[\cdot]$ denotes stop-gradient, $e$ is the selected codebook embedding, and $\beta_c = 0.25$ is the commitment cost. Codebook entries are updated via exponential moving average (decay 0.99) rather than gradient descent, and dead codes (cluster size below 1) are periodically re-initialised from live encoder outputs to prevent codebook collapse.
Latent Space Diagnostics
To compare the two latent spaces directly, 500 days are encoded and projected to 2D with PCA. For the continuous VAE, this recovers the posterior mean $\mu$ per day directly. For the VQ-VAE-2 — which has no single continuous vector per day — the pre-quantization top-level features are global-average-pooled into a summary vector instead.
After 900 epochs, the VAE's projection concentrates ~99% of variance on a single principal component: effectively, the model has learned a one-dimensional "how wet was this day" axis rather than a rich space of distinct spatial patterns. The VQ-VAE-2, despite only 20 of a planned 500 epochs of training, splits variance roughly 60%/34% across its first two components — a much less degenerate structure. It has no learned prior over the discrete codes yet, so the explorer above visualises reconstructions of real days rather than freely generated samples.
Why Move to a Discrete Latent Space?
The original design used a standard β-VAE recipe: encode to $(\mu, \log\sigma^2)$, reparameterise, and reconstruct with a per-pixel Gaussian negative log-likelihood. That works well for smooth, roughly-Gaussian fields, but precipitation is heavy-tailed and mostly zero — a loss that minimises expected error under uncertainty is satisfied by predicting something close to the local average, not a specific plausible day. In practice this showed up as the near-degenerate latent space described above, and reconstructions that looked more like a smoothed climatology than an individual weather pattern.
Switching to a VQ-VAE-2 replaces the continuous Gaussian bottleneck with a discrete codebook. Because the model must reconstruct from one of a finite set of learned codes rather than a continuous mean, it can no longer default to "predict the average" — it has to commit to a specific pattern. Early results support this: a 60/34% PCA variance split versus 99/0.4% for the original VAE suggests meaningfully different rainfall regimes are being captured, rather than just overall wetness. That said, this model is only 20 of a planned 500 epochs into training, and a learned prior over the discrete codes still needs to be added before it can generate free-running sequences the way the original VAE does.