PaperlayerView exact revision on arXiv

Source-linked AI summary

Climate-Invariant Machine Learning

Tom Beucler, Pierre Gentine, Janni Yuval, Ankitesh Gupta, Liran Peng, Jerry Lin, Sungduk Yu, Stephan Rasp, Fiaz Ahmed, Paul A. O'Gorman, J. David Neelin, Nicholas J. Lutsko, Michael Pritchard

arXiv:2112.08440v5cs.LGphysics.ao-phphysics.comp-ph

TL;DR

Machine-learning models can struggle to generalize to climate conditions beyond their training data, while future climates cannot be sampled for training. This paper uses physically informed transformations to make learned mappings approximately climate-invariant and reports improved generalization across climate and configuration changes.

  • Problem

    ML models often face out-of-distribution climate conditions, while future climates cannot be observed for training, motivating accurate prediction beyond the training range.

    Source: §1.1 · pp. 2–3

  • Method

    The framework transforms input and output variables so their distributions shift minimally across climates, making the learned mapping approximately climate-invariant.

    Source: §1.2 · p. 3

  • Results

    Climate-invariant neural networks with transformed inputs learn mappings robust to temperature and configuration changes and show superior generalization skill.

    Source: §5 · p. 18

  • Takeaways & Limitations

    The results suggest the possibility of neural-network-powered subgrid closures for Earth-like settings under substantially different climate conditions.

    Source: §5 · p. 18

    The proposed workflow assumes the target climate will not be used to retrain the ML algorithm, for example when its data are insufficient or less reliable.

    Source: §6 · p. 18

Abstract

from arXiv · show

Projecting climate change is a generalization problem: we extrapolate the recent past using physical models across past, present, and future climates. Current climate models require representations of processes that occur at scales smaller than model grid size, which have been the main source of model projection uncertainty. Recent machine learning (ML) algorithms hold promise to improve such process representations, but tend to extrapolate poorly to climate regimes they were not trained on. To get the best of the physical and statistical worlds, we propose a new framework - termed "climate-invariant" ML - incorporating knowledge of climate processes into ML algorithms, and show that it can maintain high offline accuracy across a wide range of climate conditions and configurations in three distinct atmospheric models. Our results suggest that explicitly incorporating physical knowledge into data-driven models of Earth system processes can improve their consistency, data efficiency, and generalizability across climate regimes.

1 Introduction

Machine-learning models can make unreliable out-of-distribution predictions, including under climate conditions beyond their training range. The paper asks how physical knowledge can make those predictions more accurate and proposes transforming inputs and outputs to make their mapping approximately climate-invariant.

1.1 Background

  • ML algorithms extrapolate according to implicit assumptions, and neural networks have no general reason to perform well far outside their training data.The paper defines extrapolation as prediction outside the training data range.Source: §1.1 · p. 2
  • Prior climate studies report both poor and successful generalization, with successful cases linked to inputs whose distributions change little or to particular layer choices.Some studies consequently recommend training across multiple climates.Source: §1.1 · p. 2
  • The paper focuses on physically informed methods that need no samples from the target climate and can complement domain adaptation and regularization methods.The motivation includes historical-only training because future climates cannot be observed.Source: §1.1 · pp. 2–3

1.2 Problem Definition

  • The central question is how to use physical knowledge to make accurate predictions when standard-variable climate conditions lie far outside training data.The proposed framework transforms input and output distributions to shift minimally across climates, making the learned mapping approximately climate-invariant.Source: §1.1 · p. 3, §1.2 · p. 3
  • The application is subgrid atmospheric thermodynamics, with inputs describing large-scale climate state and outputs grouping large-scale thermodynamic tendencies.The mapping assumes locality in space and time and determinism; it does not predict cloud-water changes and excludes greenhouse gases other than water vapor from inputs.Source: §1.2 · p. 3

2 Data

The study tests climate-invariant ML across three storm-resolving atmospheric models with distinct configurations and temperature conditions. It separates training, validation, and test data and evaluates both climate and configuration shifts.

2 Data

  • Three model setups—SPCAM3 aquaplanet, SPCESM2 Earth-like, and quasi-global hypohydrostatic SAM—provide tests across distinct configurations.SPCAM3 and SPCESM2 assume scale separation; SAM does not and is coarse-grained for ML parameterization.Source: §2 · pp. 3, 5

