pnasresearcharticle Tom Beucler Predicting climate change is inherently a
generalization problem: we extrapolate observations from the recent past using
physical models we believe hold across past, present, and future climates.
Current climate models are based on fundamental physical equations but require
representations of processes that occur at scales smaller than the model grid
size. Recent machine learning (ML) algorithms hold promise to improve such
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 — for incorporating
physical knowledge of climate processes into ML algorithms and show that they
can maintain high accuracy across a wide range of climates and geographies.
TB, MP, SR, and PG initially designed the research. TB, MP, JY, AG, and LP
performed the research. All authors contributed new analytic tools. TB, MP,
JY, AG, and FA analyzed data. All authors wrote the paper. The authors have no
conflicts of interest to declare. See SI A for code and data availability. 1To
whom correspondence should be addressed. E-mail<EMAIL_ADDRESS>
# Climate-Invariant Machine Learning

Tom Beucler  Department of Earth System Science, University of California,
Irvine, CA 92697, USA  Institute of Earth Surface Dynamics, University of
Lausanne, Lausanne, VD 1015, Switzerland  Michael Pritchard  Department of
Earth System Science, University of California, Irvine, CA 92697, USA  Janni
Yuval  Department of Earth, Atmospheric, and Planetary Sciences, Massachusetts
Institute of Technology, Cambridge, MA 02139, USA  Ankitesh Gupta  Department
of Earth System Science, University of California, Irvine, CA 92697, USA
Liran Peng  Department of Earth System Science, University of California,
Irvine, CA 92697, USA  Stephan Rasp  Climate AI, San Francisco, CA 94111, USA
Fiaz Ahmed  Department of Atmospheric and Oceanic Sciences, University of
California, Los Angeles, Los Angeles, CA 90095, USA  Paul A. O’Gorman
Department of Earth, Atmospheric, and Planetary Sciences, Massachusetts
Institute of Technology, Cambridge, MA 02139, USA  J. David Neelin  Department
of Atmospheric and Oceanic Sciences, University of California, Los Angeles,
Los Angeles, CA 90095, USA  Nicholas J. Lutsko  Scripps Institution of
Oceanography, University of California, San Diego, La Jolla, CA 92037, USA
Pierre Gentine  Department of Earth and Environmental Engineering, Columbia
University, New York, NY 10027, USA

###### Abstract

Data-driven algorithms, in particular neural networks, can emulate the effects
of unresolved processes in coarse-resolution climate models when trained on
high-resolution simulation data; however, they often make large generalization
errors when evaluated in conditions they were not trained on. Here, we propose
to physically rescale the inputs and outputs of machine learning algorithms to
help them generalize to unseen climates. Applied to offline parameterizations
of subgrid-scale thermodynamics in three distinct climate models, we show that
rescaled or “climate-invariant” neural networks make accurate predictions in
test climates that are 4K and 8K warmer than their training climates.
Additionally, “climate-invariant” neural nets facilitate generalization
between Aquaplanet and Earth-like simulations. Through visualization and
attribution methods, we show that compared to standard machine learning
models, “climate-invariant” algorithms learn more local and robust relations
between storm-scale convection, radiation, and their synoptic thermodynamic
environment. Overall, these results suggest that explicitly incorporating
physical knowledge into data-driven models of Earth system processes can
improve their consistency and ability to generalize across climate regimes.

###### keywords:

Machine Learning $|$ Climate Change $|$ Earth System Science $|$ Transfer
Learning $|$ Physics-Informed Machine Learning

###### doi:

www.pnas.org/cgi/doi/10.1073/pnas.XXXXXXXXXX

This manuscript was compiled on

Following its demonstrated success in computer vision and natural language
processing, machine learning (ML111here broadly describing algorithms that
learn a task from data without being explicitly programmed for that task) is
rapidly percolating through climate science (e.g., reviews by 1, 2, 3).
Examples include the emulation of radiative transfer (e.g., 4, 5, 6, 7),
gravity wave momentum fluxes (e.g., 8, 9, 10), and microphysical schemes
(e.g., 11, 12, 13), the bias correction of climate predictions (e.g., 14, 15),
the detection and classification of clouds and storms (e.g., 16, 17, 18, 19),
and the development of subgrid closures from high-resolution simulation data
(e.g., 20, 21, 22), which is the main application discussed in this
manuscript.

While ML algorithms are typically designed to maximize accuracy on their
training datasets, they make strong implicit assumptions when extrapolating
outside of their training sets. To give common examples, multiple linear
regressions (MLR) assume that the linear relationship that best describes the
training set is universally valid. When confronted with out-of-distribution
inputs, random forests (RF, 23) find the closest inputs in their training sets
and assign the corresponding outputs regardless of the value of these out-of-
distribution inputs. Neural networks (NN), which are powerful nonlinear
regression and classification tools, rely on their nonlinear activation
functions and fitted weights to extrapolate. Except in specific situations
(e.g., out-of-distribution samples in the close neighborhood of the training
set or described by the same nonlinear mapping as the training set), there is
no reason why NNs should generalize well far outside of their training sets.

In climate applications, this means that ML algorithms typically fail when
exposed to dynamic or thermodynamic conditions that differ substantially from
the range of conditions they were trained on, which remains under-reported in
the research literature. Exceptions include (24), who showed that an RF-based
moist convection scheme generalizes poorly in the tropics of a climate 6.5K
warmer than the training climate although it could generalize well to a colder
climate; and (25), who showed that an NN-based thermodynamic subgrid closure
generalizes well to climates 1-2K warmer than the training one, but makes
large errors as soon as the test climate is 4K warmer than the training one.
(26) confirmed that these generalization errors remain even when the NN
subgrid closure is modified to enforce conservation laws to within machine
precision. Interestingly, (27) found that convolutional NNs (28) trained to
classify thunderstorms in high-resolution model outputs generalize well to a
warmer climate representing the end of the twenty-first century following the
RCP8.5 emission scenario (29). While this may be because classifying
thunderstorms relies heavily on wind velocity inputs, whose distributions
change only slightly under an RCP8.5 warming scenario, the authors noted that
using two types of ML layers, namely batch normalization (BN, 30) followed by
dropout (DP, 31), was key to this successful generalization.

BN and DP are two examples of a larger set of methods used by the ML community
to help NNs generalize and avoid overfitting, broadly referred to as
“regularization” (32). Most empirical regularization methods (e.g., L1 and L2
regularizations) rely on the principle of parsimony, i.e. that simpler models,
accurately describing the training set with fewer fitted parameters, are
preferable to more complex models and generalize better to unseen conditions.
More systematic approaches to regularization have been developed in order to
use ML models in out-of-distribution situations that still require the same
inputs/outputs, referred to as domain adaptation (e.g., 33, 34, 35). However,
most domain adaptation approaches (e.g., active, semi-supervised, and
transductive learning) require at least a few samples in the generalization
domain.

Here, we purposely focus on domain adaptation methods that do not require
samples in the generalization domain for two reasons: (1) a long-term goal of
the climate science community is to train ML models that rely only on
observations and by definition we cannot observe the future climate; (2) even
if we have access to simulation data across climates, ML models that
generalize to climates they have not been trained on by construction are
likely to be more robust to other changes (e.g., a change in model geography)
as we will show in this manuscript. Motivated by this challenge, we ask:

How can we modify ML algorithms so that they make accurate predictions in
climate conditions that – in standard variables – lie far outside of their
training sets?

### Problem Definition

Our main contribution is to transform the emulated mapping within the ML
framework from a mapping that changes with climate, here referred to as
“brute-force”, to a mapping that approximately remains constant across
climates, here referred to as “climate-invariant”. Inspired by invariants in
physics and self-similarity in fluid mechanics, we change the emulated mapping
from “brute-force” to “climate-invariant” by rescaling the input and output
vectors so that their distributions do not shift across climates (see Fig 1).

Figure 1: By rescaling inputs and outputs to match their univariate PDFs
across climates, algorithms can learn a mapping that holds across climates. To
illustrate this, we show the PDFs of humidity-related inputs and outputs in
two different climates using blue and red lines.

We demonstrate the utility of this framework by adapting several ML closures
of subgrid-scale atmospheric thermodynamics (i.e., large-scale thermodynamic
tendencies resulting from subgrid convection, radiation, and turbulence) so
that they generalize better across climates. The motivation for this
application is two-fold. First, purely physically-based subgrid closures
remain one of the largest sources of error in Earth system models (36, 37,
38). While ML-based closures have emerged as a promising alternative to
traditional semi-empirical models (39), they lack robustness (40, 41) and, as
discussed earlier, usually fail to generalize across climates (24, 25, 26).
Second, atmospheric thermodynamic processes are directly affected by global
temperature changes, e.g. in response to anthropogenically-forced climate
change (42). Therefore, predicting subgrid thermodynamics in a warm climate
with a ML model trained in a cold climate leads to very apparent failure modes
(26) that we can transparently tackle.

In mathematical terms, our goal is to build a “climate-invariant” mapping
between the input vector $\boldsymbol{x}$ representing the large-scale
($\approx$100km) thermodynamic state:

$\boldsymbol{x}=\left[\begin{array}[]{cccccc}\boldsymbol{q\left(p\right)}&\boldsymbol{T\left(p\right)}&p_{s}&S_{0}&\mathrm{SHF}&\mathrm{LHF}\end{array}\right]^{T},$
(1)

