Enforcing Equity in Neural Climate Emulators

Abstract

Neural network emulators have become an invaluable tool for a wide variety of climate and weather prediction tasks. While showing incredibly promising results, these networks do not have an inherent ability to produce equitable predictions. That is, they are not guaranteed to provide a uniform quality of prediction along any particular class or group of people. This potential for inequitable predictions motivates the need for explicit representations of fairness in these neural networks. To that end, we draw on methods for enforcing analytical physical constraints in neural networks to bias networks towards more equitable predictions. We demonstrate the promise of this methodology using the task of climate model emulation. Specifically, we propose a custom loss function which punishes emulators with unequal quality of predictions across any prespecified regions or category, here defined using human development index (HDI). This loss function weighs a standard loss metric such as mean squared error against another metric which captures inequity along the equity category (HDI), allowing us to adjust the priority of each term before training. Importantly, the loss function does not specify a particular definition of equity to bias the neural network towards, opening the door for custom fairness metrics. Our results show that neural climate emulators trained with our loss function provide more equitable predictions and that the equity metric improves with greater weighting in the loss function. We empirically demonstrate that while there is a tradeoff between accuracy and equity when prioritizing the latter during training, an appropriate selection of the equity priority hyperparameter can minimize loss of performance.

Keywords AI for climate  $\cdot$ machine learning  $\cdot$ equity  $\cdot$ fairness

1 Introduction

Modern Earth system models (ESMs) are key in characterizing the Earth’s response to anthropogenic forcing, and their predictions have been widely used to guide global climate policy such as United Nations Paris Agreement. As such, the impact of the predictions from such models has the potential to reach every human being on Earth. By numerically solving equations which describe our understanding of the climate system, ESMs are able to predict the state of the planet under various future scenarios. The large computational cost of such ESMs [1, 2], however, has motivated recent applications of new machine learning (ML) techniques for climate prediction, in particular deep learning with neural networks [3, 4, 5, 6, 7, 8]. While successive iterations of these neural climate emulators have each pushed the frontier of stable climate predictions using machine learning, the metrics used to assess these models only focus on accuracy or, more rarely, sensitivity to forcing and physical consistency/plausibility [9, 10, 11].

One important aspect of climate predictions which is discussed even less frequently is their fairness or equity. This consideration is of particular importance for emulators given their increasing application across a wide variety of climate prediction tasks including atmospheric forecasting [5, 12, 4], subgrid-scale parameterization [13, 14, 15, 16], precipitation nowcasting [17, 18, 19], equation and knowledge discovery [20, 21, 22], data assimilation [23, 24], downscaling [25, 26, 27], bias correction [28, 29, 30], and more [31]. Such principles also align themselves with the explicit fairness goals of many climate policies as well as recent literature on policy development [32, 33]. Despite this, there has been little work to date discussing equity in neural climate emulator predictions. Nevertheless, the inherent process of training a neural network provides a unique opportunity to address this issue. Since researchers already define advanced loss functions to bias their emulators towards specific goals, such as long-term stability and forecast sharpness [12, 34], the loss function itself could open a path towards more equitable predictions if notions of fairness could be integrated within it.

In this work, we draw on methods for enforcing physical constraints (e.g., conservation of mass and energy) [35] in neural networks via the loss function to bias neural climate emulators of global temperature towards more equitable predictions. Specifically, we use a two-part loss function which weighs a standard error metric such as MSE against a measure of the model’s fairness on a sliding scale. In doing so we can adjust the model’s preference for predictions with equal error throughout different prespecified regions of the globe, defined here using Human Development Index [36]. Importantly, given the many existing quantitative measures of equity, our method does not specify a specific fairness metric and allows for any such metric to be weighed against error in the loss function. This is in contrast to previous methods which focus specifically on spatial biases in neural network predictions [37] or non-spatial tabular data [38]. To demonstrate our method’s capability, we train a neural climate emulator to predict surface air temperature and the diurnal temperature range through the year 2100 using our custom loss function. We show that emulators trained using the loss function behave as desired, providing more equitable predictions and improving the equity metric as it is given greater weighting in the loss function. Additionally, we demonstrate that the sliding scale nature of the loss function can give rise to a tradeoff between accuracy and equity when the latter is weighted heavily during training. We demonstrate that a sufficiently small choice of equity weighting could vastly improve the fairness metric while achieving only a small loss in prediction accuracy.