2.1 Super-Parameterized Aquaplanet Simulations

  • SPCAM3 uses two-year aquaplanet simulations at reference, 4 K colder, and 4 K warmer surface temperatures, extending the generalization gap to 8 K between cold and warm runs.The simulations use fixed boreal-summer insolation with a full diurnal cycle.Source: §2.1 · p. 5, §2 · p. 5

2.2 Super-Parameterized Earth-like Simulations

  • SPCESM2 supplies three two-year Earth-like simulations with realistic surface boundary conditions, seasonality, and aerosol conditions representative of 2000.The setup includes a land surface model and 32 SAM columns per atmospheric grid cell to represent deep convection.Source: §2.2 · p. 5

2.3 Quasi-Global Aquaplanet Hypohydrostatic Simulations

  • Hypohydrostatic SAM tests configurations without prescribed scale separation, but its idealized setup provides only the cold-to-reference generalization experiment.The SAM mapping also differs in inputs, outputs, and vertical coordinate; its warm simulation is unavailable.Source: §2.3 · pp. 5, 7

2.4 Normalization and Training, Validation, and Test Split

  • Training and validation use cold simulations for climate tests or SPCAM3 for configuration tests, with separate test sets and non-contiguous three-month periods to reduce temporal correlation.The climate tests evaluate on warmer simulations; configuration tests evaluate SPCAM3-trained models on Earth-like SPCESM2.Source: §2.4 · p. 7

3 Theory

Climate-invariant mappings are made approximately constant across climates by transforming input and output variables whose distributions otherwise shift. The proposed transformations use physical relationships to reduce those shifts, and all three input transformations improve generalization.

3 Theory

  • A climate-invariant mapping remains approximately constant across climates, achieved with physically informed transformations that make input and output distributions vary little.Unlike nondimensionalization, these transformations may yield variables with physical units.Source: §3 · p. 7
  • Relative humidity normalizes specific humidity by its saturation value, keeping most of its distribution within [0, 1] with little change as climate warms.The transformation also helps represent subgrid effects of dry-air entrainment.Source: §3 · p. 8
  • A moist-plume buoyancy transformation addresses the quasi-linear temperature shift while preserving approximate invariance of tropopause temperatures with warming.It links vertical temperature structure and near-surface humidity to a measure correlated with deep convective activity.Source: §3 · p. 8
  • Scaling latent heat flux by near-surface saturation deficit yields a variable approximately proportional to near-surface wind magnitude and density, whose distributions vary less with warming.The scaling is less effective over land, where evapotranspiration changes do not follow Clausius-Clapeyron scaling.Source: §3 · p. 8
  • The relative-humidity, buoyancy, and latent-heat-flux transformations each lead to statistically significant improvements in ML generalization.Source: §3 · p. 8

4 Results

Across the experiments, physically informed input transformations reduce climate-shift generalization errors, and climate-invariant models also transfer across configurations and multi-climate training data. SHAP analysis suggests their mappings are more spatially local.

4.1 Benefits of Incremental Input Transformations

  • Transforming q to RH, T to plume buoyancy, and LHF to LHF∆q brings warm-climate MSE within ≈25% of a raw-data NN trained in the warm climate.The final LHF transformation further reduces generalization MSE except for shortwave heating in the warm climate.Source: Figure 4 · p. 10, §4.1 · p. 11

4.2 Learning across Climates and Configurations

  • CI NNs continue learning warm-aquaplanet subgrid thermodynamics from cold training and transfer across configurations; dropout and batch normalization help difficult transfers.CI transformations also improve the MLR baseline’s generalization, though less dramatically.Source: §4.2 · pp. 11–12

4.3 Global Performance after Training

  • CI NNs with dropout and batch normalization often show the highest generalizability, and input transformations improve standard NN regularization.Raw-data MLRs generalize better than raw-data NNs, while fitting their training sets less well.Source: §4.3 · p. 12, Figure 5 · p. 13, §4.4 · p. 13

4.4 Understanding Climate Invariance

  • SHAP mappings suggest CI networks are more vertically local: plume buoyancy and relative humidity effects concentrate near the output level, unlike raw-data mappings.The raw-data network shows broad tropospheric dependencies, including moisture effects that may not be causal.Source: Figure 7 · p. 15, §4.4 · p. 16