and the output vector $\boldsymbol{y}$ representing the large-scale
thermodynamic tendencies due to explicitly-resolved convection and
parameterized radiative transfer at the $\sim$1km-scale, as well as
parameterized turbulent mixing:

$\boldsymbol{y}=\left[\begin{array}[]{cccc}L_{v}\boldsymbol{\dot{q}\left(p\right)}&c_{p}\boldsymbol{\dot{T}\left(p\right)}&c_{p}\boldsymbol{\mathrm{lw}\left(p\right)}&c_{p}\boldsymbol{\mathrm{sw}\left(p\right)}\end{array}\right]^{T},$
(2)

where $\boldsymbol{q\left(p\right)}$ is the vertical profile of specific
humidity in units kg/kg (written as a function of the background pressure
coordinate $p$ in units Pa), $\boldsymbol{T\left(p\right)}$ is the
temperature’s vertical profile in units K, $p_{s}$ is surface pressure in
units Pa, $S_{0}$ is solar insolation in units $\mathrm{W/m^{2}}$, SHF is
surface sensible heat flux in units $\mathrm{W/m^{2}}$, LHF is surface latent
heat flux in units $\mathrm{W/m^{2}}$, $L_{v}\approx 2.50\times\
10^{6}\mathrm{J\ kg^{-1}}$ is the latent heat of vaporization of water in
standard conditions, $\boldsymbol{\dot{q}\left(p\right)}$ is the subgrid
moistening vertical profile, $c_{p}\approx 1.00\times\ 10^{3}\mathrm{J\
kg^{-1}\ K^{-1}}$ is the specific heat of water at constant pressure in
standard atmospheric conditions, $\boldsymbol{\dot{T}\left(p\right)}$ is the
total subgrid heating vertical profile,
$\boldsymbol{\mathrm{lw}\left(p\right)}$ is the subgrid longwave radiative
heating vertical profile, and $\boldsymbol{\mathrm{sw}\left(p\right)}$ is the
subgrid shortwave radiative heating vertical profile. Following (43), all
components of the output vector $\boldsymbol{y}$ are mass-weighted and
vertically integrated throughout each vertical layer to yield energy flux
units ($\mathrm{W/m^{2}}$). Assuming vertical profiles have $N_{p}$ vertical
levels, $\boldsymbol{x}$ is of length $\left(2N_{p}+4\right)$ and
$\boldsymbol{y}$ is of length $4N_{p}$. The details of $\boldsymbol{x}$ and
$\boldsymbol{y}$ are somewhat different for the convection-permitting model
SAM (see “Data” Section).

We keep the overall structure of the mapping
$\boldsymbol{x}\mapsto\boldsymbol{y}$ fixed throughout the manuscript and aim
to predict $\boldsymbol{y}$ as accurately as possible in training and
generalization climates. Note that this mapping makes implicit assumptions
based on successful past attempts (25, 44), including locality in horizontal
space and time (outputs only depend on inputs in the same atmospheric column
at the same time step) and determinism (only one possible output vector for a
given input vector). For simplicity, we do not predict changes in cloud liquid
water and ice and exclude cloud water and greenhouse gases other than water
vapor from the input vector $\boldsymbol{x}$.

After introducing the climate simulations and training/validation/test split
(“Data” Section), we define the “climate invariant” mapping and physical
input/output rescalings (“Theory” Section), and demonstrate and explain their
ability to generalize (“Results” Section) before concluding. We refer the
reader to the Supplementary Information (SI) for data availability (SI A),
additional derivations and descriptions of the physical rescalings (SI B), the
implementation of our ML framework (SI C), and additional results (SI D).

## Data

To test the robustness of our framework across climate models, we use three
climate models and experimental set-ups : super-parameterized aquaplanet
simulations using the Super-Parameterized Community Atmosphere Model version
3.0 (SPCAM3), super-parameterized Earth-like simulations using the Super-
Parameterized Community Atmosphere Model version 5 (SPCAM5), and quasi-global
aquaplanet hypohydrostatic simulations using the the System for Atmospheric
Modeling version 6.3 (SAM). SPCAM assumes a strict scale separation between
the resolved and subgrid processes, making it an ideal testbed to
unambiguously machine learn local subgrid closures (45, 25). In contrast, SAM
does not assume scale separation as a global convection-permitting model,
which is more realistic but requires carefully coarse-graining its high-
resolution simulation output (46). For each climate model, we run three
simulations with three different prescribed surface temperature distributions:

  * •

(+0K) A reference simulation with a range of surface temperatures analogous to
the present climate

  * •

(-4K) A cold simulation with surface temperatures 4K cooler than the (+0K)
simulations

  * •

(+4K) A warm simulation with surface temperatures 4K warmer than the (+0K)
simulations

with the exception of SAM for which only the (-4K) and (+0K) simulations are
available. We summarize the characteristics of each simulation set in Tab S1
and refer the reader to Fig 2 for a visualization of surface temperatures in
each model (Fig 2a) and snapshots of mid-tropospheric subgrid heating (Fig
2b).

Figure 2: (a) Prescribed surface temperature (in K) for (left) the aquaplanet
SPCAM3 model and (right) the hypohydrostatic SAM model. (center) Annual-mean,
near-surface air temperatures in the Earth-like SPCAM5 model. (b) Snapshots of
near-surface subgrid heating rate (in K per day) for the same three models.
For each model, we show the cold (-4K), reference (+0K), and warm (+4K)
simulations.

### Super-Parameterized Aquaplanet Simulations

We use two-year SPCAM3 (47) climate simulations in an aquaplanet
configuration222an “ocean world” with no land (48), with zonally-symmetric
surface temperatures fixed to a realistic meridionally-asymmetric profile
(49). The insolation is fixed to boreal summer conditions with a full diurnal
cycle. A two-dimensional convection-permitting model is embedded in each grid
cell of SPCAM3, namely 8 SAM atmospheric columns, a spatiotemporal resolution
of 4km$\times$30lev$\times$25s, and the default one-moment microphysical
scheme (50). The (+0K) SPCAM3 simulation was first presented in (45) and
subsequently used to train ML subgrid closures in (25, 41, 43). Inspired by
the generalization experiment of (24), the (+4K) simulation was introduced in
(25), and we ran the (-4K) simulation for the work presented here to
ambitiously increase the surface temperature generalization gap from 4K to 8K.

### Super-Parameterized Earth-like Simulations

We run three two-year SPCAM5 (51) climate simulations in an Earth-like
configuration, where the surface boundary conditions are realistic, including
a land surface model, seasonality, aerosol conditions representative of the
year 2000, and a zonally asymmetric annual climatology of sea surface
temperatures derived from the “HadOIB1” dataset (52). We use the Community
Earth Systems Model Version 2.1.3 to couple the Community Atmosphere Model
Version 4.0 with the Community Land Model, and similarly embed 32 SAM columns
in each atmospheric grid cell to explicitly represent deep convection. Our
(+0K) simulation is almost identical to the Earth-like simulation (44)
recently used to show the potential of ML for subgrid closures in Earth-like
conditions.

### Quasi-Global Aquaplanet Hypohydrostatic Simulations

While super-parameterization is well-adapted to statistically learning subgrid
closures thanks to its explicit scale separation, this scale separation comes
at the cost of distorted mesoscale systems and momentum fluxes (53).
Furthermore, most machine-learned subgrid closures are based on coarse-
graining high-resolution simulations (e.g., 54, 55). This motivates us to test
the “climate-invariant” framework in hypohydrostatic SAM simulations in which
the dynamics are not affected by a scale separation. Computational expense is
reduced through hypohydrostatic scaling which multiplies the vertical
acceleration in the equations of motion by a factor of 16 to increase the
horizontal scale of convection without overly affecting the larger-scale flow
(56, 57). While these simulations use idealized settings, such as aquaplanet
configurations, a quasi-global equatorial beta plane domain, and perpetual
equinox without a diurnal cycle, (58) showed that they produce tropical
rainfall intensity and cluster-area distributions that are close to satellite
observations. The prescribed surface temperature distribution in the control
simulation of (58) is designed to be close to zonal-mean observations (59),
and its maximum value is roughly 2K colder than that of the distribution used
for the 0K SPCAM3 simulation. To better match the SPCAM maxima of
distributions of upper-level temperatures and humidities, we choose to treat
this SAM control simulation as the (-4K) SAM simulation, and the warm
simulation of (58) as the (+0K) SAM simulation. We refer the reader interested
in the details of the simulations and the coarse-graining (here by a factor of
8) to (58) and the methods section of (46). We summarize key differences in
the mappings that are learned for SAM as compared to SPCAM below:

  * •

the input vector does not contain the specific humidity, surface pressure,
sensible heat flux, or the latent heat flux, but instead contains the total
water concentration and uses distance to the equator as a proxy for solar
insolation,

  * •

the output vector includes the subgrid total non-precipitating water tendency
instead of the subgrid specific humidity tendency,

  * •

the output vector does not contain the subgrid longwave and shortwave heating;
instead the radiative heating is included in the temperature tendency
$\boldsymbol{\dot{T}}$,

  * •

SAM uses a height-based vertical coordinate rather than a pressure-based one,
and

  * •

the (+4K) SAM simulation is not available and so the generalization experiment
is from the (-4K) to the (+0K) simulation.

