Transforming Weather Data from Pixel to Latent Space

Summary

The paper proposes to move weather modelling out of pixel space and into a shared, compact latent space produced by a Weather Latent Autoencoder (WLA). The conceptual move has three parts: (i) a Pressure-Variable Unified Module (PVUM) uses a hypernetwork over variable/level metadata to project any pressure-variable subset (PVS) into a common feature space; (ii) a VAEformer-style encoder–decoder pair compresses that unified feature; and (iii) a Binary Quantization Module (BSQ) maps the latent onto a unit sphere and binarises it into hardware-friendly bitwise tokens. The authors then run downstream forecasting entirely in that latent space (Window-ViT-Latent), leaving high-resolution pixel data only for final metric computation. On ERA5 the pipeline is reported to shrink 244.34 TB to 0.43 TB (a 566.3× headline ratio), to beat five compression baselines on weighted RMSE and BPSP at higher compression, to generalise to unseen pressure levels and HRES, and to yield sharper extreme-value forecasts (SEDI/RQE) than a matched pixel-space Window-ViT baseline. The framing — one autoencoder that unifies heterogeneous PVS so downstream models operate on a canonical low-rate representation — is genuinely useful for the community: it targets both the smoothness problem in pixel-space forecasting and the storage/PVS-brittleness cost that hampers medium-range ML weather models. My overall assessment is that the direction is sound and the compression numbers are architecturally consistent — the shape-and-quantisation arithmetic behind 566.3× reproduces almost exactly — but the case that the framework itself is responsible for the sharpness advantage is under-argued theoretically (a fixed decoder does not guarantee sharp outputs when the predicted latent is smoothed) and under-controlled empirically (the WTG adversarial control is not tuned). Reproducibility is the biggest concrete concern: the release is model-definition code only, and execution shows the central BSQModel is unconstructible as shipped, the downstream forecasters, baselines, ERA5-Latent artifact, metric harnesses, and pretrained checkpoints are all absent, so the trained-reconstruction and forecasting numbers cannot be verified end-to-end; the compression-ratio numbers that did reproduce did so with random weights.

Strengths

• Conceptual contribution. Decoupling weather reconstruction from prediction, and unifying arbitrary PVSs through a metadata-conditioned hypernetwork, is a clean idea and directly addresses two pain points (single-PVS models, over-smoothed forecasts) that hobble current pixel-space systems.
• Storage story has architectural backing. Independent recomputation of the encoder token grid at ERA5's full 721×1440 resolution recovers 566.30× / 0.4315 TB against the paper's 566.3× / 0.43 TB, and the WLA_small row's 1251.7× / 0.026 BPSP likewise recovers exactly — so the shape-and-quantisation arithmetic of the compression claim is internally consistent.
• Careful rate–distortion sweep. The Fig. 8 ablation over five codebook sizes ($2^{16}$–$2^{128}$) crossed with three level counts (6/13/25) is exactly the analysis one wants for a quantised autoencoder, and the discrete-vs-continuous bitrate-matched comparison in Fig. 8(b) is a helpful design justification for BSQ.
• Honest scope on precipitation. Rather than papering over it, the authors flag that the latent precipitation model (WTLP) remains below Operational-IFS on RQE, which is the appropriate level of candour for a purely data-driven system.
• Cross-modality generalisation experiments. The Section 4.3 study on unseen pressure levels and downsampled HRES data is the right way to probe whether PVUM's metadata conditioning actually generalises — even if the current evidence is thin (see Weaknesses).

Weaknesses