4.5 Advantages of Climate Invariance with Multi-Climate Data

  • CI NNs outperform RD NNs offline in ≈95% of cases with data from both climates, and show less sensitivity to the training-data partition.The reported benefits include reduced variability in test errors.Source: Figure 8 · p. 17, §4.5 · p. 18

5 Discussion

The discussion reports that climate-invariant neural networks generalize across temperature and configuration changes, with attribution maps showing more spatially local mappings. The authors connect this approach to domain-informed feature engineering and describe possible implications for subgrid thermodynamics.

5 Discussion

  • Physically transformed neural networks generalize better as more inputs are transformed, and climate-invariant networks are robust to temperature and configuration changes across three atmospheric models.They show superior generalization skill almost everywhere, including when training data span multiple climates.Source: §5 · p. 18
  • Attribution maps indicate that climate-invariant networks learn more spatially local mappings that facilitate generalization across climates and configurations.The authors present locality alongside control of feature distributions as a characteristic of these mappings.Source: §5 · p. 18
  • Incorporating climate knowledge to extract more informative predictors is an example of domain-informed feature engineering.The authors also say this aids interpretability of the mapping.Source: §5 · p. 18
  • For subgrid thermodynamics, the results suggest the possibility of neural-network-powered closures that work in Earth-like settings under vastly different climate conditions.The authors also suggest that more local climate-invariant representations could support new analytic representations of convection.Source: §5 · p. 18

6 Materials and Methods

The workflow seeks physically informed transformations that make input and output distributions more similar across climates, then tests whether those transformations improve generalization. It assumes target-climate retraining is unavailable or unwanted, and uses progressive transformations because effects can interact.

6 Materials and Methods

  • The workflow assumes the model will not be retrained in the target climate, for example because target data are insufficient or less reliable, or to seek relations that hold across wider climates.The workflow uses physical transformations rather than finding them by training in the target climate.Source: §6 · p. 18
  • Propose transformations from robust physical or statistical relations, then check whether transformed-variable PDFs approximately match between training and target climates.Joint PDF matching is ideal; when data or practicality constrain this, matching univariate PDFs is necessary but not sufficient.Source: §6 · pp. 18–19
  • Normalize PDFs and their support variables to [0, 1] using the same factors across climates, then compare distributions with a PDF distance metric.The study uses the symmetric Hellinger distance, bounded from 0 for identical PDFs to 100% for non-overlapping PDFs.Source: §6 · p. 19
  • After transformed PDFs match, compare inexpensive raw-data and climate-invariant baseline models on the target climate; reject transformations when the transformed baseline performs worse.The baseline model should be tailored to the desired final model; the study used multiple linear regression baselines for climate-invariant neural networks.Source: §6 · p. 19
  • If the baseline improves, train raw-data and climate-invariant versions of the desired model on the reference climate and retain the transformation only if target-climate validation favors it.The workflow can be repeated for additional input or output variables.Source: §6 · p. 19
  • Transform inputs progressively because transformations interact: the study ordered them by importance as specific humidity, temperature, then surface energy fluxes.Testing temperature transformation independently initially suggested a negative effect, later invalidated by jointly transforming specific humidity and temperature.Source: §6 · p. 19

8 Supplementary materials

The supplementary materials are listed as supplementary text, figures, and tables.

8 Supplementary materials

  • The supplementary materials include supplementary text sections SM A through SM D.The supplied list also identifies Figures S1 to S15 and Tables S1 to S3.Source: §8 · p. 20
  • The supplementary materials list Figures S1 to S15.These accompany supplementary text sections SM A through SM D and Tables S1 to S3.Source: §8 · p. 20
  • The supplementary materials list Tables S1 to S3.These accompany supplementary text sections SM A through SM D and Figures S1 to S15.Source: §8 · p. 20

Supplementary Materials for Climate-Invariant Machine Learning

The supplementary materials describe reproducibility resources, transformation derivations, implementation guidance, and additional results for climate-invariant machine learning.

Supplementary Materials for Climate-Invariant Machine Learning

  • The supplement covers code and data availability, characteristics of the emulated mapping, and derivations of input and output transformations.It also discusses possible vertical-coordinate transformations.Source: Supplementary Materials for… · p. 27
  • A guide to finding new climate-invariant transformations and practical implementation details are included in the supplement.These are covered in SM B5 and SM C, respectively.Source: Supplementary Materials for… · p. 27
  • Supplementary results include distribution distances, learning curves, near-surface generalization skill, and methods for comparing raw-data and climate-invariant mappings.The comparison methods address cold (-4K) and warm (+4K) climates.Source: Supplementary Materials for… · p. 27