### Training, Validation, and Test Split

To understand which solutions have the most promise for helping ML algorithms
generalize to unseen datasets, we purposely designed two difficult
generalization experiments: (1) Training and validating ML models on cold
simulations (-4K) and testing them on warm simulations (+4K); and (2) training
and validating ML models on Aquaplanet simulations (SPCAM3) and testing them
on Earth-like simulations (SPCAM5). Both generalization experiments expose ML
models to out-of-distribution inputs they have not been trained on. Following
best ML practices, we used the training set to optimize the ML model’s free
parameters, saved the free parameters that led to the best performance on the
validation set to avoid overfitting the training set, and evaluated the final
model on samples from a separate test set. We split each of the 9 simulations
into training/validation/test sets by using separate 3-month periods (reported
in Tab S1) to avoid high temporal correlations between
training/validation/test set samples (60).

## Theory

We formally define a “climate-invariant” mapping as a mapping that is
unchanged across climates. In practice, it is difficult, if not impossible, to
find mappings that are exactly invariant and we will use the terminology
“climate-invariant” for any mapping that remains approximately constant across
climates. To achieve climate invariance, we introduce physical rescalings,
defined as physically-informed functions that map the inputs/outputs to
different inputs/outputs whose distributions vary less across climates. We
deem the physical rescalings to be “climate-invariant” if they are so
successful that the distributions of the inputs/outputs vary minimally across
climates. Note that “climate-invariant” rescalings are distinct from
dimensional analysis in fluid mechanics, as dimensional analysis does not
necessarily affect distributions while “climate-invariant” rescalings may
yield variables that are not dimensionless.

After introducing a workflow for identifying climate-invariant rescalings, we
compare different possible input rescalings. We discuss output rescalings in
SI B2 and possible vertical coordinate rescalings in SI B3.

### Finding Climate-Invariant Rescalings

Fig 3 illustrates our proposed workflow for finding robust inputs/outputs
rescalings that transform the initial “brute-force” mapping into a “climate-
invariant” mapping when combined. Note that this workflow assumes that we
cannot or do not want to retrain ML algorithms in the target climate, which
excludes automatically finding a rescaling by training a model. This
limitation could arise because the data in the target climate are insufficient
or less reliable, or because we seek to uncover new physical relations that
hold across an even wider range of climates.

Figure 3: Proposed five-step workflow to find climate-invariant rescalings
that help Machine Learning (ML) models generalize from a reference (ref.)
climate to a target one, using (top) a baseline Multiple Linear Regression
(MLR) as an initial guide.

The first step is to propose a physical rescaling to implement. We can do this
through knowledge of robust physical or statistical relations that link and/or
preserve distributions (e.g., state equations, self-similarities, conservation
laws, accurate empirical relations, etc). These relations help derive
invariants (e.g., 61) under a change in thermodynamic conditions. Before
taking the time to implement this rescaling in the ML workflow, we can verify
that the PDFs of the rescaled inputs/outputs match in the training and target
climates. Ideally, the joint PDFs of the rescaled inputs/outputs would
perfectly match. In practice, because it is easier to rescale one variable at
a time and the data are often insufficient in the target climate, we can fall
back on the necessary (but not sufficient) condition that the univariate PDFs
of the rescaled inputs/outputs must match in the training and target climates.
Mathematically, this match can be quantified using PDF distance metrics.

An additional challenge is that the original and rescaled variables may have
different units and range, meaning that any nonlinear distance metric will
complicate the PDF comparison. To address this, we normalize the PDFs and
their support variables $X$ so that the PDFs’ domains strictly lie within
$\left[0,1\right]$. For a given variable, we use the same normalization
factors across climates:

$X_{\mathrm{norm}}\overset{\mathrm{def}}{=}\frac{X-\min_{\mathrm{cl}}X}{\max_{\mathrm{cl}}X-\min_{\mathrm{cl}}X},$
(3)
$\mathrm{PDF_{norm}}\overset{\mathrm{def}}{=}\frac{\mathrm{PDF}}{\int_{0}^{1}dX_{\mathrm{norm}}\times\mathrm{PDF}\left(X_{\mathrm{norm}}\right)},$
(4)

where $\mathrm{PDF_{norm}}$ is the rescaled PDF and $X_{\mathrm{norm}}$ is its
rescaled support, and $\max_{\mathrm{cl}}$ and $\min_{\mathrm{cl}}$
respectively refer to the maximum and minimum operators over the variables’
domains and across climates, i.e. over the (-4K), (+0K), and (+4K)
simulations.

Once the PDFs of each variable are normalized, we may pick any informative PDF
distance metric to quantify how PDFs match across climates. Here, we pick the
commonly-used Hellinger distance between two PDFs $p$ and $q$, formally
defined (62) as:

$\mathrm{\cal
H}\left(p,q\right)\overset{\mathrm{def}}{=}\sqrt{\int_{0}^{1}\frac{dx}{2}\left[\sqrt{p\left(x\right)}-\sqrt{q\left(x\right)}\right]^{2}}.$
(5)

This distance is symmetric (i.e., the arguments’ order does not affect the
outcome) and easy to interpret: $\mathrm{\cal H}\left(p,q\right)$ is bounded
by 0 (when $p=q$) and $100\%$ (when $p$ is zero whenever $q$ is positive and
vice-versa). In SI D1, we show that our results using the Hellinger distance
(see Tab 1) are consistent with those using the Jensen-Shannon distance (63),
a PDF distance metric giving large weights to the PDFs’ tails that tend to be
particularly problematic for generalization purposes.

Row | Input | SPCAM3 | SPCAM5 | SAM
---|---|---|---|---
1 | $q_{600\mathrm{hPa}}$ |  20.3, 35.1 |  17.1, 29.5 | 22.1
2 | $q_{\mathrm{deficit,}600\mathrm{hPa}}$ |  24.9, 36.5 |  18.1, 31.0 | 30.0
3 | $\mathrm{RH}_{600\mathrm{hPa}}$ |  3.6, 8.2 |  3.2, 5.3 | 4.3
4 | $T_{850\mathrm{hPa}}$ |  53.2, 64.3 |  25.7, 37.3 | 51.9
5 | $T_{\mathrm{from\ NS},850\mathrm{hPa}}$ |  5.1, 6.2 |  3.3, 6.4 | 10.5
6 | $B_{\mathrm{plume},850\mathrm{hPa}}$ |  9.4, 14.7 |  3.6, 7.6 | 5.8
7 | $T_{150\mathrm{hPa}}$ |  30.2, 33.5 |  31.6, 34.4 | 65.6
8 | $T_{\mathrm{from\ NS},150\mathrm{hPa}}$ |  38.0, 53.7 |  14.5, 28.7 | 51.0
9 | $B_{\mathrm{plume},150\mathrm{hPa}}$ |  35.1, 42.2 |  10.4, 20.9 | 21.1
10 | $\mathrm{LHF}$ |  8.6, 14.5 |  9.7, 12.5 |
11 | $\mathrm{LHF}_{q}$ |  4.7, 9.5 |  10.0, 10.7 |
12 | $\mathrm{LHF}_{\Delta q}$ |  6.3, 9.9 |  9.9, 14.0 |
Table 1: Hellinger distance (in $\%$) away from the (-4K) simulation for the
PDFs of
$\left(q_{600\mathrm{hPa}},T_{850\mathrm{hPa}},T_{150\mathrm{hPa}},\mathrm{LHF}\right)$
and their rescalings: (+0K) distance in gray and (+4K) distance in red.

Once the univariate PDFs of the physically rescaled variables match across
climates, the second step is to train two inexpensive or “baseline”
models333Because we here develop climate-invariant neural nets, we suggest
training MLR baselines, but note that the ML model used to define the baseline
should be tailored to the desired final ML model. on the reference climate to
quickly check whether the rescaling improves an ML model’s generalization
ability: (1) A “brute-force” model without the rescaling; and (2) a “climate-
invariant” model with the rescaling. If the rescaling does not improve the
baseline model’s generalization abilities (i.e., (2) performs worse than (1)
in the target climate), then the rescaling may not be appropriate.

If the rescaling improves the inexpensive baseline model’s performances, the
last step is to train the “brute-force” and “climate-invariant” versions of
the desired ML model (usually nonlinear) on the reference climate. If the
physical rescaling improves the desired ML model’s generalization abilities
(i.e., the climate-invariant model beats the brute-force model in the target
climate using the same performance metric), then we may keep the rescaling.
This workflow may be repeated for the ML model’s additional input/output
variables until the emulated mapping is as “climate-invariant” as possible.

Before applying this workflow to subgrid thermodynamics closures, we underline
one of its key challenges: Because some rescalings are much more impactful
than others, it is often not possible to develop each physical rescaling
independently. In our case, the specific humidity inputs vary the most across
climates, meaning that rescaling specific humidity affects the model’s
generalization abilities the most. As a result, initial experiments that
independently tested the effect of rescaling temperature suggested a negative
impact of temperature rescaling on generalization ability (not shown). This
initial result was later invalidated by experiments that jointly rescaled
specific humidity and temperature (see below). Following this, we adopt a
progressive input rescaling approach, where the most important inputs are
rescaled first: Specific humidity, then temperature, and finally surface
energy fluxes.

### Input Rescaling