2 Equitable Loss Functions

The ultimate goal of this work is to bias neural climate emulators towards more equitable predictions by integrating notions of fairness into their objective functions. This is accomplished in a manner similar to that of previous methods for enforcing analytical physical constraints in neural networks [35, 39, 26]. The key idea is to balance a traditional measure of error, in our case mean squared error, against a penalty term which captures some representation of fairness in the neural networks. Specifically, the loss function is formulated as

$$\mathcal{L}=\alpha\mathcal{P}+(1-\alpha)\text{MSE},$$

where $\mathcal{P}$ is the equity penalty, MSE is global mean squared error, and $\alpha$ is the equity weighting coefficient. Under this construction, $\mathcal{P}$ may be based on any quantitative fairness metric. Given the lack of a singular definition for fairness, this flexibility is particularly important. We illustrate the effect of our loss function by defining fairness based on prediction accuracy in various predefined regions across the globe. Specifically, we divide the land of Earth into $n$ equally sized regions based on Human Development Index, a composite measure of a country’s development [36, 40]. Then for the term $\mathcal{P}$, we punish the coefficient of variation, or relative standard deviation, between the errors across those $n$ regions, encouraging the emulator to have equal predictive error in each region. While this is the only equity penalty presented in this work, we found that other definitions such as simple standard deviation and performance deviation from the most well predicted region lead to similar conclusions, with modestly worse overall performance.

Figure 1: The 5 HDI regions over which equity is defined in the penalty $\mathcal{P}$.

The particular penalty explored in this work seeks to equalize the loss function for different groups, which is akin to parity-based fairness metrics in statistical applications [41, 42]. More precisely, if $\text{MSE}_{i}$ is the MSE of the emulator’s prediction in region $i$, then the emulator’s MSE over all land is

$$\text{MSE}_{l}=\frac{1}{n}\sum_{i=0}^{n-1}\text{MSE}_{i}.$$

Furthermore, the standard deviation of the MSE’s for each region is

$$\sigma_{R}=\sqrt{\frac{\sum_{i=0}^{n-1}(\text{MSE}_{i}-\text{MSE}_{l})^{2}}{n}}.$$

Finally, our equity penalty for the climate emulator is

$$\mathcal{P}=\frac{\sigma_{R}}{\text{MSE}_{l}}.$$

We split the land grid cells into quantiles each containing 20% of the cells (i.e., $n=5$) based on the prespecified equity measure of interest, HDI. The spatial distribution of these 5 HDI regions are shown in Fig. 1.

3 Results

To demonstrate our equitable loss function methodology, we train an ensemble of neural networks to predict one of either surface air temperature (TAS) or diurnal temperature range (DTR) for the years 2080-2100 under an intermediate climate forcing scenario [2] using the ClimateBench dataset [3]. We emphasize again, however, that our method is agnostic to the training dataset and chosen definition of quantitative fairness. These neural climate emulators take as input annual means of carbon dioxide (CO₂), methane (CH₄), sulfur dioxide (SO₂), and black carbon (BC) on a 96 latitude $\times$ 144 longitude global grid and predict TAS and DTR at the same spatiotemporal resolution. (See Materials and Methods for further details.)

We generate ensembles of 15 models for each of $\alpha=0,0.01,0.05,0.1,0.25,0.5,0.75,$ and $1$. Then we compute the MSE of each trained ensemble member’s prediction, as well as the equity penalty. The results are shown in Fig. 5. Immediately obvious is the near monotonic increasing trend for MSE with increasing $\alpha$ and a corresponding near monotonic decreasing trend for the equity penalty $\mathcal{P}$. Given the construction of the loss function with equity weighting coefficient $\alpha$, this tradeoff between accuracy and fairness is expected. Interestingly, for low values of $\alpha$, there appear to be large relative gains in the equity penalty in exchange for very little relative loss in predictive power. This is highlighted in Fig. 5 which shows the same plot as Fig. 5 but only for $\alpha\leq 0.25$. For example, compared to $\alpha=0$, the ensemble trained with $\alpha=0.1$ sees a 32% improvement in the equity penalty for only an 8% loss in accuracy.