A. Code and Data Availability

The appendix documents code and simulation resources, describes the emulated mapping, and previews physically based input transformations evaluated for climate invariance.

A. Code and Data Availability

  • The code repository, reproducible training example, accompanying data, and simulation rerun resources are identified for SPCAM3, SPCESM2, and SAM.The appendix notes that full simulation outputs are several terabytes and available upon request.Source: A. Code and Data Availability · pp. 27–28

B1. Mapping

  • Inputs encode the large-scale climate state, while outputs group subgrid thermodynamic tendencies expressed as energy fluxes.For SPCAM3 and SPCESM2, x has length 2Np + 4 and y has length 4Np; for SAM, the lengths are 2Np + 1 and 2Np.Source: B1. Mapping · pp. 28–29

B2. Physically-Based Input Transformations

  • The input-transformation appendix compares three options by their univariate PDFs, then examines relative humidity, plume buoyancy, and other inputs’ distribution shifts.It identifies the bottom option as the most successful and the middle option as second best.Source: B2. Physically-Based Input… · p. 29

B2a. Second Best Input Transformations

The appendix evaluates alternative input transformations and idealized output quantile mapping, reporting which transformations show climate invariance and where their assumptions or performance are limited.

B2a. Second Best Input Transformations

  • Relative humidity yields a climate-invariant specific-humidity PDF, with distributions mostly overlapping across the three climates.Saturation deficit and temperature-minus-near-surface-temperature do not yield climate invariance; plume buoyancy yields approximate invariance.Source: B2a. Second Best Input Transformations · p. 29
  • The LHF transformation using near-surface saturation deficit improves generalization over the near-surface-specific-humidity alternative despite appearing slightly worse by PDF distances.The latter can fail in very dry atmospheres when latent heat flux is negative or near-surface humidity is very small.Source: B2a. Second Best Input Transformations · p. 30

B2b. Relative Humidity

  • Relative humidity is the selected specific-humidity transformation, defined as water-vapor partial pressure divided by its saturation value.The saturation-pressure expression blends liquid- and ice-phase values across the specified temperature interval.Source: B2b. Relative Humidity · p. 30

B2c. Plume Buoyancy

  • Plume buoyancy estimates a parcel’s buoyancy from environmental humidity, temperature, and pressure using a moist-static-energy-conserving plume.The derivation assumes a non-entraining, hydrostatic plume lifted from near the surface and a saturated parcel.Source: B2c. Plume Buoyancy · pp. 31–32

B2d. Sensible Heat Fluxes and Surface Pressure

  • Sensible heat-flux PDFs change little with warming, while the surface-pressure PDF’s left tail expands slightly as the most extreme low-pressure systems intensify.The authors hypothesize that these shifts are small enough not to require a dedicated transformation.Source: B2d. Sensible Heat Fluxes and Surface… · p. 32

B3a. Theory

  • Output transformations using marginal CDFs improve generalization only marginally in the tested idealized setting.The approach assumes outputs are more likely invariant in quantile than physical space and tests quantile mapping before and after training.Source: B3a. Theory · pp. 32–33

B3b. Results

Post-training output quantile mapping slightly improved warm-climate generalization, whereas training directly in probability space performed poorly. Vertical-coordinate transformations showed no visible benefit in this case.

B3b. Results

  • R2 rose from 0.58 to 0.62 for the generalization (+4K) set when outputs were transformed after training.Transforming outputs before training produced equally bad results on training and generalization sets; the authors say further investigation is needed to assess probability-space training.Source: B3b. Results · p. 33

B4. Spatial Coordinate Transformation

  • Possible vertical-coordinate choices include temperature, saturation specific humidity, geopotential height or altitude, and levels tied to atmospheric features such as the tropopause.The proposed coordinate is the hybrid pressure coordinate p.Source: B4. Spatial Coordinate Transformation · pp. 33–34
  • Coordinate transformations of the vertical pressure coordinate showed no visible benefit in this case, possibly because input transformations already address some upward shift of convective activity.The authors note this as a possible explanation, not a demonstrated cause.Source: B4. Spatial Coordinate Transformation · p. 34

