Samuel Green: TRACE: Transformer Reanalysis Atmospheric Compression Engine for MERRA2
Introduction and scientific context
MERRA-2 reanalysis provides hourly global atmospheric data from 1980 to present across dozens of variable collections. Multi-collection research projects on Gadi routinely accumulate 60+ TB of reanalysis output, placing severe pressure on project storage quotas. Existing compression methods, int16 packing, ZFP, and zstd, yield only 2-5x reduction because MERRA-2 files are already encoded near their classical entropy ceiling.
Methodology
TRACE applies neural image compression (NIC) to atmospheric reanalysis, adapting the scale hyperprior VAEformer architecture for MERRA-2’s native 361x576 global grid. The model pairs a Swin Transformer encoder-decoder with a learned scale hyperprior entropy model for arithmetic coding, exploiting spatial and cross-channel correlations that classical codecs cannot capture. Per-variable importance weighting in the rate-distortion loss prioritises scientifically critical fields such as surface temperature. Training uses distributed data-parallel processing with BF16 mixed precision on GPUs, built on the CompressAI library.
Results
Applied to the 36-variable radiation product (1980–2022), TRACE currently achieves approximately 700x compression, encoding hourly global fields to ~100 KB versus a 15 MB int16 original. Surface temperature reconstructs at under 10% of its climatological standard deviation. 15 of 36 radiation variables meet an RMSE < 15% of standard deviation quality threshold, with cloud fractions and optical depths remaining the most challenging variables due to their high spatial variability. Training is ongoing and quality continues to improve.
Audience
Climate data infrastructure researchers, ML researchers working with Earth system data, and large-scale reanalysis users facing storage constraints at NCI.
Keywords
Neural compression, atmospheric reanalysis, Gadi, data infrastructure