Throughout the following section, we compare three options for each input,
whose univariate PDFs are depicted for all three climate models in Fig 4: No
rescaling (top), our most successful rescaling (bottom), and our second best
rescaling (middle).

Figure 4: Univariate PDFs of the (a) 600hPa specific humidity, (b) 850hPa
temperature, and (c) latent heat flux in the cold (blue), reference (gray),
and warm (red) simulations of each model (SPCAM3, SPCAM5, and SAM). For each
variable, we also show the PDFs of the two rescalings discussed in the
“Theory” section. From top to bottom, the variables are $q$ (g/kg),
$q_{\mathrm{deficit}}$ (g/kg), RH, T (K), $T_{\mathrm{from\ NS}}$ (K),
$B_{\mathrm{plume}}$ (m/s2), LHF (W/m2), $\mathrm{LHF}_{q}$ ($\mathrm{kg\
m^{-2}s^{-1}}$), and $\mathrm{LHF}_{\Delta q}$ ($\mathrm{kg\ m^{-2}s^{-1}}$).
For a given variable and rescaling, we use the same vertical logarithmic scale
across models.

#### Specific Humidity

Without any rescaling, the PDF of specific humidity $\boldsymbol{q}$ (Fig 4a,
top) extends through a considerably larger range as the climate warms (Tab 1,
row 1). This is because $\boldsymbol{q}$’s upper bound, namely the saturation
specific humidity $\boldsymbol{q_{\mathrm{sat}}\left(T,p\right)}$, increases
quasi-exponentially with temperature through the Clausius-Clapeyron relation
(e.g, 64, 65). Saturation deficit (Fig 4a, middle), defined as the amount by
which the water vapor concentration must be increased to achieve saturation
without changing the environmental temperature and pressure:

$\boldsymbol{\tilde{q}_{\mathrm{deficit}}\left(T,p\right)}\overset{\mathrm{def}}{=}\boldsymbol{q_{\mathrm{sat}}\left(T,p\right)}-\boldsymbol{q},$
(6)

still has a corresponding expansion of the PDF with warming as a result of the
Clausius-Clapeyron relation (Tab 1, row 2). In contrast, the relative humidity
rescaling (Fig 4a, bottom):

$\boldsymbol{\tilde{q}_{\mathrm{RH}}\left(p\right)}\overset{\mathrm{def}}{=}\boldsymbol{\mathrm{RH}\left(q,T,p\right)}$
(7)

normalizes specific humidity by its saturation value. As a result, most of the
RH PDF lies within $\left[0,1\right]$, except for a few atmospheric columns
exhibiting super-saturation in SPCAM, and that PDF changes little as the
climate warms (Tab 1, row 3). In addition to capturing grid-scale saturation,
$\boldsymbol{\tilde{q}_{\mathrm{RH}}}$ helps predict the subgrid-scale effects
of dry-air entrainment, known to regulate tropical convection (66, 67, 68)
(See SI B1a for details of RH and saturation specific humidity calculations).

#### Temperature

As expected, the PDF of temperature $\boldsymbol{T}$ (Fig 4b, top) shifts
quasi-linearly as the climate warms (Tab 1, row 4). Assuming that this shift
is almost uniform with height, we can derive an (approximate) invariant by
subtracting the temperature at all levels $\boldsymbol{T\left(p\right)}$ from
the near-surface temperature $T\left(p_{\mathrm{NS}}\right)$ (Fig 4b, middle):

$\boldsymbol{\tilde{T}_{\mathrm{from\
NS}}\left(p\right)}\overset{\mathrm{def}}{=}T\left(p_{\mathrm{NS}}\right)-\boldsymbol{T\left(p\right)},$
(8)

where $p_{\mathrm{NS}}$ is the lowest atmospheric pressure level (Tab 1, row
5).

However, this linear rescaling fails in the upper atmosphere, especially near
the tropopause where temperatures are approximately invariant with warming
(69, 70, 71) and therefore decoupled from surface temperature changes (Tab 1,
row 8). This motivates deriving a second temperature rescaling directly
relevant for moist convection: the buoyancy of a non-entraining, moist static
energy-conserving plume (Fig 4b, bottom, see SI B1b for its derivation):

$\boldsymbol{\tilde{T}_{\mathrm{buoyancy}}\left(p\right)}\overset{\mathrm{def}}{=}\boldsymbol{B_{\mathrm{plume}}}\left(q_{\mathrm{NS}},\boldsymbol{T,p}\right),$
(9)

where $q_{\mathrm{NS}}=q(p_{\mathrm{NS}})$ is the near-surface specific
humidity. This rescaling is inspired by recently introduced lower-tropospheric
buoyancy measures (72, 73), but with an extension to the full troposphere
(74). While Eq 9 does not explicitly include entrainment effects, the mapping
of $\boldsymbol{\tilde{T}_{\mathrm{buoyancy}}\left(p\right)}$ and
$\boldsymbol{\tilde{q}_{\mathrm{RH}}\left(p\right)}$ to heating and moisture
sink will implicitly include these. This rescaling captures leading order
effects needed to yield approximate climate invariance (Fig 4b) (Tab 1, rows 9
and 6). It also increases physical interpretability: Eq 9 links the vertical
temperature structure and near-surface humidity changes to a metric that
correlates well with deep convective activity (75). The presence of
$q_{\mathrm{NS}}$ in Eq 9 captures the role of near-surface humidity relative
to the temperature structure aloft in contributing to moist convective
instability.

#### Latent Heat Flux

The last input whose distribution significantly changes with warming is the
latent heat flux LHF (Fig 4c, top), and we refer the curious reader to SI B1c
for the remaining inputs that change less with warming (sensible heat flux and
surface pressure). Similarly to specific humidity, the increase of latent heat
fluxes with warming is directly linked to the Clausius-Clapeyron relation
(e.g., 76) (Tab 1, row 10), which motivates scaling LHF by near-surface
specific humidity $q\left(p_{\mathrm{NS}}\right)$ (Fig 4c, middle):

$\tilde{\mathrm{LHF}}_{q}\overset{\mathrm{def}}{=}\frac{\mathrm{LHF}}{L_{v}\max\left\\{\epsilon_{q},q\left(p_{\mathrm{NS}}\right)\right\\}},$
(10)

where $\epsilon_{q}$ is a user-chosen parameter that we can set to $10^{-4}$
to avoid division by zero. While better than directly using LHF (Tab 1, row
11) this rescaling fails for very dry atmospheres when the latent heat flux is
negative, e.g. in polar oceans where atmospheric water vapor may be condensing
on the surface, or when the near-surface specific humidity is very small, e.g.
in subtropical regions. To address this, we leverage the bulk aerodynamic
formula to provide a physics-motivated rescaling of LHF using the near-surface
saturation deficit (Fig 4c, bottom):

$\tilde{\mathrm{LHF}}_{\Delta
q}\overset{\mathrm{def}}{=}\frac{\mathrm{LHF}}{L_{v}\max\left\\{\epsilon_{q},q_{\mathrm{sat}}\left[T\left(p_{\mathrm{NS}}\right),p_{\mathrm{NS}}\right]-q\left(p_{\mathrm{NS}}\right)\right\\}}.$
(11)

Eq 11 rescales LHF, a thermodynamic variable, into
$\tilde{\mathrm{LHF}}_{\Delta q}$, approximately proportional to the magnitude
of near-surface horizontal winds and density (e.g., 76), whose distributions
vary less with warming. While both the Jensen-Shannon and Hellinger distances
would suggest that $\tilde{\mathrm{LHF}}_{\Delta q}$ is a slightly less good
rescaling than $\tilde{\mathrm{LHF}}_{q}$ (Tab 1, row 12) the
$\tilde{\mathrm{LHF}}_{\Delta q}$ rescaling leads to improved generalization
performance compared to $\tilde{\mathrm{LHF}}_{q}$ (not shown). This confirms
that only considering the PDF distances is insufficient to find the optimal
rescalings and underlines the importance of Fig 3’s workflow.

We now show that all three input rescalings
($\boldsymbol{\tilde{q}_{\mathrm{RH}}\left(p\right)}$,
$\boldsymbol{\tilde{T}_{\mathrm{buoyancy}}\left(p\right)}$, and
$\tilde{\mathrm{LHF}}_{\Delta q}$) lead to statistically significant
improvements in the ML models’ ability to generalize.

## Results

This section is organized as follows. After demonstrating the benefits of
progressively rescaling the ML models’ inputs (Fig 5), we show how “climate-
invariant” models learn subgrid closures across climates and geographies
during training (Fig S5, Tab 2). We then discuss the global skill of different
models after training (Fig 6, Fig S6, Fig 7) before finally investigating the
structure of “climate-invariant” mappings to understand why they generalize
better across climates of different temperatures (Fig 8).

Figure 5: Mean-squared error of subgrid thermodynamic tendencies versus
pressure, horizontally-averaged over the Tropics of SPCAM3 aquaplanet
simulations, for the four model outputs: total moistening ($\dot{q}$), total
heating ($\dot{T}$), longwave heating (lw) and shortwave heating (sw). While
the error is small for all NNs in the test set of the cold training climate
(a), it is much larger in the test set of the warm test climate (b), but
becomes progressively smaller as the inputs are successively rescaled: First
no rescaling (blue), then specific humidity (orange), then temperature
(green), then latent heat fluxes (red). For reference, the purple line depicts
an NN trained in the warm climate. Given that the generalization error of the
brute-force NN (blue line) greatly exceeds that of the rescaled NNs, we zoom
in on each panel to facilitate visualization.

