Post-Training Quantization of Autoregressive Weather Models

1 Introduction

• NWP computation costs challenge scaling toward higher spatio-temporal resolution, including efforts toward global forecasting at 1 km resolution.The passage contrasts CMIP studies at 100–200 km with HRRR at 3 km over CONUS.Source: §1 · p. 1
• Computational loads in emulator training and testing constrain high-resolution modeling and out-of-the-box use in resource-constrained environments.Reducing memory footprint, inference time, and power consumption is presented as a way to improve feasibility of scaling emulators.Source: §1 · p. 2
• Lower-bit representations can reduce memory and accelerate matrix multiplications, while quantization alters the precision of model inputs, weights, and activations.The study applies PTQ to two DL architectures to investigate model optimization for weather forecasting.Source: §1 · p. 2
• PTQ is implemented across weight and activation precisions in pretrained DLWP U-Net and FCN vision-transformer emulators to optimize matrix multiplication for deterministic forecasting.The study also characterizes quantization error during autoregressive inference, where it compounds across recursive steps in a nonlinear system.Source: §1 · pp. 2–3

2.1 PTQ Workflow

• Earth2Studio loads pretrained DLWP and FCN checkpoints, while ModelOpt configures weight and activation quantization and the resulting forecasts are evaluated against reference truth, full precision, and climatology.ERA5 supplies the initial conditions for emulator inference.Source: §2.1 · p. 3, Figure 1 · p. 4

2.2 Model and PTQ Configurations

• The study tests DLWP’s cubed-sphere U-Net and FCN’s Fourier-mixing vision transformer using FP32, INT8, INT4, and INT2 configurations, including SmoothQuant and AWQ.Activations remain full precision in A32 configurations; input, output, and normalization layers are also kept full precision.Source: §2.2 · p. 3

2.3 Evaluation setup

• PTQ is calibrated on a 48-hour window of eight inference steps and evaluated over ten steps, or 60 hours, on four seasonal dates using RMSE and ACC.The setup uses ERA5 initial conditions, WeatherBench 2 ERA5 climatology for ACC, and tests on NVIDIA V100 and A100 GPUs.Source: §2.3 · pp. 3–4, Figure 1 · p. 4

3.1 Skill Scores and Forecast Quality

• FCN's W8A32 and SmoothQuant configurations perform best over 60 hours, with ACC above 0.65 despite variation across initial conditions.DLWP RMSE shows a similar trend to FP32 in most PTQ experiments; its temperature and total-column-water-vapor ACC trends closely follow FP32 in the best configurations.Source: §3.1 · pp. 4–5

3.2 Spatial Maps of Forecast Variables

• AWQ and SmoothQuant preserve DLWP spatial quality relative to FP32, while reducing precision introduces coarse features and biases.For FCN, W8A32 and INT8 SmoothQuant retain spatial skill; INT2 maps show periodic patterns reflecting Fourier-transform-based embedding mixing.Source: §3.2 · p. 5
• DLWP's intermediate-precision configurations with special quantization algorithms preserve spatial features and forecast skill within the 60-hour lead time.Quantized temperature fields retain spatial trends relative to FP32, though reducing precision leads to less accurate forecast values and the fields contain hot and cold biases.Source: §3.2 · p. 5

3.3 Hardware Metrics

• Mean power does not change significantly across PTQ configurations, and V100 consumes less power than A100 across configurations.Inference time is hardware-dependent; PTQ adds time relative to FP32 in most cases, while DLWP computations are 2–4 times faster than FCN.Source: §3.3 · p. 6

• Pre-trained DLWP and FCN emulators produce qualitatively meaningful forecasts over short-range horizons under PTQ, despite expected degradation in autoregressive multi-step inference.The evaluation uses RMSE and ACC skill scores to assess the configurations.Source: §4 · p. 6
• Weight-bit quantization plays a more dominant role in forecast-skill degradation than changes in activation-bit quantization.Source: §4 · p. 6
• SmoothQuant and AWQ perform better than naive absolute-maximum-based PTQ calibration.Source: §4 · p. 7
• The experiments provide a first benchmark of PTQ for autoregressive weather emulators and point to future QAT and idealized-circulation studies.The paper proposes assessing physical interpretation and consistency with idealized general-circulation test cases.Source: §4 · p. 7

• For DLWP, the best PTQ configurations have geopotential RMSE below 500 m2s−2 and temperature RMSE below 2 K relative to FP32.SmoothQuant and AWQ outperform absolute-max calibration at lower weight-bit representations across variables.Source: §A.1 · p. 10
• FCN RMSE varies substantially for z500, t850, t2m, and u10m, while z1000, v10m, and tcwv show no significant variation across initial conditions.W2A32 has non-monotonic wind-speed RMSE evolution and does not provide a qualitatively meaningful forecast.Source: §A.1 · p. 10

• Forecast quality degrades as weight and activation bit representation decreases, and output fields are smooth relative to ERA5 truth.Source: §A.2 · p. 10
• Figures 10 and 11 track RMSE error relative to FP32 over forecast lead time for DLWP and FCN PTQ configurations.The figures characterize degradation for different weight- and activation-quantization cases.Source: Figure 10 · p. 15, Figure 11 · p. 16
• Total-column-water-vapor maps at 60 hours compare DLWP and FCN PTQ configurations with ERA5 truth and each model's FP32 reference.Source: Figure 12 · p. 17, Figure 13 · p. 18