Figure 2: MSE (black) and $\mathcal{P}$ (blue) for each neural network ensemble trained to predict DTR with varying $\alpha$. Error bars represent standard error of the ensemble mean.

Figure 3: As in Fig. 5, but only for $\alpha\leq 0.25$.

Figure 4: MSE (black) and $\mathcal{P}$ (blue) for each neural network ensemble trained to predict TAS with varying $\alpha$. Error bars represent standard error of the ensemble mean.

Figure 5: As in Fig. 5, but only for $\alpha\leq 0.25$.

Figs. 5 and 5 show a similar tradeoff between performance and equity for the TAS prediction task. It seems, however, that the $\alpha$ value for which large gains in the equity penalty come at low cost in performance is between $\alpha=0.25$ and $\alpha=0.5$, as opposed to around $\alpha=0.1$ for DTR. We do not necessarily expect that this “sweet spot" value for $\alpha$ should be the same for both prediction tasks, as the errors for each will have different scales which can affect both the MSE and $\mathcal{P}$ values in the loss function.

For both DTR and TAS, certain ensembles trained with low $\alpha$ values exhibit interesting behavior where both the accuracy and equity penalty of the ensemble are modestly improved compared to the $\alpha=0$ ensemble, though the $\alpha$ value where this occurs is not consistent between DTR and TAS. This is especially surprising since the MSE component of the loss function is defined globally while the equity penalty is only defined over land, meaning that the emulator is not simply achieving more equitable predictions by significantly raising error over the ocean. This is illustrated in Fig. 6 which shows the emulator’s MSE over both land and sea, as well as the equity penalty, for $\alpha\leq 0.25$. From this, it is seen that the emulator’s performance over the ocean does not degrade with increasing $\alpha$ and in fact improves for certain values.

Figure 6: MSE (black) and $\mathcal{P}$ (blue) for each neural network ensemble trained to predict DTR with varying $\alpha$. The square and triangular points represent MSE over land and ocean, respectfully. Error bars represent standard error of the ensemble mean.

Overall, Fig. 5-5 highlight that it is indeed possible to bias neural climate emulators towards more equitable predictions with a loss function that explicitly encodes a fairness metric. While this comes at the cost of accuracy, an appropriate choice of a small equity weighting coefficient $\alpha$ (here, $\alpha\lessapprox 0.25$) can mitigate losses in performance or even improve global predictions, narrowing the $\alpha$ hyperparameter search space for those wishing to make use of equitable loss functions. This finding for low $\alpha$ values aligns with those of prior work [35] for enforcing analytical physical constraints in neural networks.

Figure 7: DTR MSE in each HDI region for several values of $\alpha$. The black dashed line shows the mean of the colored bars, or simply the MSE over land.

Figure 8: As in Fig. 8, but TAS.

Figure 9: MSE of the neural network’s DTR predictions averaged over the test data for $\alpha=0$, $0.05$, $0.1$, and $0.25$.

To further understand the performance of neural climate emulators trained with the equitable loss function, we break down the predictive error at several small $\alpha$ values by the HDI regions (see Figure 1). This is illustrated in Fig. 8, which shows the emulator’s predictive error in each HDI region for five increasing values of $\alpha$ (0, 0.01, 0.05, 0.1, and 0.25) for DTR, and Fig. 8, which shows the same results but for the set of emulators trained to predict TAS. These results highlight how error is redistributed between the HDI regions in order to lower the equity penalty as $\alpha$ increases. Specifically, the ensemble achieves more equitable predictions at higher values of $\alpha$ by lowering the relative deviation of the error in each region from the mean (black dashed line). This is accomplished by simultaneously lowering the error in poorly predicted regions and raising the error in well predicted regions. For both predicted variables, the $\alpha=0.1$ neural network ensemble achieves similar MSE performnce as the $\alpha=0$ case while also successfully reducing the spread in error between the HDI regions.