Figure 6: Coefficient of determination $R^{2}$ for 500hPa subgrid heating of
brute-force (a) and climate-invariant (b) NNs trained using the cold (-4K)
training set of each model (SPCAM3, SPCAM5, and SAM). $R^{2}$ is calculated
over the test sets of each cold (-4K) and warm (+4K) simulations.

Figure 7: Same as Fig 6, but averaging the samples used to calculate $R^{2}$ for subgrid heating over (time,longitude) to obtain a (latitude,pressure) map. |  | Training Set | Validation Set | Training Set | Training Set
---|---|---|---|---|---
Trained | Model | Same temperature | Same temperature | Different temperature | Same temperature
on |  | Same geography | Same geography | Same geography | Different geography
| MLR BF | 287 (288) | 288 (289) | 759 (785) | 242 (265)
Cold | MLR CI | 294 (294) | 295 (295) | 671 (671) | 245 (265)
Aquaplanet | NN BF | 172 (172) | 172 (173) | 2167 (2242) | 1098 (1117)
(-4K) | NN BF+DN | 222 (222) | 205 (206) | 3225 (3509) | 258 (318)
SPCAM3 | NN CI | 168 (169) | 169 (169) | 422 (425) | 297 (316)
| NN CI+DN | 214 (214) | 198 (198) | 483 (487) | 224 (227)
| MLR BF | 642 (644) | 643 (645) | 328 (330) | 507 (543)
Warm | MLR CI | 650 (651) | 650 (651) | 315 (315) | 529 (598)
Aquaplanet | NN BF | 363 (364) | 362 (364) | 386 (472) | 1354 (1612)
(+4K) | NN BF+DN | 473 (473) | 434 (436) | 248 (248) | 728 (1275)
SPCAM3 | NN CI | 354 (355) | 355 (356) | 199 (202) | 1241 (1293)
| NN CI+DN | 451 (451) | 412 (414) | 210 (212) | 567 (659)
| MLR BF | 197 (197) | 197 (197) | 508 (562) | 299 (299)
Cold | MLR CI | 203 (203) | 202 (202) | 459 (459) | 307 (307)
Earth-like | NN BF | 114 (114) | 116 (116) | 709 (803) | 247 (251)
(-4K) | NN BF+DN | 165 (165) | 152 (152) | 512 (575) | 247 (247)
SPCAM5 | NN CI | 112 (112) | 114 (114) | 292 (294) | 312 (323)
| NN CI+DN | 158 (158) | 144 (144) | 346 (347) | 229 (233)
Table 2: Mean-Squared Error (MSE, in units W2 m-4) of six models trained in
three simulations (first column) and evaluated over the training or validation
set of the same and two other simulations (last four columns). The models
(second column) are brute-force (BF) or climate-invariant (CI), multiple
linear regressions (MLR) or neural nets (NN), and sometimes include dropout
layers preceded by a batch normalization layer (DN). The models are trained
for 20 epochs and we give the MSE corresponding to the epoch of minimal
validation loss followed by the MSE averaged over the 5 epochs with lowest
validation losses (in parentheses). In each application case, we highlight the
best model’s error using bold font.

Figure 8: SHAP feature importance matrix for the (a) brute-force and (b)
climate-invariant neural nets trained in the cold (-4K) aquaplanet simulation
(left), the warm (+4) warm aquaplanet simulation (middle), and their
difference (right). To calculate these matrices, we sample inputs from the
(+4K) warm aquaplanet simulation for all ML models to facilitate inter-model
comparison. The x-axes indicate the vertical levels of the inputs, from the
surface (left, 103hPa) to the top of the atmosphere (right, 0hPa), while the
y-axes indicate the vertical levels of the outputs, from the surface (bottom,
103hPa) to the top of the atmosphere (top, 0hPa). We additionally indicate the
200hPa vertical level with dotted black lines.

### Benefits of Incremental Input Rescalings

In this section, we demonstrate that incrementally rescaling the inputs of
neural networks progressively improves their generalization abilities from the
cold (-4K) aquaplanet (SPCAM3) simulation to the warm (+4K) aquaplanet
simulation. The largest surface temperature jump tested in this study is
between the cold aquaplanet simulation and the tropics of the warm aquaplanet
simulation (“warm Tropics” for short), defined as the regions with out-of-
distribution surface temperatures, whose latitudes are between -15°S and 23°N.
These regions approximately correspond to the red regions in the top-left
subplots of Fig 2. Therefore, to expose the failure modes of the “brute force”
model and the benefits of progressively rescaling the inputs, we first trained
several NNs on the cold aquaplanet until they reached high accuracy (Fig 5a)
before testing their generalization abilities out-of-distribution in the warm
Tropics (Fig 5b).

In the cold Tropics, the vertical profiles of the mean-squared error (MSE) are
nearly indiscernible for all types of NNs and roughly follow the vertical
structure of subgrid variance, as discussed in (45, 43). When evaluated in the
warm Tropics, the MSE of the “brute force” NN (blue line) increases by a
factor of $\approx$10 and peaks above 100W 2m-4, underlining how standard NNs
fail to generalize across climates! As discussed in the “Theory” Section, we
progressively rescaled the inputs starting with the most important one:
specific humidity, which is rescaled to relative humidity (orange line). This
first rescaling decreases the MSE so much (by a factor of 5-10) that we need
to zoom in on each panel to distinguish the generalization abilities of
additional NNs. Adding the rescalings of temperature to plume buoyancy (green
line) and then of LHF to $\tilde{\mathrm{LHF}}_{\Delta q}$ (red line)
incrementally decrease the MSE, which is particularly clear for the subgrid
longwave heating. Impressively, rescaling all three inputs decreases the MSE
so much that the resulting “climate-invariant” NN’s MSE (red line) is within
$\approx 25\%$ of the MSE of a “brute-force” NN that was directly trained in
the warm climate (purple line).

Hereafter, we use “climate-invariant” to refer to models for which all three
inputs ($q$, $T$, LHF) but no outputs (i.e., we do not assume knowledge of the
output CDF across climate and geographies) were physically rescaled; after
demonstrating their success in the aquaplanet case, we are now ready to
investigate how these “climate-invariant” models learn in other climates and
simulations introduced in the “Data” Section.

### Learning across Climates and Geographies

In this section, we show that “climate-invariant” models learn mappings that
are valid across climates and geographies, especially in conjunction with BN
and DP layers.

Tab 2 shows the MSE of ML models trained in three different datasets, and
evaluated over their training sets, validation sets, and over test sets of
different temperature and geography. As discussed previously, “climate-
invariant” NNs (NN CI) generalize better to warmer climates than “brute-force”
NNs (NN BF). We can go one step further by examining learning curves, defined
as the MSE of an ML model at the end of each epoch 444One epoch corresponds to
the ML model being fed the entire training set once. Impressively, the
learning curve of the “climate-invariant” NN trained in the cold aquaplanet
but tested in the warm aquaplanet (starred blue line in Fig S5a) is mostly
decreasing, which is a key result of this manuscript: “climate-invariant” NNs
are able to continuously learn about subgrid thermodynamics in the warm
aquaplanet as they are trained in the cold aquaplanet. In contrast, the
“brute-force” NN trained in the cold aquaplanet but tested in the warm
aquaplanet makes extremely large generalization errors, which worsen as the
model is trained in the cold aquaplanet (see SI D2 for details).

“Climate-invariant” NNs also facilitate learning across geographies, i.e.,
from the aquaplanet to the Earth-like simulations and vice-versa (see NN CI
rows in Tab 2). “Climate-invariant” rescalings additionally improve the MLR
baseline’s generalization ability, albeit less dramatically. This smaller
improvement in MLR’s generalization abilities is linked to its relatively
small number of free parameters, resulting in (1) “brute-force” MLRs
generalizing better than “brute-force” NNs; and (2) MLRs having lower
representation power and fitting their training sets less well, limiting the
maximal accuracy of “climate-invariant” MLRs on test sets.

There are a few cases in which rescaling inputs does not fully solve the
generalization problem, e.g., when trying to generalize from the aquaplanet to
the Earth-like simulation. In that case, we leverage the fact that input
rescalings can easily be combined with standard ML techniques to improve
generalization, such as DP layers before each activation function and a single
BN layer before the first DP layer (30). As DP layers randomly drop a fixed
proportion of the free parameters during training, NNs with DP fit their
training set less well (see NN CI+DN row of Tab 2). However, they improve
generalization in difficult cases (e.g., between cold aquaplanet and Earth-
like simulations) and do not overly deteriorate generalization in cases where
the input rescalings work particularly well (e.g., from cold to warm
aquaplanet). Our results suggest that combining physics-guided generalization
methods (e.g., physical rescaling of the inputs/outputs) with standard ML
generalization methods (e.g., DP) is advantageous and deserves further
investigation. After analyzing the overall MSE during training, we now turn to
the spatial characteristics of our ML models’ skill after training.

### Global Performance after Training