• Mechanistic argument for sharpness is thin. Sections 1 and 3.1 assert that because reconstruction is decoupled, task uncertainty has 'minimal impact' or 'does not affect' extreme-value reconstruction. But the WLA decoder is deterministic conditional on the predicted latent, so if WTL predicts an averaged/smoothed latent (the standard outcome of MSE/BCE training), the decoder will faithfully reconstruct that smoothed latent. The paper needs either an argument (e.g., that BSQ's near-categorical support makes averaging less harmful) or a direct comparison of predicted-vs-true latents at long lead times.
• Adversarial control (WTG) is not a fair control. The Section A.5 ablation adds a GAN loss to WT while keeping all other hyperparameters identical to the MSE baseline. GAN training is famously sensitive to loss weighting and discriminator schedule; concluding from this single un-tuned run that adversarial training in pixel space is inherently non-trivial overreaches — it is at least as consistent with 'this particular weighting was bad'.
• Aggregate 566.3× is not decomposed and is accounting-sensitive. Table 1's per-modality ratios (e.g., 257.88× for WLA upper-air, 625.9× for precipitation) cannot be trivially combined into 566.3×, and the paper does not show how the ERA5-Latent PVS mix plus the 0.117 TB LZMA pixel companion sum to 0.43 TB. Attempting the arithmetic from the eight PVS subsets listed in Section A.3.1 recovers ~197.86× and ~1.235 TB, not 566.3× / 0.43 TB — the headline ratio holds only under a maximal-single-PVS accounting.
• Foundation-model compression comparison is asserted but not run. The related-work paragraph closes with 'WLA leads in data compression efficiency' relative to Aurora / GenCast / Prithvi-WxC, but none of those systems appears in Table 1 or any compression benchmark, so the ordering is not evidenced.
• PVUM description is under-specified. The mapping from $C_1$ metadata tokens to a $C_1 \times C_2$ weight matrix $W$ is described as coming from 'another linear layer' — whether this is a per-token projection producing one row of $W$ each, or a single projection over a pooled representation with a reshape, is left implicit. Reproduction depends on the choice.
• Baseline attribution is inconsistent. Table 1 attributes VQVAE/VQGAN to Mirowski et al. 2024, whereas Section 4.2 attributes them to van den Oord 2017 / Esser 2021 (and mis-spells 'VQGAN' as 'VAGAN'), while Appendix A.1.3 says the authors fine-tuned them themselves. A reader cannot tell whose numbers are shown.
• FourCastNet is treated as a benchmark but never set up. Section 4.5.1 declares IFS and Pangu-Weather as benchmarks; the discussion then mentions surpassing FourCastNet with no citation, checkpoint, or configuration, and repeats this in the precipitation section.
• Codebook choice contradicts the compression-first framing. The final model uses $2^{128}$ — the largest codebook tested — chosen because increasing further gives diminishing returns in reconstruction. The trade-off with compression ratio (the paper's headline) is not discussed, and $2^{128}$ over $2^{96}$ is not justified on the rate-distortion Pareto frontier.
• Generalisation claim rests on one variable. Fig. 5 shows temperature only; extending 'strong generalization' to all variables and to HRES on the basis of a single-variable thermodynamic curve is a large inferential leap, and unseen-level and HRES OOD effects are conflated in one figure.
• Extreme-value tension in Section A.7. The appendix acknowledges that reconstruction error is 'systematically correlated with regions of extreme values' but concludes the model is 'highly effective at preserving' those same extremes; no magnitude numbers (e.g., MAE inside the >99th-percentile mask vs. everywhere) reconcile the two.
• Precipitation superiority claim comes from one metric. WTLP's 'stronger capability in characterizing the heavy-tailed distribution' is inferred from RQE alone; RQE measures preservation of extreme magnitudes, not distributional characterisation.
• 'Superior' visualisation without a comparator. Section A.7 concludes 'superior reconstruction quality' from side-by-side ERA5 / ERA5-Latent panels that contain no baseline reconstruction.
• Terminology. The abstract/intro count of 164 'variables' is really 164 variable-level fields ($6 \times 25 + 8 + 6$); the related work uses 'fields' for the analogous CRA5 count.
• Figure numbering. 'Figure S1' is referenced for both the precipitation RQE plot (Sec. 4.5.2) and the appendix visualisation series (Sec. A.7), and Figure S2's caption is copy-pasted from Fig. 6(b) and does not describe the WTG ablation it supposedly plots.

Reproducibility & code

The release is model-definition code only. The reproducibility statement is a bare 'available at WLA' anchor with no URL, DOI, or manifest, and the README's TODO boxes for the pretrained model, the ERA5-Latent dataset, and the training code are all still unchecked. On this basis:

• Headline compression numbers reproduce architecturally, not empirically. Running the encoder on the paper's exact 15-cell grid at full 721×1440 resolution recovers 566.30× / 0.4315 TB (vs 566.3× / 0.43 TB) and 625.87× / 0.05113 BPSP for WLA-25L (vs 625.9× / 0.051) and 1251.7× / 0.02556 BPSP for WLA_small (vs 1251.7 / 0.025). All of this used random weights, because no checkpoint ships; the numbers reflect shape and 1-bit-per-dimension BSQ arithmetic, not a trained model.
• Central BSQModel is unconstructible as shipped. The paper's central class cannot be instantiated because the transcoder.scheduler subpackage is not in the release, and the downstream forecasters bsq_era5_fengwu.py / bsq_era5_fengwu_latent.py fail to import for the same reason. Compression measurements had to go through a hand-reassembled surrogate encode path built from ViT_Encoder + quant_embed + BinarySphericalQuantizer.
• Reconstruction-quality half of Table 1 could not be checked. With no pretrained checkpoint, no ERA5 loader (zero torch.utils.data.Dataset subclasses and zero xarray/zarr/netCDF imports in the tree), no baseline code (zero source hits for ELIC/IEN/VQVAE/VQGAN/VAEformer), and no wired weighted-RMSE/BPSP evaluator (the cos-latitude weighting exists inside losses/vqperceptual.py but is buried in training_step), none of the seven weighted-RMSE cells and none of the baseline rows could be produced. The paper's ordering claim ('lowest weighted RMSE on all seven variables against all five baselines') is entirely unverified.
• Downstream forecasting evidence (Figs 6, 7, S1) could not be exercised at all. No SEDI or RQE implementation (grep returns zero hits), no autoregressive 7-day/6-hour rollout, no Pangu-Weather/FourCastNet/IFS integration, no ERA5-Latent inputs, and no WT/WTL training loop — only a single-step $t \to t+1$ smoke test with random weights. The claimed SEDI/RQE advantages of latent-space forecasting are effectively not testable from the release.
• PVUM implementation may not match the paper. The primary encoder path uses fixed per-modality nn.Embedding altitude tokens; a TransformerWeightGenerator exists but is not clearly wired into the run path. Because the metadata-hypernetwork story is central to the Section 4.3 generalisation claim, this is a substantive uncertainty about which module produced the reported numbers.
• HRES generalisation pathway is unimplemented. No HRES ingestion, no downsampling script, no per-pressure-level evaluator, and no dates/cycles pointer accompany Fig. 5.
• Ablations partially reproduce. The compression-ratio axis of Fig. 8(a) recomputes cleanly on the full 15-cell grid (all rows within ~0.14% of Eq. 1), but the reconstruction-quality axis is null everywhere, and the continuous-float32 baseline for Fig. 8(b) is not implemented.
• GAN-loss ablation cannot be re-run. Discriminators ship, but vqperceptual.py depends on transcoder.losses.lpips (absent) and StyleGAN utilities that are not vendored; the WT backbone against which the GAN loss is added is not released.
• LZMA companion is unverifiable and code was inactive. The July-2023 pixel artifact does not ship, no manifest sizes it, and the compress_binary_array LZMA path in bsq_era5_cra5.py is entirely commented out — it had to be re-activated verbatim before it could be exercised.

Recommended Changes

Essential

• Strengthen the sharpness argument. In response to the mechanism weakness, add either a theoretical argument or an empirical diagnostic — e.g., the L2 distance between the true and WTL-predicted latent codes for extremes at 3-, 5-, 7-day lead times — that explains why latent prediction sharpens outputs when a deterministic decoder is applied to a predicted latent, not just to a ground-truth one.
• Ship the missing pieces so the central model imports. In response to the BSQModel-unconstructibility weakness, add the missing transcoder.scheduler subpackage (and any other unresolved imports such as main.instantiate_from_config) so BSQModel and the two forecaster entry points can be instantiated from the shipped code.
• Release the ERA5-Latent build pipeline and at least one checkpoint per modality. In response to the ERA5-Latent-artifact reproducibility gap, ship build_era5_latent.py (ERA5 fetch → normalisation → WLA encoding → BSQ bit-packing → manifest), the dataset (or a hashed downloadable pointer), and at least one pretrained WLA checkpoint per modality so the pipeline can be exercised end-to-end without training.
• Release Window-ViT, Window-ViT-Latent, and SEDI/RQE. In response to the downstream-forecaster reproducibility gap, publish both forecaster architectures, their training scripts, the WLA checkpoint they consume, an autoregressive 6-hour rollout loop, the SEDI (per-percentile) and RQE (per-quantile-grid) implementations, and either the exact IFS/Pangu-Weather/FourCastNet references or the frozen forecast fields used.
• Release baselines and metrics for Table 1. In response to the Table-1 reproducibility gap, ship CompressAI configs for the retrained Elic/IEN, the fine-tuning script and hyperparameters for VQVAE/VQGAN, the exact VAEformer checkpoint pointer (with hash), and the weighted-RMSE and BPSP evaluation code.
• Decompose the 566.3× aggregate and reconcile with the PVS mix. In response to the aggregation weakness, add a per-PVS accounting table (raw bytes, compressed bytes, ratio, contribution) that reproduces 244.34 TB, 0.43 TB and 566.3× line by line, including the 0.117 TB LZMA pixel companion, and clarify whether the eight PVS subsets share underlying tokens (deduplication) so that a reader can reconcile the maximal-single-PVS 566.3× with the ~197.86× figure implied by summing all subsets naïvely.
• Confirm and release the PVUM used for the reported numbers. In response to the code/paper mismatch, state whether Fig. 5 and Table 1 numbers were produced by the hypernetwork-based PVUM or by the fixed altitude-embedding path, and ship the actual module used.
• Redo the WTG control with a small hyperparameter sweep. In response to the adversarial-control weakness, tune the GAN-loss weight (and discriminator warm-up) for WTG over a modest grid and report the best point, or soften the claim about the intrinsic non-triviality of pixel-space adversarial training.
• Un-comment and ship the LZMA companion. In response to the LZMA weakness, re-enable compress_binary_array in bsq_era5_cra5.py, ship the July-2023 pixel artifact (or a downloadable pointer), and include the accounting script that reproduces 0.117 TB.

Suggested

• Add a compression benchmark against foundation-model tokenisers. In response to the unsupported 'WLA leads' claim, benchmark against at least one foundation-model tokenizer (e.g., Aurora) on compression ratio and BPSP, or soften the wording.
• Justify the $2^{128}$ codebook on the rate-distortion axis. In response to the codebook-choice weakness, add compression-ratio numbers to the ablation and pick the operating point that dominates on the Pareto frontier, or explicitly state why maximum reconstruction is preferred despite the compression cost.
• Broaden the generalisation figure. In response to the single-variable Fig. 5 weakness, add at least one dynamical variable (e.g., $u$ or $w$), report per-level confidence bands over dates, and split HRES OOD from unseen-level performance so the two effects can be read off independently.
• Ship configs and seeds for the Fig. 8 ablation. In response to the Fig. 8 reproducibility gap, add a sweep.sh and per-cell configs for the 15 (codebook × level) grid, seeds per point, and a continuous_baseline.py implementing the bitrate-matched float32 variant used in Fig. 8(b).
• Release the HRES pipeline. In response to the HRES-unimplemented weakness, ship the downsampling/normalisation script and specify the HRES cycles/dates used so Fig. 5's OOD half can be checked.
• Vendor LPIPS/StyleGAN and release the WT trainer. In response to the GAN-ablation reproducibility weakness, either vendor the LPIPS weights (with license) and pin the exact StyleGAN commits, and release the WT trainer, or substitute a fully bundled perceptual/adversarial term.
• Quantify extreme-value error in Section A.7. In response to the surface tension between 'systematic error at extremes' and 'preserves extremes', report MAE (or MAPE) inside the >99th-percentile mask vs. everywhere for each of the six variables shown.
• Add a distributional metric for precipitation. In response to the single-metric superiority claim, complement RQE with (e.g.) CRPS or a POD/FAR / FSS sweep so 'characterises the heavy-tailed distribution' is supported.
• Fix baseline attributions in Section 4.2 and Table 1. Correct 'VAGAN' → 'VQGAN' and align the citations for VQVAE/VQGAN so Table 1 and the prose agree on whether numbers come from the original authors' code (fine-tuned) or from Mirowski et al. 2024.
• Terminology: 'fields' vs 'variables'. In response to the '164 variables' weakness, use 'fields' throughout to match Section 2's usage.
• Renumber appendix figures. Fix the 'Figure S1' collision between the precipitation RQE figure and the visualisation series (S3–S8), and rewrite Figure S2's caption to describe the WTG ablation.
• Downgrade 'superior' in Section A.7. In response to the visualisation-only comparative claim, replace 'superior reconstruction quality' with 'high reconstruction quality'; the comparative ordering belongs to Table 1.
• Soften 'does not affect' to 'has minimal impact'. In response to the internal inconsistency between Section 1 and Section 3.1, harmonise the wording so the mechanistic guarantee is not overstated.
• Clarify PVUM's weight-generation mechanism. In response to the specification gap, state explicitly whether each of the $C_1$ metadata tokens is projected to a length-$C_2$ row of $W$ by a shared linear layer, or whether a single projection produces the entire matrix.