C. Implementation

The workflow implements physical transformations in data pipelines and uses MSE-based training and benchmarking. The supplementary implementation specifies the regression and neural-network designs, evaluation metrics, and learning-curve procedure.

C1. Overall Workflow

  • The workflow compares three implementation approaches: transforming datasets before training, adding model layers, or customizing a data generator to transform batches.Pre-transforming datasets uses more disk space; custom layers slow training, while data-generator transformations can run in parallel on the CPU.Source: C1. Overall Workflow · p. 34

C2. Multiple Linear Regressions

  • The MLRs use TensorFlow 2.0 and Adam, train for 20 epochs, and retain weights and biases with the lowest validation loss.The climate-invariant MLR has a weight matrix of size 4Np × (2Np + 4) and a bias vector of length 4Np.Source: C2. Multiple Linear Regressions · pp. 34–35

C3. Neural Network Design

  • SPCAM neural networks use seven 128-neuron layers, while SAM networks use five 128-neuron layers and a cyclic learning-rate schedule.SPCAM networks train for 20 epochs; SAM networks train for 10. Some SPCAM networks also use dropout and batch normalization, without tuning those settings on the generalization test set.Source: C3. Neural Network Design · p. 35

C4. Benchmarking

  • MSE and R2 benchmark prediction skill, with averaging chosen according to whether scores are needed by output, location, or as a single value.The workflow also uses MSE curves on different-temperature and different-geography datasets to track generalization during training.Source: C4. Benchmarking · p. 35
  • Generalization-learning-curve callbacks evaluate approximately 100M samples per epoch and are best avoided when seeking performance alone.They return validation loss and losses on the two generalization sets as a function of training epochs.Source: C4. Benchmarking · p. 36

D. Supplementary Results

The supplementary material uses Jensen-Shannon distance to compare probability distributions across climates, emphasizing symmetry and differences in distribution tails.

D1. Jensen-Shannon Distance between PDFs across Climates

D2. Learning across Climates and Geographies

Climate-invariant models learn across climate and geography more successfully than raw-data models in the reported tests, with improvements extending to warm-Tropics skill. Supplementary diagnostics compare distribution shifts and model sensitivities.

D2. Learning across Climates and Geographies

  • The climate-invariant NN’s warm-aquaplanet learning curve mostly decreases when trained in the cold aquaplanet, while the raw-data NN’s large generalization errors worsen during training.This indicates that the climate-invariant NN continues learning about warm-aquaplanet subgrid thermodynamics from cold-aquaplanet training.Source: D2. Learning across Climates and… · p. 36
  • Climate-invariant NNs improve transfer between aquaplanet and Earth-like simulations, while transformations improve MLR generalization less dramatically.The smaller MLR improvement is linked to fewer free parameters and lower representation power; dropout helps difficult cases but fits training data less well.Source: D2. Learning across Climates and… · pp. 36–37

D3. Geographic Skill

  • Climate-invariant mappings alleviate raw-data NNs’ poor warm-Tropics generalization across the atmosphere, without reducing skill in warm mid-latitudes and poles.The buoyancy-temperature transformation helps in the Tropics without overly compromising other latitudes; a slight high-latitude skill compromise is especially apparent for SAM and is partly traced to ice-sedimentation challenges.Source: D3. Geographic Skill · p. 37
  • Near-surface subgrid-heating results show the same pattern: raw-data NNs generalize poorly, while climate-invariant NNs mostly solve that error.Training-climate skill is highest over land for all NNs, where most variability comes from the diurnal cycle.Source: D3. Geographic Skill · p. 37

D4. Visualizing Climate-Invariant Mappings

  • MLR weights and NN Jacobians suggest non-local vertical mappings, while SHAP indicates subgrid heating and moistening can be predicted mostly from local plume-buoyancy information.Linear sensitivity to lower-tropospheric plume buoyancy is high, but that does not imply the prediction uses broad vertical information.Source: D4. Visualizing Climate-Invariant… · p. 37
  • Supplementary PDF comparisons cover key variables and transformations across models and climates; for temperature and latent heat flux, the best options do not reduce distribution distance more than the second-best options.The distance tables report values away from the cold simulation for key input PDFs and their transformations.Source: Table S1 · p. 38, Table S3 · p. 38, Figure S2 · p. 40

Related papers

Loading 2112.08440v5…