In this section, we first compare the spatial characteristics of the “brute
force” and “climate-invariant” NNs’ skill across climates of different
temperatures to further establish the advantages of our “climate-invariant”
rescaling. These advantages are clearly visible in Fig 6, where the “brute-
force” models struggle to generalize to the warm Tropics for all simulations
despite fitting the cold training set well (Fig 6a). We can trace these
generalization errors to warm surface temperature conditions the NNs were not
exposed to during training, visible when comparing Fig 6a with Fig 2a. In
contrast, the “climate-invariant” models fit the warm climate almost as well
as the cold climate they were trained in (Fig 6b)! Note that Fig 6 focuses on
the horizontal map of a single output, i.e., the total subgrid heating at
500hPa, but that the horizontal map of other outputs, such as the near-surface
subgrid heating (see Fig S5) all exhibit the same pattern of “brute-force”
models failing in the warm Tropics and the “climate-invariant” models mostly
correcting these generalization errors. Finally, the spatial distribution of
the skill in the training set (e.g., middle column of Fig S6) is reassuringly
consistent with the skill map of highly tuned NNs trained in similar
conditions (e.g., 44). This confirms that our “brute-force” models are
representative of state-of-the-art ML subgrid closures, and yet fail to
generalize.

This failure is confirmed in the latitude-pressure map of the subgrid heating
at all vertical levels shown in Fig 7a. In contrast, Fig 7b underlines how
“climate-invariant” NNs improve generalization throughout the atmosphere in
the warm Tropics without deteriorating skill in the mid-latitudes and poles of
the warm simulation. This consideration helped us choose our final input
rescaling, as the $\boldsymbol{\tilde{T}_{\mathrm{from\ NS}}}$ temperature
rescaling significantly deteriorated generalization in the mid-latitudes,
while the $\boldsymbol{\tilde{T}_{\mathrm{buoyancy}}}$ rescaling helps
generalization in the Tropics without overly compromising skills at other
latitudes. There is a slight skill compromise at high latitudes, as can be
seen by comparing the second rows of Fig 7a and Fig 7b, which is especially
apparent in the SAM case and can be partially traced back to challenges in
generalizing subgrid ice sedimentation (not shown here, see 46, for details).

To fully compare ML models across climate and geographies, we list their
overall MSEs in the training, validation, and both generalization test sets in
Table 2. In addition to the MSE of minimal validation loss, we show the MSE
averaged over the 5 epochs of minimal validation loss in parentheses to
confirm that our models have converged. Consistent with the learning curves
shown in Fig S5, the generalization ability of “climate-invariant” NNs with DP
and BN layers is best (two rightmost columns of each row’s NN CI+DN models).
While they fit their training sets less well, “brute-force” MLRs generalize
better than “brute-force” NNs because they have fewer free parameters (see MLR
BF and NN BF models). In Tab 2, we also show that while DP and BN layers
generally increase the generalization performance of “brute-force” NNs555We
have verified that in our case, DP and BN layers combined outperform DP layers
only and BN layers only (not shown). (NN BF+DN models), we can systematically
improve these standard ML generalization methods by combining them with input
rescalings (NN CI+DN models).

Now that we have fully established the advantages of physical rescalings, we
ask: Why do “climate-invariant” models generalize better than “brute force”
ones?

### Understanding Climate Invariance

In this section, we use a game theory-based explainable ML library, namely
SHAP (77, 78), to dissect “climate-invariant” mappings and provide insight
into why they generalize better across climates and geographies. Note that as
MLRs are interpretable by construction, we can draw preliminary conclusions by
visualizing MLR weights without the need for explainable ML libraries.
Similarly, we can directly visualize the linear responses (71, 79, 80, 81) of
NNs by calculating their Jacobians via automatic differentiation (e.g., 40).
However, the difference between “brute-force” and “climate-invariant” MLRs is
small and Jacobians cannot be reliably used to explain nonlinear NN
predictions as they are first-order derivatives (82). Therefore, we present
these two visualization methods in SI D4, and focus here on explaining the
generalization abilities of NNs using SHAP, recalling that “climate-invariant”
NNs are our best models at generalizing from cold to warm climates (see NN CI
model in the first row and “Different temperature” column of Tab 2).

We select SHAP’s Kernel Explainer because of its versatility, as it can be
used for any multi-inputs/outputs ML model, and specifically its “special
weighted linear regression” option as our outputs are continuous making
emulating subgrid closures a regression problem. For each output $y_{j}$ of
our ML model, SHAP estimates the effect of an input taking a particular value
$x_{i}$; SHAP is built based on a local accuracy property ensuring that the
sum of each individual input’s effect equals how different $y_{j}$ is from its
average over the training set because of the whole input vector taking a
particular value $\boldsymbol{x}$:

$\sum_{i=1}^{2N_{p}+4}\mathrm{SHAP}\left(x_{i},y_{j}\right)\overset{\mathrm{def}}{=}y_{j}^{\prime},$
(12)

where we have introduced the deviation $y^{\prime}_{j}$, defined as the
difference between $y_{j}$ and its training set average
$y_{i}^{\prime}\overset{\mathrm{def}}{=}y_{i}-\left\langle
y_{i}\right\rangle_{{\cal E}}$. We use these “Shapley values” to build a
nonlinear feature matrix $\cal M$ representing the effect of an input $x_{i}$
on an output $y_{j}$:

${\cal
M}_{\mathrm{ij}}\overset{\mathrm{def}}{=}\left\langle\mathrm{sign}\left(x_{i}^{\prime}\right)\times\mathrm{SHAP}\left(x_{i},y_{j}\right)\right\rangle_{{\cal
E}},$ (13)

where we use the sign of the input deviation $x^{\prime}_{i}$ to make ${\cal
M}_{\mathrm{ij}}$ positive if $x^{\prime}_{i}$ and $y^{\prime}_{j}$ have the
same sign, e.g. if a positive input deviation leads to a positive output
deviation666In the particular case of the MLR defined in Eq 18 of the SI, the
nonlinear feature matrix becomes the regression weight matrix multiplied by
the absolute value of the input deviations: ${\cal
M}_{\mathrm{ij}}=\left\langle
A_{ij}\left|x_{i}^{\prime}\right|\right\rangle_{{\cal E}}$, confirming that
the feature matrix ${\cal M}$ is a nonlinear extension of the Jacobian ($A$ in
the MLR case)..

In Fig 8, each rectangle depicts the SHAP feature matrix $\cal{M}$ for a given
model: “brute-force” (a) and “climate-invariant” (b), both trained in the warm
aquaplanet simulation (second column). Focusing on a single feature matrix,
the inputs of each model (e.g., specific humidity $q$ and temperature $T$) are
organized on the x-axis from the surface ($10^{3}$hPa) to the top of the
atmosphere (0hPa). The outputs of each model (subgrid moistening $\dot{q}$ and
subgrid heating $\dot{T}$; see Fig S11 for subgrid longwave heating
$\mathrm{lw}$ and subgrid shortwave heating $\mathrm{sw}$) are organized on
the y-axis, from the surface to the top of the atmosphere. We use dotted lines
to depict the 200hPa isobars. Following a horizontal line shows how different
inputs contribute to a given output, while following a vertical line shows how
a given input influences different outputs.

Fig 8 contains a wealth of information about subgrid closures trained in
aquaplanet simulations; we focus here on visualizing how the “climate-
invariant” NNs (b) operate in ways that generalize better than their “brute-
force” counterpart (a). Consider the row for subgrid heating $\dot{T}$. In the
“brute-force” case (a), $\cal{M}$ has large coefficients in the entire
troposphere (in the entire square below the dashed lines depicting the
approximate tropopause level). This means that specific humidity and
temperature deviations at all levels impact subgrid heating at a given level,
i.e. there are large non-local relations in the vertical. Some non-local
relations are physically plausible for convection, since buoyant plumes tend
to rise from the surface, and indeed near-surface $T$ and $q$ influence
$\dot{T}$ through the entire troposphere. However, moisture at higher
altitudes influences $\dot{T}$ at lower altitudes, raising suspicions that
some aspects of this non-locality are not causal, as (40) show can happen in
brute force NNs. Temperature variations are observed to have strong vertical
correlations (83) in part because of deep convective effects. Because
temperature affects the saturation threshold for moisture, the NN will have to
correctly capture the effects of both temperature and moisture wherever either
has influence. In contrast, in the $B_{plume}$-RH “climate-invariant” case
(b), leading nonlocal effects between the boundary layer and the free
troposphere have already been taken into account in the buoyancy formulation,
and the temperature-dependent saturation threshold is built into RH. Thus,
$\cal{M}$ for $\dot{T}$ tends to be concentrated near the red diagonal,
meaning that positive deviations of plume buoyancy and relative humidity
increase subgrid heating near the same vertical level. The use of domain
knowledge has effectively reduced the effects that must be estimated by the NN
for the “climate-invariant” models. This tends to yield differences between
the models trained in the cold and the warm climates that are much smaller
than for the “brute-force” models (see last column of Fig S11).

## Conclusion