By redistributing error throughout the HDI regions as shown in Figs. 8 and 8, neural climate emulators trained with the equitable loss function make the spatial distribution of predictive error more uniform. That is, they make predictions of equal quality throughout the globe. This is illustrated in Fig. 9 which shows the DTR MSE averaged over the test period (2080-2100) for four increasing values of $\alpha$ (0, 0.05, 0.1, and 0.25). Most obviously for the $\alpha=0.25$ ensemble compared to $\alpha=0$, the error in the poorly predicted regions over the Tibetan plateau and southeast Asia decreases while error in other regions of the globe increases, effectively lowering the equity penalty. For the regions where predictions become worse at higher $\alpha$ values, the ensemble appears to distribute error randomly within them. This trend towards a random uniform distribution of error across the globe becomes more and more obvious with increasing $\alpha$.

4 Discussion and Conclusion

As weather and climate predictions from both pure ML and hybrid physics-ML models continue to improve according to global accuracy metrics, it remains equally important to ensure that their predictions are accurate for everyone, regardless of their location on the globe. By utilizing a novel loss function which balances a traditional error metric with a penalty for inequitable predictions, our method shows that neural climate emulators can be biased towards more fair predictions, regardless of the chosen definition of quantitative equity. Moreover, with proper weighting of the equity penalty term in the loss function, this comes at little cost to global accuracy. Overall, our results highlight that neural climate emulators can achieve fair predictions when trained towards an appropriate goal.

While we focus on one particular definition of equity in this work, another may easily be substituted in this proposed loss function penalty framework. There are interesting implications for the loss landscape when using different definitions of numerical equity. Just as mean squared error is preferred over mean absolute error when optimizing neural networks for certain problems [43, 44], some equity penalties will be more favorable from an optimization standpoint than others. Similarly, our choice to define equity regions based on HDI is not the only possible option. A more spatially coherent choice (e.g., using the continents as regions) is easily achievable using our method, and there are interesting open questions regarding the effects this may have on training and performance. Furthermore, many recent AI weather models tend to be optimized towards multi-objective loss functions, often to achieve accuracy on both short and long lead times. In a similar manner, future work may investigate the effects of training a neural climate emulator towards both equity and another goals such as obeying a physical constraint or maintaining sharpness in predictions at long lead times.

5 Materials and Methods

5.1 Data

We train neural networks to emulate the outputs of the Norwegian Earth System Model version 2 (NorESM2) [45] using the ClimateBench dataset [3]. For several shared socioeconomic pathways (SSPs) covering a wide range of future emissions scenarios [2], the ClimateBench dataset provides four key inputs to NorESM2, namely carbon dioxide (CO₂), methane (CH₄), sulfur dioxide (SO₂), and black carbon (BC). These inputs are given as annual means on a 96 latitude $\times$ 144 longitude global grid. We use these inputs to predict the surface air temperature (TAS) and diurnal temperature range (DTR) outputs of NorESM2 at the same 96 latitude $\times$ 144 longitude global grid spatiotemporal resolution.

5.2 Neural Network Model and Training

In this work, we use the best performing model analyzed in [3], a CNN-LSTM. This type of neural network combines a convolutional neural network (CNN) with a long short-term memory (LSTM) network, which respectively capture spatial and temporal relations in the data making them well suited for climate prediction.

For each of the temperature variables we wish to predict (TAS and DTR) and $\alpha$ value explored ($0,0.01,0.05,0.1,0.25,0.5,0.75,$ and $1$), we train an ensemble of 15 CNN-LSTMs on data provided by ClimateBench for a total of $2\times 8\times 15=240$ neural networks. The training data includes NorESM2’s historical experiment as well as its output for the SSP126, SSP370, and SSP585 experiments of the Scenario Model Intercomparison Project (ScenarioMIP) [2]. We also include the hist-GHG and hist-aer experiments from the Detection and Attribution Model Intercomparison Project (DAMIP) [46] in our training data. Each of the 3 SSP and 3 historical experiments contain data from 2015-2100 and 1850-2014 for a total of 753 training data points. From this training dataset, we reserve the first two years of data from every decade in each experiment as validation data. This is done because climate data is highly autocorrelated in time, so it best practice to form the validation data from continuous subsets of dataset rather than at random. Finally, to evaluate the CNN-LSTM’s performance, we use the SSP245 scenario as a holdout test dataset because it represents intermediate future anthropogenic forcing.

5.3 Data, Materials, and Software Availability

The complete codebase for processing the ClimateBench data, defining the equitable loss function, training the neural networks, and analyzing the results will be made available soon.