In the context of climate change, we hypothesized that ML models emulating
“climate-invariant” mappings (Fig 1, Fig 3), for which the distributions of
inputs/outputs change little across climates (Fig 4), generalize much better
than ML models emulating “brute-force” mappings, for which the distributions
of inputs/outputs change significantly across climates. Tested on a suite of
storm-resolving simulations with different surface temperatures in three
distinct climate models with different geographies (Fig 2), physically-
rescaled NNs generalize better to a warmer climate as their inputs are
progressively rescaled (Fig 5). “Climate-invariant” NNs whose inputs have all
been rescaled learn mappings that hold across temperatures and geographies
during training (Tab 2), and hence exhibit superior generalization skill
almost everywhere on the globe (Fig 6, Fig 7). Finally, attribution maps
reveal that “climate-invariant” NNs learn more spatially local mappings that
facilitate generalization across climates (Fig 8).

From a computational perspective, incorporating physical knowledge of climate
change into an ML framework to improve its generalization skill is a
successful example of using domain knowledge to extract more informative
predictors, informally referred to as “feature engineering” (e.g., 84). This
also aids interpretability of the mapping. From a climate science perspective,
requiring that a nonlinear statistical model of the atmosphere generalize
across climate is a stringent test that helped us discover a new mapping. As
this “climate-invariant” mapping is more robust to climate and geography
changes, it appears to be more advantageous than directly using model and
observational outputs (e.g., specific humidity, temperature), even when data
are available in various climate regimes. In the particular case of subgrid
thermodynamics, our generalization results suggest the possibility of NN-
powered closures that could work in Earth-like settings, even in vastly
different climate conditions. The attribution maps suggest the possibility of
new analytic representations of convection from data, facilitated by the more
local “climate-invariant” representation of subgrid thermodynamics.

TB and MP acknowledge funding from NSF grants OAC1835863 and AGS-1734164. MP
acknowledges partial funding from the DOE “Enabling Aerosol-cloud interactions
at GLobal convection-permitting scalES (EAGLES)” project (74358), funded by
the U.S. Department of Energy, Office of Science, Office of Biological and
Environmental Research, Earth System Model Development (ESMD) program. JY
acknowledges support from the EAPS Houghton-Lorenz postdoctoral fellowship. PG
thanks NSF grant AGS-1734164 as well as USMILE European Research Council
grant. PG, POG, and JY acknowledge funding from M2LInES Schmidt Future
funding. FA and JDN acknowledge NSF AGS-1936810. We also thank the Extreme
Science and Engineering Discovery Environment supported by NSF grant number
ACI-1548562 (charge numbers TG-ATM190002 and TG-ATM170029), and HPC support
from Cheyenne (doi:10.5065/D6RX99HX) provided by NCAR’s Computational and
Information Systems Laboratory (sponsored by the NSF) for computational
resources. We thank William Boos for providing the SAM hypohydrostatic
simulation output, Andrea Jenney, Imme Ebert-Uphoff, and the CBRAIN discussion
group for help and advice that greatly benefited the present manuscript.

## References

  * (1) M Reichstein, et al., Deep learning and process understanding for data-driven earth system science.  Nature 566, 195–204 (2019).
  * (2) T Beucler, I Ebert-Uphoff, S Rasp, M Pritchard, P Gentine, Machine learning for clouds and climate in Clouds and Climate.  (AGU Geophysical Monograph Series), (2021).
  * (3) C Irrgang, et al., Towards neural earth system modelling by integrating artificial intelligence in earth system science.  Nature Machine Intelligence 3, 667–674 (2021).
  * (4) P Ukkonen, R Pincus, RJ Hogan, K Pagh Nielsen, E Kaas, Accelerating radiation computations for dynamical models with targeted machine learning and code optimization.  Journal of Advances in Modeling Earth Systems 12, e2020MS002226 (2020).
  * (5) A Belochitski, et al., Tree approximation of the long wave radiation parameterization in the ncar cam global climate model.  Journal of Computational and Applied Mathematics 236, 447–460 (2011).
  * (6) JJ Gristey, G Feingold, IB Glenn, KS Schmidt, H Chen, On the relationship between shallow cumulus cloud field properties and surface solar irradiance.  Geophysical Research Letters 47, e2020GL090152 (2020).
  * (7) R Lagerquist, D Turner, I Ebert-Uphoff, J Stewart, V Hagerty, Using deep learning to emulate and accelerate a radiative transfer model.  Journal of Atmospheric and Oceanic Technology 38, 1673–1696 (2021).
  * (8) M Chantry, S Hatfield, P Dueben, I Polichtchouk, T Palmer, Machine learning emulation of gravity wave drag in numerical weather forecasting.  Journal of Advances in Modeling Earth Systems 13, e2021MS002477 (2021).
  * (9) Z Espinosa, A Sheshadri, G Cain, E Gerber, K DallaSanta, Machine learning emulation of parameterized gravity wave momentum fluxes in an atmospheric global climate model in EGU General Assembly Conference Abstracts.  pp. EGU21–1398 (2021).
  * (10) D Matsuoka, et al., Application of deep learning to estimate atmospheric gravity wave parameters in reanalysis data sets.  Geophysical Research Letters 47, e2020GL089436 (2020).
  * (11) H Morrison, M van Lier-Walqui, MR Kumjian, OP Prat, A bayesian approach for statistical–physical bulk parameterization of rain microphysics. part i: Scheme description.  Journal of the Atmospheric Sciences 77, 1019–1041 (2020).
  * (12) A Seifert, S Rasp, Potential and limitations of machine learning for modeling warm-rain cloud microphysical processes.  Journal of Advances in Modeling Earth Systems 12, e2020MS002301 (2020).
  * (13) A Gettelman, et al., Machine learning the warm rain process.  Journal of Advances in Modeling Earth Systems 13, e2020MS002268 (2021).
  * (14) B François, S Thao, M Vrac, Adjusting spatial dependence of climate model outputs with cycle-consistent adversarial networks.  Climate Dynamics 57, 3323–3353 (2021).
  * (15) B Pan, et al., Learning to correct climate projection biases.  Journal of Advances in Modeling Earth Systems 13, e2021MS002509 (2021).
  * (16) V Zantedeschi, et al., Cumulo: A dataset for learning cloud classes.  arXiv preprint arXiv:1911.04227 (2019).
  * (17) S Rasp, H Schulz, S Bony, B Stevens, Combining crowdsourcing and deep learning to explore the mesoscale organization of shallow convection.  Bulletin of the American Meteorological Society 101, E1980–E1995 (2020).
  * (18) D Watson-Parris, S Sutherland, M Christensen, R Eastman, P Stier, A large-scale analysis of pockets of open cells and their radiative impact.  Geophysical Research Letters 48, e2020GL092213 (2021).
  * (19) L Denby, Discovering the importance of mesoscale cloud organization through unsupervised classification.  Geophysical Research Letters 47, e2019GL085190 (2020).
  * (20) Y Han, GJ Zhang, X Huang, Y Wang, A moist physics parameterization based on deep learning.  Journal of Advances in Modeling Earth Systems 12, e2020MS002076 (2020).
  * (21) ND Brenowitz, CS Bretherton, Prognostic validation of a neural network unified physics parameterization.  Geophysical Research Letters 45, 6289–6298 (2018).
  * (22) VM Krasnopolsky, MS Fox-Rabinovitz, AA Belochitski, Using ensemble of neural networks to learn stochastic convection parameterizations for climate and numerical weather prediction models from data simulated by a cloud resolving model.  Advances in Artificial Neural Systems 2013 (2013).
  * (23) L Breiman, Random forests.  Machine learning 45, 5–32 (2001).
  * (24) PA O’Gorman, JG Dwyer, Using machine learning to parameterize moist convection: Potential for modeling of climate, climate change, and extreme events.  Journal of Advances in Modeling Earth Systems 10, 2548–2563 (2018).
  * (25) S Rasp, MS Pritchard, P Gentine, Deep learning to represent subgrid processes in climate models.  Proceedings of the National Academy of Sciences 115, 9684–9689 (2018).
  * (26) T Beucler, M Pritchard, P Gentine, S Rasp, Towards physically-consistent, data-driven models of convection in IGARSS 2020-2020 IEEE International Geoscience and Remote Sensing Symposium.  (IEEE), pp. 3987–3990 (2020).
  * (27) MJ Molina, DJ Gagne, AF Prein, A benchmark to test generalization capabilities of deep learning methods to classify severe convective storms in a changing climate.  Earth and Space Science Open Archive ESSOAr (2021).
  * (28) Y LeCun, Y Bengio, , et al., Convolutional networks for images, speech, and time series.  The handbook of brain theory and neural networks 3361, 1995 (1995).
  * (29) C Liu, et al., Continental-scale convection-permitting modeling of the current and future climate of north america.  Climate Dynamics 49, 71–95 (2017).
  * (30) S Ioffe, C Szegedy, Batch normalization: Accelerating deep network training by reducing internal covariate shift in International conference on machine learning.  (PMLR), pp. 448–456 (2015).
  * (31) N Srivastava, G Hinton, A Krizhevsky, I Sutskever, R Salakhutdinov, Dropout: a simple way to prevent neural networks from overfitting.  The journal of machine learning research 15, 1929–1958 (2014).
  * (32) X Ying, An overview of overfitting and its solutions in Journal of Physics: Conference Series.  (IOP Publishing), Vol. 1168, p. 022022 (2019).
  * (33) SJ Pan, Q Yang, A survey on transfer learning.  IEEE Transactions on knowledge and data engineering 22, 1345–1359 (2009).
  * (34) VM Patel, R Gopalan, R Li, R Chellappa, Visual domain adaptation: A survey of recent advances.  IEEE signal processing magazine 32, 53–69 (2015).
  * (35) D Tuia, C Persello, L Bruzzone, Domain adaptation for the classification of remote sensing data: An overview of recent advances.  IEEE geoscience and remote sensing magazine 4, 41–57 (2016).
  * (36) D Randall, M Khairoutdinov, A Arakawa, W Grabowski, Breaking the Cloud Parameterization Deadlock.  Bulletin of the American Meteorological Society 84, 1547–1564 (2003).
  * (37) S Bony, et al., Clouds, circulation and climate sensitivity.  Nature Geoscience 8, 261–268 (2015).
  * (38) T Schneider, et al., Climate goals and computing the future of clouds.  Nature Climate Change 7, 3–5 (2017).
  * (39) P Gentine, V Eyring, T Beucler, Deep learning for the parametrization of subgrid processes in climate models.  Deep Learning for the Earth Sciences: A Comprehensive Approach to Remote Sensing, Climate Science, and Geosciences pp. 307–314 (2021).
  * (40) ND Brenowitz, T Beucler, M Pritchard, CS Bretherton, Interpreting and stabilizing machine-learning parametrizations of convection.  Journal of the Atmospheric Sciences 77, 4357–4375 (2020).
  * (41) J Ott, et al., A fortran-keras deep learning bridge for scientific computing.  Scientific Programming 2020 (2020).
  * (42) P Forster, et al., The earth’s energy budget, climate feedbacks, and climate sensitivity.  Climate Change (2021).
  * (43) T Beucler, et al., Enforcing analytic constraints in neural networks emulating physical systems.  Physical Review Letters 126, 098302 (2021).
  * (44) G Mooers, et al., Assessing the potential of deep learning for emulating cloud superparameterization in climate models with real-geography boundary conditions.  Journal of Advances in Modeling Earth Systems 13, e2020MS002385 (2021).
  * (45) P Gentine, M Pritchard, S Rasp, G Reinaudi, G Yacalis, Could machine learning break the convection parameterization deadlock?  Geophysical Research Letters 45, 5742–5751 (2018).
  * (46) J Yuval, PA O’Gorman, Stable machine-learning parameterization of subgrid processes for climate modeling at a range of resolutions.  Nature communications 11, 1–10 (2020).
  * (47) M Khairoutdinov, D Randall, C DeMott, Simulations of the atmospheric general circulation using a cloud-resolving model as a superparameterization of physical processes.  Journal of the Atmospheric Sciences 62, 2136–2154 (2005).
  * (48) MS Pritchard, CS Bretherton, Causal evidence that rotational moisture advection is critical to the superparameterized madden–julian oscillation.  Journal of the Atmospheric Sciences 71, 800–815 (2014).
  * (49) JA Andersen, Z Kuang, Moist static energy budget of mjo-like disturbances in the atmosphere of a zonally symmetric aquaplanet.  Journal of Climate 25, 2782–2804 (2012).
  * (50) MF Khairoutdinov, DA Randall, A cloud resolving model as a cloud parameterization in the ncar community climate system model: Preliminary results.  Geophysical Research Letters 28, 3617–3620 (2001).
  * (51) M Wang, et al., A multiscale modeling framework model (superparameterized cam5) with a higher-order turbulence closure: Model description and low-cloud simulations.  Journal of Advances in Modeling Earth Systems 7, 484–509 (2015).
  * (52) JW Hurrell, JJ Hack, D Shea, JM Caron, J Rosinski, A new sea surface temperature and sea ice boundary dataset for the community atmosphere model.  Journal of Climate 21, 5145–5153 (2008).
  * (53) SN Tulich, A strategy for representing the effects of convective momentum transport in multiscale models: Evaluation using a new superparameterized version of the weather research and forecast model (SP-WRF).  J. Adv. Model. Earth Syst. 7, 938–962 (2015).
  * (54) ND Brenowitz, CS Bretherton, Spatially extended tests of a neural network parametrization trained by coarse-graining.  Journal of Advances in Modeling Earth Systems 11, 2728–2744 (2019).
  * (55) T Bolton, L Zanna, Applications of deep learning to ocean data inference and subgrid parameterization.  Journal of Advances in Modeling Earth Systems 11, 376–399 (2019).
  * (56) ST Garner, D Frierson, I Held, O Pauluis, G Vallis, Resolving convection in a global hypohydrostatic model.  Journal of the atmospheric sciences 64, 2061–2075 (2007).
  * (57) WR Boos, A Fedorov, L Muir, Convective self-aggregation and tropical cyclogenesis under the hypohydrostatic rescaling.  Journal of the Atmospheric Sciences 73, 525–544 (2016).
  * (58) P O’Gorman, Z Li, W Boos, J Yuval, Response of extreme precipitation to uniform surface warming in quasi-global aquaplanet simulations at high resolution.  Philosophical Transactions of the Royal Society A 379, 20190543 (2021).
  * (59) RB Neale, BJ Hoskins, A standard test for agcms including their physical parametrizations: I: The proposal.  Atmospheric Science Letters 1, 101–107 (2000).
  * (60) A Karpatne, I Ebert-Uphoff, S Ravela, HA Babaie, V Kumar, Machine learning for the geosciences: Challenges and opportunities.  IEEE Transactions on Knowledge and Data Engineering 31, 1544–1554 (2018).
  * (61) TG Shepherd, Symmetries, conservation laws, and hamiltonian structure in geophysical fluid dynamics in Advances in Geophysics.  (Elsevier) Vol. 32, pp. 287–338 (1990).
  * (62) R Beran, Minimum hellinger distance estimates for parametric models.  The annals of Statistics pp. 445–463 (1977).
  * (63) DM Endres, JE Schindelin, A new metric for probability distributions.  IEEE Transactions on Information theory 49, 1858–1860 (2003).
  * (64) JR Holton, An introduction to dynamic meteorology.  American Journal of Physics 41, 752–754 (1973).
  * (65) AP Siebesma, S Bony, C Jakob, B Stevens, Clouds and Climate: Climate Science’s Greatest Challenge.  (Cambridge University Press), (2020).
  * (66) RG Brown, C Zhang, Variability of midtropospheric moisture and its effect on cloud-top height distribution during TOGA COARE.  Journal of the atmospheric sciences 54, 2760–2774 (1997).
  * (67) CS Bretherton, ME Peters, LE Back, Relationships between water vapor path and precipitation over the tropical oceans.  Journal of Climate 17, 1517–1528 (2004).
  * (68) CE Holloway, JD Neelin, Moisture vertical structure, column water vapor, and tropical deep convection.  J. Atmos. Sci. 66, 1665–1683 (2009).
  * (69) JT Seeley, N Jeevanjee, DM Romps, Fat or fitt: Are anvil clouds or the tropopause temperature invariant?  Geophysical Research Letters 46, 1842–1850 (2019).
  * (70) DL Hartmann, K Larson, An important constraint on tropical cloud-climate feedback.  Geophysical research letters 29, 12–1 (2002).
  * (71) Z Kuang, DL Hartmann, Testing the fixed anvil temperature hypothesis in a cloud-resolving model.  Journal of Climate 20, 2051–2057 (2007).
  * (72) F Ahmed, JD Neelin, Reverse engineering the tropical precipitation–buoyancy relationship.  Journal of the Atmospheric Sciences 75, 1587–1608 (2018).
  * (73) F Ahmed, ÁF Adames, JD Neelin, Deep convective adjustment of temperature and moisture.  Journal of the Atmospheric Sciences 77, 2163–2186 (2020).
  * (74) KA Schiro, F Ahmed, SE Giangrande, JD Neelin, GoAmazon2014/5 campaign points to deep-inflow approach to deep convection across scales.  Proceedings of the National Academy of Sciences 115, 201719842 (2018).
  * (75) F Ahmed, JD Neelin, Protected convection as a metric of dry air influence on precipitation.  Journal of Climate 34, 3821–3838 (2021).
  * (76) DL Hartmann, Global physical climatology.  (Newnes) Vol. 103, (2015).
  * (77) SM Lundberg, SI Lee, A unified approach to interpreting model predictions in Proceedings of the 31st international conference on neural information processing systems.  pp. 4768–4777 (2017).
  * (78) LS Shapley, 17\. A value for n-person games.  (Princeton University Press), (2016).
  * (79) MJ Herman, Z Kuang, Linear response functions of two convective parameterization schemes.  Journal of Advances in Modeling Earth Systems 5, 510–541 (2013).
  * (80) T Beucler, T Cronin, K Emanuel, A linear response framework for radiative-convective instability.  Journal of Advances in Modeling Earth Systems 10, 1924–1951 (2018).
  * (81) Z Kuang, Linear stability of moist convecting atmospheres. part i: From linear response functions to a simple model and applications to convectively coupled waves.  Journal of the Atmospheric Sciences 75, 2889–2907 (2018).
  * (82) PJ Kindermans, et al., The (un) reliability of saliency methods in Explainable AI: Interpreting, Explaining and Visualizing Deep Learning.  (Springer), pp. 267–280 (2019).
  * (83) CE Holloway, JD Neelin, The convective cold top and quasi equilibrium.  Journal of the atmospheric sciences 64, 1467–1487 (2007).
  * (84) A Zheng, A Casari, Feature engineering for machine learning: principles and techniques for data scientists.  (" O’Reilly Media, Inc."), (2018).