# CHEXANET: a novel approach to fast-tracking disequilibrium chemistry calculations for exoplanets using neural networks

Antonia Vojtekova, Ingo Waldmann, Kai Hou Yip, Bruno Merín, Ahmed Faris Al-Refaie, Olivia Venot

► To cite this version:

Antonia Vojtekova, Ingo Waldmann, Kai Hou Yip, Bruno Merín, Ahmed Faris Al-Refaie, et al. CHEXANET: a novel approach to fast-tracking disequilibrium chemistry calculations for exoplanets using neural networks. Monthly Notices of the Royal Astronomical Society, 2025, 538 (3), pp.1690-1719. 10.1093/mnras/staf297. hal-05381988

HAL Id: hal-05381988

https://hal.science/hal-05381988v1

Submitted on 25 Nov 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# CHEXANET: a novel approach to fast-tracking disequilibrium chemistry calculations for exoplanets using neural networks

Antonia Vojtekova<fnref n="1" />,<fnref n="1" />,<fnref n="2" />* Ingo Waldmann,<fnref n="1" /> Kai Hou Yip,<fnref n="1" /> Bruno Merín,<fnref n="2" /> Ahmed Faris Al-Refaie<fnref n="1" /> and Olivia Venot <fnref n="1" />,<fnref n="3" />

Accepted 2025 February 2. Received 2025 January 20; in original form 2024 June 6

## ABSTRACT

In the rapidly evolving field of exoplanetary science with missions like JWST and Ariel, there is a pressing need for fast and accurate simulations of disequilibrium chemistry in exoplanet atmospheres. Methods which assume chemical equilibrium simplify calculations but fail to capture the more complex chemical dynamics observed in actual exoplanets. Accurate estimations require complex kinetic codes that are time-intensive due to the need to solve ordinary differential equations. Given an extensive parameter space that needs to be explored to calculate forward models, kinetic codes will become a significant bottleneck, especially when used to match observations or in retrieval processes. We introduce CHEXANET, a novel U-Net-based neural network architecture designed to efficiently simulate disequilibrium chemistry in exoplanetary atmospheres. By utilizing the thermochemical equilibrium state of a hot-Jupiter's atmosphere as an initial condition, which computes in seconds, alongside a set of initial parameters such as the C/O ratio, the network effectively learns to predict the atmosphere in disequilibrium. It significantly enhances computational efficiency, reducing the prediction time for atmospheric disequilibrium states to just one second per atmosphere on a standard personal computer – over a hundred times faster than traditional kinetic models like Full and Reduced Exoplanet Chemical Kinetics distiLLed (FRECKLL).

**Key words:** methods: data analysis – exoplanets – planets and satellites: atmospheres.

## 1 INTRODUCTION

The discovery of exoplanets has exponentially expanded our understanding of the Universe, challenging our conceptions of planetary systems and their potential to maintain life. Examining exoplanetary atmospheres offers significant insights into the underlying physical and chemical processes that regulate these remote worlds (Tinetti et al. 2007; Haynes et al. 2015; Madhusudhan et al. 2016; Nikolov et al. 2018; Brogi & Line 2019; Mikal-Evans et al. 2020; Swain et al. 2021; Changeat et al. 2022; Matcheva & Castillo 2022; Edwards et al. 2023; Louca & Miguel 2023; Changeat et al. 2024). Traditional spectroscopic methods have been instrumental in detecting atmospheric signatures, revealing the presence of molecules such as water vapour, methane, and carbon dioxide (Louca & Miguel 2023). However, the intrinsic complexity of exoplanetary spectra and the limitations imposed by current observational technologies necessitate advanced analytical approaches (Matcheva & Castillo 2022).

In exoplanet atmosphere chemistry, assuming a chemical equilibrium state is often necessary due to computational demand, data resolution, and volume limitations. Equilibrium chemistry simplifies atmospheric modelling by assuming that reactions progress until forward and reverse rates balance, resulting in a stable composition

dictated by temperature, pressure, and elemental abundances. A system reaches thermochemical equilibrium when its Gibbs energy is minimal.

Hotter planets better fit the equilibrium assumption due to their high temperatures driving rapid chemical reactions that push the atmospheric composition towards thermochemical equilibrium (Heng & Tsai 2016). While computationally convenient, equilibrium state of atmosphere calculations often oversimplify the complex chemical processes in exoplanetary atmospheres, highlighting the need for more complicated disequilibrium models (Knutson et al. 2012; Venot et al. 2018; Roudier et al. 2021a).

These more complex models include atmospheric processes such as photochemistry and vertical and horizontal mixing due to atmospheric movement or quenching (Drummond et al. 2020; Baeyens et al. 2021; Moses et al. 2022; Zamyatina et al. 2023). Advancements in computational models, such as VULCAN (Tsai et al. 2021; Polman et al. 2023; Tsai et al. 2024) and FRECKLL (Al-Refaie et al. 2022; Veillet et al. 2024), have greatly enhanced our ability to simulate complex chemical networks across various exoplanetary environments. This understanding is crucial for analysing dynamic processes that impact atmospheric compositions, climate variability, and exoplanet habitability.

Studies like Roudier et al. (2021b) discuss the importance of such models. They highlight how disequilibrium chemistry impacts the detectability and composition of critical molecular species in

<footnote n="1">¹Department of Physics and Astronomy, University College London, Gower Street, WC1E 6BT London, UK</footnote>

<footnote n="2">²European Space Agency, ESAC, Camino Bajo del Castillo, Villanueva de la Cañada, E-28692 Madrid, Spain</footnote>

<footnote n="3">³CNRS, Université Paris Cité and Univ Paris Est Creteil, LISA, F-75013 Paris, France</footnote>

<footnote>* E-mail: vojtekova.antonia@gmail.com</footnote>
many exoplanet atmospheres, thereby playing a pivotal role in the observable characteristics of these distant worlds. Recent JWST studies of WASP-39b underscore the significance of disequilibrium processes in understanding exoplanetary atmospheres. Discoveries, such as the detection of SO₂, reveal complex chemical interactions and highlight the dynamic nature of planetary atmospheres beyond equilibrium assumptions (Ahrer et al. 2023; Rustamkulov 2023). Although kinetic models can fit data from state-of-the-art missions such as JWST (Tsai et al. 2023) or future ones like Ariel, calculation speed will become a significant bottleneck in disequilibrium atmospheric retrievals. The challenge arises from the necessity to explore an extensive parameter space and the slow computation of disequilibrium forward models. A single simulation using the FRECKLL model typically requires an average of 2.5 min; thus, handling 100 000 spectra for a single retrieval would necessitate approximately 173 d on a single CPU or nearly 3 d using 64 cores. Enhancing the speed and efficiency of these calculations without sacrificing accuracy is crucial for keeping pace with the rapid growth in exoplanet discoveries and the increasing complexity of observational data. At this point, exploring innovative methods for calculating disequilibrium states becomes imperative, ensuring we are well prepared for future challenges. One promising avenue is the utilization of machine learning, particularly neural networks.

Machine learning (ML) has emerged as a valuable tool for handling the vast and complex data sets characteristic of astronomical research (Lintott et al. 2008; Domínguez-Sánchez et al. 2018; Pasquet et al. 2019; Vojtekova et al. 2020; Yip et al. 2021; Kruk et al. 2023). Moreover, recent advancements in ML techniques, including deep learning and neural networks, have shown promise in enhancing the detection of atmospheric components, deciphering the subtle signals embedded in exoplanetary data, offering new perspectives on atmospheric chemistry and composition, improving direct imaging (Yip et al. 2020), atmospheric retrievals (Zingales & Waldmann 2018; Vasist et al. 2023; Changeat & Yip 2023; Ardévol Martínez et al. 2024; Yip et al. 2024), or acceleration of disequilibrium chemistry calculations (Hendrix, Louca & Miguel 2023).

This study delves into the innovative application of advanced neural network architectures, focusing on our novel CHEXANET (CHEmistry of EXoplanet Atmospheres NETwork). CHEXANET employs a U-Net-based architecture, which enhances the model's ability to simulate disequilibrium state computations and accelerate atmospheric retrieval processes. CHEXANET aims to bridge the gap between the increasing complexity of atmospheric data and the need for rapid, accurate modelling tools. As integrating artificial intelligence into exoplanet exploration remains a young field, there is a pressing need to initiate a dialogue to establish standards for such research. To this end, we delve into various numerical and visual evaluation techniques to comprehensively assess advanced deep neural network architectures' capabilities.

We present CHEXANET as a proof-of-concept study that these U-Net architectures are well suited for the modelling of exoplanet chemistry. For modelling simplicity and as discussed later in the paper, our model is trained only on isothermal atmospheres. This assumption can be lifted using additional transfer learning on additional data, using CHEXANET presented here as foundation model.

In the subsequent section (Section 2), we detail the methodologies employed for data generation, data preprocessing for machine learning applications, and selected network architecture. Additionally, this section meticulously outlines the steps undertaken to evaluate our models, ensuring a comprehensive understanding of our approach. Section 3 presents results accompanied by various visualization

techniques. Section 4 discusses patterns that emerged in the previous section and looks for connections between trained neural networks and phenomena happening in the atmosphere due to chemical reactions. Final conclusions can be found in Section 5.

## 2 METHODS

The Methods section describes our systematic approach to modelling exoplanetary atmospheric chemistry, as shown in Fig. 1. We define initial planetary and stellar parameters, feeding them into the TauREx ACE chemistry code (Agúndez et al. 2014, 2020) to generate equilibrium atmospheric samples. Conventionally, the next step involves using the kinetic code FRECKLL to produce disequilibrium steady-states of atmospheres. Our alternative is to analyse and pre-process equilibrium data and reach the disequilibrium steady-state of the atmosphere via CHEXANET. This approach allows for a more efficient transition from thermochemical equilibrium to disequilibrium steady-state, bypassing the computationally intensive steps traditionally required.

### 2.1 Data

#### 2.1.1 Data simulation

This section outlines our approach to generating a data set of exoplanet chemistries in thermochemical equilibrium and disequilibrium steady-state, see Fig. 2 as examples. In this study, we focus on warm Neptunes and hot Jupiters and draw our planetary system and chemistry parameters from a uniform distribution defined as follows:

(i) **Planet parameters** – mass, radius [Predicted by Forecaster (Chen & Kipping 2017) using mass], isothermal temperature of the atmosphere, and semi-major axis of the orbit.

(ii) **Atmospheric parameters** – carbon to oxygen ratio, metallicity, number of atmospheric layers, pressure grid.

(iii) **Star parameters** – radius, mass, type, temperature.

The pressure level in all simulations ranges from $1 \times 10^{-5}$ bars to $1 \times 10^1$ bars, maintaining an isothermal temperature profile across the atmosphere's hundred layers. The simulations are conducted using 1D models. Table 1 summarizes the selection and ranges of initial planetary parameters, utilizing data generation techniques from the Ariel Data Challenge (Changeat & Yip 2023).

**Equilibrium chemistry:** The second step consists of calculating the volume mixing ratios for species based on initial parameters, leading to a data set that results in an atmosphere in thermochemical equilibrium. This paper leverages TauREx 3 ACE (Agúndez et al. 2012, 2020) chemistry code for atmospheric equilibrium chemistry modelling, minimizing Gibbs free energy. This principle ensures that the system reaches a state of maximum stability and minimum energy, a condition where the atmosphere's chemical composition is in equilibrium. The model contains 108 C/H/O/N-bearing chemical species from the full Venot2020 chemical network (Venot et al. 2020) and three condensation species: H₂Oc, CH₄c, and NH₃c. The network contains a total of 1906 reactions (948 reversible and 10 irreversible ones). In this study, we follow species notation as defined in (Venot et al. 2020). Furthermore, the TauREx 3 framework (Al-Refaie et al. 2019, 2020) offers dynamic and extendable capabilities for spectral retrievals, providing a comprehensive suite for atmospheric modelling. In addition to the ACE chemistry code, the framework supports alternative chemical equilibrium solutions such as FastChem (Stock et al. 2018) and GGchem (Woitke et al. 2018).
<figure><img src="image_2.png" /><figcaption>**Figure 1.** The schematics of project workflow encapsulate parts of the projects into sections mentioned in the flowchart. The dashed rectangle named Traditional workflow denotes our current state-of-the-art pathway of generating disequilibrium steady-states of atmospheres. The dotted rectangle and arrow mark alternative workflow using a convolutional neural network is presented in this study. Finally, the paper cited in the flowchart refers to the subsequent paper, which provides a more in-depth analysis of the data used.</figcaption></figure>

<figure><img src="image_3.png" /><figcaption>**Figure 2.** Example of two atmospheric samples (index A3 and A6 in the test set) with different initial parameter values. The plots show the atmospheres in equilibrium (dashed line) and disequilibrium (full line) states. The detailed atmospheric parameters are provided in Table 7.</figcaption></figure>

**Disequilibrium chemistry:** In the third step, an equilibrium state of the atmosphere is used as input for the Full and Reduced Exoplanet Chemical Kinetics distiLLed (FRECKLL) chemical kinetic model. The model evolves a system of differential equations representing reaction rates, molecular diffusion and vertical mixing until a steady state is reached. Here, the initial abundance is set to the ACE-derived equilibrium abundances (Al-Refaie et al. 2024). To address the stiffness in chemical kinetics equations and improve numerical stability, FRECKLL utilizes stiff ODE algorithms like DLSODES from the ODEPACK (Hindmarsh 1983) or VODE

(Brown, Byrne & Hindmarsh 1989), alongside the K-fold summation method, for precise molecular production and loss rate calculation.

Disequilibrium chemistry introduces a new parameter, the Eddy diffusion coefficient ($K_{zz}$), which describes the rate of vertical mixing along the atmosphere. For this initial study, we generate data sets with $K_{zz}$ fixed to $10^9 \text{ cm}^2\text{s}^{-1}$ (Venot et al. 2020; Al-Refaie et al. 2024), representing an atmosphere with strong vertical mixing. This choice provides a balance of reaction-dominated regions in the lower atmosphere and vertical mixing-dominated upper atmosphere, creating
<table><caption>**Table 1.** Initial parameters, input for the ACE code to simulate an atmosphere in equilibrium, were chosen based on uniform distributions, except for the planet radius, predicted by Forecaster (Chen & Kipping 2017) based on planet mass.</caption><thead><tr><th>Parameter</th><th>Range</th><th>Units/Notes</th></tr></thead><tbody><tr><td>C/O Ratio</td><td>Uniform(0.4–1.5)</td><td>–</td></tr><tr><td>Metallicity</td><td>Uniform(0.5–100)</td><td>–</td></tr><tr><td>Planet mass</td><td>Uniform(0.3–10)</td><td>M<sub>J</sub> (Jupiter masses)</td></tr><tr><td>Planet temperature</td><td>Uniform(1100–2000)</td><td>K (Kelvin)</td></tr></tbody></table>

a difference between the final abundances from their equilibrium counterparts.

This detailed approach enables the study of time-dependent chemical variations, which is essential for understanding atmospheric compositions beyond thermochemical equilibrium.

The data set comprises 50 000 pairs of input (atmosphere in thermochemical equilibrium) and output (atmosphere at steady state) data. This data set was randomly divided into training (80 per cent), validation (10 per cent), and test (10 per cent) sets. The training set is used to train a selected machine learning algorithm, while the validation set is employed during training to monitor and validate the algorithm's progress. The test set is reserved for final evaluation, ensuring that it remains unseen during training. The test set represents an independent evaluation of the network's performance once training has been completed. This extensive data set is crucial for training neural networks, enabling them to learn the complex patterns associated with the transition from equilibrium to disequilibrium states in planetary atmospheres. Initially, the data included 111 species, refined to 64 of the most prominent species, as detailed in the pre-processing section (Section 2.1.2). This refinement was essential to focus our analysis on those species that significantly contribute to atmospheric chemistry and are potentially detectable, ensuring that our model remains computationally efficient and scientifically relevant. We also reduced the number of atmospheric layers to 96, optimizing the input size for the neural network architecture. The entire data set generation took approximately two weeks using one Apple M2 CPU, with the disequilibrium modelling process taking about 3 to 5 min for a hundred-layer atmosphere to achieve a steady state.

### 2.1.2 Data pre-processing

The initial data set spans more than 50 orders of magnitude in dynamic range (i.e. anywhere from 0.90 to $1 \times 10^{-50}$ of volume mixing ratio or VMR). This range requires pre-processing in multiple steps. We start with a logarithm transformation [log(VMR)] and continue with normalization. However, this step is complex due to the limited VMR ranges of some species and the large spans of others.

Hendrix et al. (2023) suggested species-specific normalization, considering each species-specific mean and standard deviation. Additionally, the same scaling parameters derived from the input properties are used to normalize the output mixing ratios. However, this method was unsuitable for our data set due to significant differences in the statistical properties between the input and output data and because species-specific normalization was not ideal for our objectives. Instead, we normalized all molecules using the same parameters to highlight the abundance differences between species. This approach aligns with our network architecture, where cross-species comparisons are critical.

We experimented with various normalization techniques and concluded that normalization would be problematic without cutting some species from the chemical network. As a lower threshold, we selected log (VMR) < -20. This threshold is far below the detection limit of JWST and Ariel Tier-3 for even highly spectrally active molecules such as methane ($CH_4$) and water (Changeat et al. 2020).

All molecules below log(VMR) = -20 were cut unless:

(i) If no more than 10 per cent of its values fall below log (VMR) = -20. All values below this threshold were clipped.

(ii) If no more than 20 per cent is below log (VMR) = -25 and no more than 10 per cent is below log (VMR) = -30, we retain the molecule and clip values below the threshold.

This pragmatic approach to data pre-processing, particularly the setting of thresholds based on the detectability and relevance of species, ensures that our neural network model, CHEXANET, is trained on data that are both scientifically meaningful and reflective of observations that could realistically be made with current and near-future astronomical instruments. By adopting a percentile-based criterion for data inclusion, we effectively balance the need for comprehensive chemical representation against the practical considerations of model training and performance. Afterwards, the chemical network is scaled to a range from -1 to 1 [minimum set to log (VMR) = -20, maximum set to log(VMR) = 0, i.e. 100 per cent abundance]. This scaling facilitates effective neural network learning by standardizing input data ranges, thereby enhancing the model's ability to discern and learn from the underlying patterns in the data.

## 2.2 Neural network

This section provides a comprehensive description of the neural network architecture used in this study. Additional technical details are available in Appendix A for readers seeking further information. We provide a brief summary below and refer the interested reader to the appendix for additional details.

In this paper, we explore applying a U-Net-based neural network (Ronneberger, Fischer & Brox 2015) architecture to simulate disequilibrium chemistry in exoplanetary atmospheres. We started with the established U-Net architecture and made several modifications to enhance its performance. U-Nets were first introduced as a highly efficient and accurate technique for biomedical image segmentation. Their ability to retain accurate spatial information makes this network architecture an interesting candidate for the prediction of chemical abundance profiles. In this application, the network takes the equilibrium state of an atmosphere as input and outputs the disequilibrium steady state. In order to improve prediction performances, we experimented with the addition of attention layers to the network, aiming to improve the network's focus on relevant features during training. We provide details about the attention mechanism in Section 2.2.2.

Additionally, we incorporated an additional network to process auxiliary planetary parameters – such as carbon-to-oxygen ratio, mean atmospheric temperature, metallicity, planet mass, and planet radius. This information was then fed into our primary equilibrium to disequilibrium model to further refine the accuracy of predictions. The final network, called Model D or CHEXANET, integrates a U-Net architecture enhanced with attention layers and an additional network for additional input parameters. It is optimized with a modified mean absolute error (MAE) loss function and trained with a small batch size of 8 input-output pairs per training iteration. A summary of models evaluated in this study is found in Table 2.
<table><caption>**Table 2.** Comparative analysis of model configurations, illustrating the presence of network enhancements, attention mechanisms, additional information network integration, and their respective batch sizes.</caption><thead><tr><td>Model</td><td>Attention</td><td>Add info</td><td>Batch size</td></tr></thead><tbody><tr><td>A</td><td>-</td><td>-</td><td>16</td></tr><tr><td>B</td><td>✓</td><td>-</td><td>16</td></tr><tr><td>C</td><td>✓</td><td>✓</td><td>16</td></tr><tr><td>D</td><td>✓</td><td>✓</td><td>8</td></tr></tbody></table>

### 2.2.1 U-Net architecture

The U-Net architecture was first introduced by Ronneberger et al. (2015). The U-Net architecture has shown versatility in iris (D’Alessandro et al. 2023) and has been fine-tuned for improved medical imaging (Younisse, Ghnemat & Al Saraireh 2023). In recent astronomical research, the U-Net architecture has been applied to complex tasks such as cosmic ray rejection, demonstrating its effectiveness in enhancing image processing (Vojtekova et al. 2020; Zhang & Bloom 2020).

A U-Net is a fully convolutional neural network without dense layers, designed primarily for image segmentation. This design choice, favouring convolutional over dense layers, optimizes U-Nets for handling spatial hierarchies of features, making it particularly effective for the detailed analysis required in atmospheric modelling. The architecture has a U-shaped encoder-decoder structure, and the most significant difference compared to typical autoencoders is skip connections, enabling them to capture features effectively. These skip connections are crucial for preserving spatial information throughout the network, facilitating the detailed reconstruction of atmospheric phenomena from the processed data. Skip connections concatenate feature maps from earlier to later layers, improving information and gradient flow throughout the network, as seen in Fig. 3. In this context, a feature map refers to the output of a convolutional layer, representing different aspects or features of the input data that the network has learned to recognize.

To understand the architecture of the used network, we recommend examining both the code and Figures simultaneously. As depicted in Fig. 3, the encoder and decoder segments each consist of four distinct steps organized into separate blocks. The input to the network, as seen in Fig. 3, is an atmospheric sample in an equilibrium state, which contains 96 pressure layers and 64 molecules (as selected based on criteria discussed earlier). The input dimensions are 64 × 96 × 1, indicating that the input is a 64 × 96 grid with a single channel.

**Encoder:** The encoder compresses the input data into a lower-dimensional representation, capturing the most relevant features while reducing redundancy and enabling more efficient processing by subsequent layers. Fig. 4 illustrates the encoder section of the U-net, called a convolutional module. This convolutional module comprises three consecutive operations: convolution, activation, and max-pooling layer. The convolution layer learns patterns in the feature maps, capturing spatial hierarchies and features. The activation function adds non-linearity, allowing the model to understand complex patterns. Our default activation function is LeakyReLU (Leaky Rectified Linear Unit, Maas et al. (2013)), chosen for its ability to maintain gradient flow in deeper network layers by preventing the vanishing gradient problem and ensuring robust learning across the network. The LeakyReLU activation function is given by:

$$f(x) = \begin{cases} x & \text{if } x \ge 0 \\ ax & \text{if } x < 0 \end{cases} \qquad (1)$$

<figure><img src="image_4.png" /><figcaption>**Figure 3.** A schematic representation of the U-Net architecture. The network contains a skip connection, denoted in the schematic as X2. In the architecture, separate steps are Convolutional (asymmetrically dashed red boxes) and Attention Upsample and Concatenate (symmetrically dashed blue boxes) modules, further explained in Figs 4 and 5. In each box, each module's output and dimension are visualized. Path K supersedes Path G when the Additional Info model is added (Models C—E). Additionally, the cursive text indicates the names of the separate layers in the code, making it easier for the reader to match the visualization with the corresponding code.</figcaption></figure>

where ‘a’ is a small constant value set to 0.2. If we changed ‘a’ to zero, the activation function would become rectified linear unit (ReLU) activation function (Nair & Hinton 2010).

Max-pooling reduces the input size, making computation more efficient and helping the model recognize features regardless of their position. This convolution, activation, and max-pooling sequence, denoted as a convolutional module, is repeated four times in the encoder, progressively reducing the input’s dimensionality and retaining the most relevant features. Specifically, the max-pooling layer decreases the input’s spatial size (width and height), effectively compressing the feature maps while preserving important information. Additionally, the output from the second activation layer in each convolutional module is saved and used later in the network for skip connections. These skip connections are crucial in the U-net architecture as they allow the model to reuse features from the encoder in the decoder, enhancing the reconstruction of the output. This is a major advantage of U-Nets over more simplistic autoencoder architectures.
<figure><img src="image_5.png" /><figcaption>**Figure 4.** Schematic of the encoder segment of the U-Net, specifically detailing the Convolutional module. It illustrates the sequence of operations within the module, including two convolution and activation operations followed by a max-pooling layer. The module's output is the result of the max-pool layer and the outcome of the previous convolutional layer. In the figure, *W* represents the width, *H* represents the height, and *F* represents the number of feature maps. The number of input feature maps (*F*₀) and output feature maps (*F*) differ, with *F* being a selected hyperparameter. The convolutional step uses a 3 × 3 kernel size.</figcaption></figure>

<figure><img src="image_6.png" /><figcaption>**Figure 5.** Visualization of the Upsample and Concatenate Block within the decoder portion of the U-Net architecture. This figure also introduces the alternative pathway (Path M) the attention module employs. Except for Model A, all other networks follow Path M (dashed line).</figcaption></figure>

<figure><img src="image_7.png" /><figcaption>**Figure 6.** Schematics of the attention module's structure and operational mechanism. The module selectively emphasizes features by assigning higher importance to elements that contribute most to the model's output. The inputs to the attention module, labelled X and G, represent the main input (X) and the guiding context (G), respectively. This distinction allows the module to focus attention on key features in X using the contextual information provided by G. In this attention module, *XL* and *GL* represent the latent, processed versions of the inputs X and G, following convolutional transformations. *F* = 1 indicates that there is just one feature map.

dimensions – width, height, and a number of feature maps of the input X and output are identical.</figcaption></figure>

Following the encoder section, the deepest layer in the U-Net, as seen in Fig. 3, is called latent space. It is particularly noteworthy, as it represents the most abstract representation of the input data, distilled through successive convolution layers.

**Decoder:** The decoder upsamples the feature maps of the latent space via transpose convolution operations, effectively restoring the lost spatial information and producing high-resolution maps. Like the encoder, each step within the decoder follows an established pattern for U-net architecture, in this study named 'Upsampling and Concatenation', as shown in Fig. 5 – Path J. The module involves a transpose convolution for spatial upsampling, followed by the concatenation of feature maps from the corresponding encoder layer to preserve high-level information. Subsequently, this concatenated output is processed through two sets of group convolution and activation layers, ensuring the refined reconstruction of the image's details. Alternatively, incorporating an attention module, as Path M depicts in Fig. 5, which offers a focused enhancement to the U-Net architecture. This addition enables the network to prioritize and refine its focus on areas of particular interest or complexity within the atmospheric data. In practice, attention mechanisms highlight regions in the input data that contribute most to the output, providing insights into the network's decision-making process. A follow on study (Vojtekova et al. in preparation) will examine the attention module as part of neural network analysis and interpretability. The attention module is described in the subsequent section.

### 2.2.2 Attention mechanism

The Attention U-Net architecture (Oktay et al. 2018; Vaswani et al. 2017) adds self-attention mechanisms to the U-Net model. The network can assess the significance of different spatial locations during feature extraction and segmentation. The primary objective of the attention block is to integrate feature maps obtained from both the encoder and decoder sections. The integration process ensures that the network effectively pays more attention to the most meaningful details in the data.

The composition of the attention module is illustrated in Fig. 6, and Fig. 5 shows its integration into the U-Net network – following Path M.

The attention module receives inputs from two sources: a feature map from the encoder (denoted as Conv1/2/3) with larger spatial dimensions for width and height (e.g. Conv3 16 × 16 × 128) and a corresponding feature map from the decoder with smaller dimensions (e.g. Conv4 8 × 8 × 256). The processing and integration of these feature maps allow the network to prioritize areas of interest, enhancing the model's predictive accuracy by focusing on regions within the atmospheric data most indicative of out-of-equilibrium states.

The module's inner workings are as follows: a convolutional layer with ReLU activation is applied to the feature map from the encoder. Notably, this convolutional layer uses a stride of two, effectively reducing the spatial dimensions of the input feature map by half to match the dimensionality of the feature map from the decoder. This step is essential for aligning the dimensionality of the encoder feature map with that of the decoder feature map, facilitating effective integration and processing. The second step is to apply a convolutional layer with ReLU activation to the feature map from the decoder. This convolutional layer uses a stride of one, which preserves the spatial dimensions of the input feature map. After combining feature maps through element-wise addition, the ReLU activation function is applied to create an activation feature map.
Next, the activation feature map is input for another convolution layer with a sigmoid activation function. The output is upsampled using the Keras UpSampling2D function with nearest neighbour interpolation to align with the encoder's feature map dimensions, acting as a weighting map. Subsequently, this map weights the encoder's feature map through element-wise multiplication. The weighted feature map feeds into the upsampling process, similar to the traditional U-Net approach.

### 2.2.3 Network for additional information input

A more compact network was designed to enhance the model's performance, incorporating additional parameters. This network integrates crucial planetary parameters directly into the modelling process, offering a nuanced approach to predicting atmospheric steady-state. This network utilizes normalized values for the carbon-to-oxygen ratio, temperature, metallicity, planet radius, and planet mass.

The architecture of this network is simple; the model contains three dense layers with LeakyReLU activation, and the last layer is subsequently reshaped to the same size as the bottom layer of the U-net (namely the latent space), as seen in the Fig. 7. We have found the choice of three dense layers to balance model complexity and computational efficiency. It offers sufficient capacity to learn complex relationships without overfitting or leading to excessive training times.

During a preliminary exploration of our training data set using Principal Component Analysis (PCA<fnref n="1" />) amongst other techniques, we encountered a number of correlations and data structures between auxiliary data and the datas' equilibrium and disequilibrium states. Whilst we reserve a detailed discussion of these correlations for an upcoming publication(Vojtekova et al., in preparation), we none the less provide a brief summary of our findings here as justification for including said auxiliary parameters in our modelling.

In our analysis, we were primarily interested in the relationship between principal components and initial parameters to identify which initial parameters significantly influence the data and therefore may help with data reconstruction in the network. Our results indicated that the C/O ratio and temperature show a strong linear correlation in equilibrium and disequilibrium chemistry data, as expected. On the other hand, the relationship between metallicity and principal components is more complex, as it does not correlate strongly with a single principal component like temperature and C/O ratio do. In the equilibrium state, planet mass and radius show almost no significant correlation with the principal components. However, planet mass and radius exhibit higher and statistically significant correlations with multiple principal components in the disequilibrium steady state, reflecting their increased influence on the system under these conditions. The above mentioned insight underscores the importance of these parameters in determining atmospheric composition, justifying their inclusion for enhancing model performance.

Similar to including planetary parameters, we considered adding PCA components of the input data to planetary parameters. The assumption being that a representation of these correlations may aid the network to learn these more explicitly. However, we found only negligible improvements. Given the additional network complexity

<figure><img src="image_8.png" /><figcaption>**Figure 7.** The network architecture for additional initial parameters data – planet radius, planet mass, carbon-to-oxygen ratio, isothermal temperature, and metallicity. The network input is normalized to the range of minus one to one, with all values considered linearly. Dense refers to a fully connected or dense layer, with the number following 'dense' indicating the number of filters at that layer. Afterwards, the network output is reshaped with standard Keras library tools to match the dimensionality of the latent space layer in the U-Net network.</figcaption></figure>

this approach would incur, we chose not to include principle components of the training data in our network.

## 2.3 Loss functions

The default selected loss function for experiments presented in this paper is mean absolute error (MAE). It calculates the average magnitude of errors between predicted and actual values. The MAE loss can be expressed as

$$MAE = \frac{1}{nLS} \sum_{i=1}^{n} \sum_{l=1}^{L} \sum_{s=1}^{S} |y_{ils} - \hat{y}_{ils}| \quad (2)$$

where n is the number of samples, L is the number of layers, and S is the number of species. $y_i$ and $\hat{y}_i$ represent the ground truth and predicted values by the network, respectively, in units of normalized log (VMR). The MAE loss is a commonly used metric for regression models. Its advantages lie in its simplicity and interpretability, which are well tested in the ML community (Goodfellow, Bengio & Courville 2016). However, a significant disadvantage, especially relevant in atmospheric studies, is its inability to effectively account for localized values. In cases where low-value predictions may have a high percentage error but not a high absolute error, MAE can overlook substantial relative discrepancies, as localized variations may average out over the whole data set. This can be especially problematic when slight variations have significant implications. Despite these limitations, we found the MAE to be the most robust loss function investigated.

### 2.3.1 Network design and hyperparameters

In order to determine the optimal configuration for our network, we systematically evaluated a wide range of hyperparameters, culminating in the selection of the network architectures discussed in this paper. The initial phase involved establishing a baseline for comparative analysis, utilizing a simplified version of the U-Net architecture with the following hyperparameters to be optimized:

(i) LeakyReLU activation functions

(ii) Mean absolute error loss function

(iii) Adam optimizer

(iv) Batch size 16

<footnote n="1">¹PCA is a statistical technique that transforms data into linearly uncorrelated variables, known as principal components.</footnote>
The selection of the adaptive moment estimation optimizer, or ADAM, (Kingma & Ba 2014) reflects a deliberate strategy to enhance the network's training process for both speed and accuracy. In simple terms, ADAM helps the network learn faster and more accurately by automatically adjusting how much the model's parameters are updated during training. This optimizer has become the standard go-to choice for deep learning applications.

In this study, we tested four variations of network architecture to explore the effects of different design choices on performance. We start with a minimal architecture and progressively include additional complexity to explore the performance impact of additional features (see Table 2).

**Model A:** is the architecturally most simplistic model. The flow through the simple network follows Path G and Path J, indicated in Fig. 3 and Fig. 5, respectively. The batch size is set to 16.

**Model B:** additionally includes an attention module discussed in Section 2.2.2 to Model A. Instead of following Path J, the network flow diverges to Path M. All other parameters remain the same.

**Model C:** adds the additional information input therefore employing Path K. The information provided for this network consists of the initial planet parameters. Parameters include carbon-to-oxygen ratio, planet temperature, metallicity, and planet radius and mass. Module C builds Model B in all other aspects.

**Model D:** Same as Model C but with an optimized batch size of 8 during training. A small batch size ensures a more granular gradient computation during training, leading to better final performances.

## 2.4 Metrics

We constructed a series of evaluation metrics to test the networks' performance. These metrics are pivotal for rigorously assessing the model's ability to accurately simulate atmospheric compositions, providing a comprehensive picture of its predictive capabilities and areas for improvement. These metrics include testing the network's output and outlier distributions, global errors, and error examination per molecule and layer. All metrics are calculated on the normalized data.

**KL-divergence:** We compared the distribution between network prediction and the ground truth via the Kullback-Leibler or KL divergence. The KL divergence measures the difference between two probability distributions over the same variable. It measures the amount of information lost when approximating one distribution with another. It is crucial to note that KL divergence cannot be used to measure one data set in isolation but rather as a comparative tool between predicted outputs and observational or simulated ground truth over the full training/test set. A lower KL divergence indicates higher similarity between the predicted and actual distributions when comparing multiple predicted outputs with an actual distribution. A value of zero indicates identical distributions. We should note that KL divergence is an improper distance due to its asymmetry; therefore, the KL divergence of P from Q is not the same as that of Q from P. This asymmetry necessitates careful consideration when interpreting the results, ensuring comparisons are consistently made with respect to the ground truth. We define the KL-divergence as

$$D_{KL}(P \parallel Q) = \sum_{x \in \mathcal{X}} P(x) \log \left( \frac{P(x)}{Q(x)} \right) \quad (3)$$

where $P(x)$ is the probability distribution of the ground truth data, $Q(x)$ is the probability distribution of the network predictions, and $x \in \mathcal{X}$ represents the individual outcomes or values in the data set of disequilibrium chemistry values.

**Error metrics:** We calculate multiple prediction versus ground-truth error measures. These are simple error distribution (i.e. ground truth minus network output), mean absolute error (MAE), mean absolute percentage error (MAPE), mean squared error (MSE), and standard deviation (STD). MAE averages absolute errors for a direct accuracy measure, whilst MSE emphasizes larger errors by squaring them, making it more sensitive to outliers. The primary distinction between MAE and MAPE is incorporating the actual value as a scaling factor in MAPE, offering a percentage-based measure of prediction accuracy relative to the actual values. MAPE is calculated as follows:

$$MAPE = \frac{100 \text{ per cent}}{nLS} \sum_{i=1}^{n} \sum_{l=1}^{L} \sum_{s=1}^{S} \left| \frac{\mathbf{y}_{ils} - \hat{\mathbf{y}}_{ils}}{\mathbf{y}_{ils}} \right| \quad (4)$$

MAPE quantifies the proportion of test atmospheres that fall below a specified error threshold. On the other hand, MAE and MSE provide insights into the average magnitude of errors across all test atmospheres.

**Outliers:** The purpose of looking for outliers determined by the metrics above is to look for potential patterns in the data and identify problematic input parameter regions. Outliers are identified through two distinct methods: (1) Data points above the 95th percentile of MAE; (2) leveraging KL divergence to measure differences between network outputs and ground truth, with significant deviations identified via the interquartile range (IQR) method. KL divergence outliers are identified as those exceeding Q3 (the value below which 75 per cent of the data points fall) plus 1.5 times the IQR.

**Visual:** Additionally, we visually compared the distribution between network prediction (network output) and the ground truth for the best, worst, and average network prediction to monitor the network improvement. Visual comparison is a supplementary method to numerical metrics, offering intuitive insights into the network's performance and the quality of its predictions.

## 2.5 Training and computational details

The CHEXANET (Model D) was trained for 60 epochs with 2000 steps per epoch and an 8-batch size. The training took approximately 4 h on an NVIDIA A100-SXM4-40GB GPU. Once trained, the model generates one atmospheric sample prediction on a standard laptop in approximately one second.

## 2.6 Application of CHEXANET to WASP-39b

To evaluate our model using more realistic data, we tested it with a temperature-pressure (T-P) profile that differs from the isothermal profile used during training. This approach allowed us to assess the network's performance when applied to conditions outside its training range.

We simulated the chemistry of WASP-39b (Faedi et al. 2011) and adopted the T-P profile from Khorshid et al. (2024) using an N-point profile (Al-Refaie et al. 2021) shown in Fig. 8. We adapted the profile to fit our pressure range (black dashed line in Fig. 8), and we refer to this adjusted profile as the WASP-39b target profile. The N-point profile defines a temperature profile at various atmospheric altitudes, requiring temperatures at both the surface and the top of the atmosphere, with the option of adding intermediate temperature and pressure points. Following Khorshid et al. (2024), we selected a radius of 1.27 $R_J$, a Mass of 0.28 $M_J$ (rounded to 0.3 due to parameter space restriction), and an orbital radius of 0.486 au. Table 3 summarizes all simulation parameters.
<figure><img src="image_9.png" /><figcaption>**Figure 8.** Temperature-pressure profile comparison for the exoplanet WASP-39b. The plot illustrates three different models: an isothermal profile at 1670K (Wasp-39b I), a modified N-point model (Wasp-39b II) and the N-point model from Khorshid et al. (2024) (Wasp-39b).</figcaption></figure>

<table><caption>**Table 3.** Summary of parameters used for Wasp-39b. The maximum temperature is used as input for the network for additional information.</caption><thead><tr><td>Parameter</td><td>Value</td></tr></thead><tbody><tr><td>Planet radius</td><td>1.27 R<sub>I</sub></td></tr><tr><td>Planet mass</td><td>0.30 M<sub>I</sub></td></tr><tr><td>C/O ratio</td><td>0.48</td></tr><tr><td>Metallicity</td><td>1.32</td></tr><tr><td>Selected isothermal temperature</td><td>1670K</td></tr><tr><td>Temperature range</td><td>760K—1670 K</td></tr><tr><td>Pressure range</td><td>10<sup>-5</sup> bar—10<sup>1</sup> bar</td></tr></tbody></table>

<figure><img src="image_10.png" /><figcaption>**Figure 9.** Comparison of disequilibrium steady-state profiles for various chemical species in the atmosphere of WASP-39b, based on both isothermal and N-point temperature profiles.</figcaption></figure>

For comparison, we also simulated the isothermal temperature profile (Wasp-39b I) and modified the N-point profile (Wasp-39b II). The Wasp-39b I has an isothermal temperature set at the maximum temperature of the N-point T-P profile at 1670 K, while the Wasp-39b II has a temperature range from 1200 to 1670 K. The remaining parameters are configured according to the values listed in Table 3. Fig. 9 compares the disequilibrium steady-state in the case of the isothermal T-P profile and N-point profile. It is essential to note that

some parameters, such as temperature, are outside the range on which we trained the model. Furthermore, in the Network for Additional Information, the input for temperature is just one number, not an array of numbers, as in the case of the N-point profile. Therefore, we use the maximum temperature of the atmosphere as an input. Introducing temperature values outside the training range leads to data drift (Humphrey et al. 2022; Simoes et al. 2024), where the input feature distribution changes. Additionally, the non-isothermal profile causes concept drift (Žliobaitė, Pechenizkiy & Gama 2016; Yip et al. 2022), altering the relationship between the inputs and outputs

## 3 RESULTS

This section details the various performances of the models discussed in the previous section. Overall, we find that CHEXANET (Model D) can accurately and robustly represent disequilibrium chemistry within its training range for the large majority of cases. In the results discussion below, we focus on two separate but interrelated aspects of the model: (1) Overall performance and (2) systematic trends and patterns. The overall performance, measured by error estimates across the whole data set and model, provides us with a good quantitative understanding of the network's capabilities and robustness. Here, we find Model D to be highly robust when measured against the (previously unseen) test data set. For the second aspect, we are interested in a more qualitative understanding of the underlying systematic errors contained within our models. These are of interest as they elucidate potential biases or regions of higher uncertainty in our overall disequilibrium chemistry predictions.

We note that despite its good performance, the model should only be seen to be valid within its original training range. We showcase this by discussing a non-isothermal data set of WASP-39b that clearly presents a 'concept drift' scenario to the isothermal-only training data. Here, we show that despite this drift scenario, CHEXANET does perform reasonably well. This leads us to believe that isothermally trained, or in other words, simpler models, can be used as foundation models. Foundation models can be subsequently augmented by additional training on, say, non-isothermal T-P profiles.

### 3.1 Overall performance

Understanding prediction errors and their causes is crucial. Our study uses various metrics to evaluate model performance in different atmospheric conditions. As outlined in Section 2.4, the initial phase of our network evaluation involved computing errors using various metrics – MAE, MSE, MAPE, STD, and the absolute difference – across specific atmospheres or the entire test set. This is supplemented with various forms of data visualizations. All metrics reported in this paper were calculated using data normalized to the range of –1 to 1, and as such, the results are presented without units. Below, we will discuss the overall performances of the individual models described in the previous section.

Table 4 summarizes the various metrics for all models. As can be seen, Model D demonstrates a substantial improvement over Model A across all metrics with reduced MAE and MSE values, demonstrating the effectiveness of implemented modifications. The performance gap between Model A and Model B (with an added attention mechanism) is minimal, suggesting that U-Net's skip connections may limit the benefits of attention layers. The most notable improvement is observed between Models B and C, where adding the network for additional parameters increases accuracy.
<table><caption>**Table 4.** Comparative performance metrics of neural network architectures: This table presents the highest and lowest mean absolute error (MAE) values for individual atmospheres within the test sample set. Additionally, it provides the average MAE, mean squared error (MSE), and standard deviation (STD) across the entire test set for each network. Lastly, the 95th percentile of MAE is shown, serving as one of the criteria for identifying outliers, with samples above this cutoff considered outliers. The lowest values in each column have been highlighted in bold to facilitate quick comparison. The first subtable presents values calculated on a normalized log (VMR), scaled from negative one to one, while the second subtable shows values in a log (VMR). Conversion to dex in log (VMR) follows a linear transformation with a factor of 10, except for MSE calculations, where the factor is 100.</caption><thead><tr><th>Model Name</th><th>Highest MAE</th><th>Lowest MAE</th><th>MAE for Test set</th><th>MSE for Test set</th><th>STD for Test set</th><th>95th Percentile MAE</th></tr></thead><tbody><tr><td>Model A</td><td>1.79 &middot; 10<sup>-1</sup></td><td>9.85 &middot; 10<sup>-4</sup></td><td>1.70 &middot; 10<sup>-2</sup></td><td>2.49 &middot; 10<sup>-3</sup></td><td>4.98 &middot; 10<sup>-2</sup></td><td>6.79 &middot; 10<sup>-2</sup></td></tr><tr><td>Model B</td><td>1.79 &middot; 10<sup>-1</sup></td><td>8.08 &middot; 10<sup>-4</sup></td><td>1.68 &middot; 10<sup>-2</sup></td><td>2.49 &middot; 10<sup>-3</sup></td><td>4.98 &middot; 10<sup>-2</sup></td><td>6.72 &middot; 10<sup>-2</sup></td></tr><tr><td>Model C</td><td>4.35 &middot; 10<sup>-2</sup></td><td>8.17 &middot; 10<sup>-4</sup></td><td>3.00 &middot; 10<sup>-3</sup></td><td>3.82 &middot; 10<sup>-5</sup></td><td>6.17 &middot; 10<sup>-3</sup></td><td>6.61 &middot; 10<sup>-3</sup></td></tr><tr><td>Model D</td><td><strong>3.56 &middot; 10<sup>-2</sup></strong></td><td><strong>5.61 &middot; 10<sup>-4</sup></strong></td><td><strong>2.11 &middot; 10<sup>-3</sup></strong></td><td><strong>2.17 &middot; 10<sup>-5</sup></strong></td><td><strong>4.65 &middot; 10<sup>-3</sup></strong></td><td><strong>4.80 &middot; 10<sup>-3</sup></strong></td></tr><tr><td colspan="7" style="text-align: center;">log (VMR)</td></tr><tr><td>Model A</td><td>1.79 &middot; 10<sup>0</sup></td><td>9.85 &middot; 10<sup>-3</sup></td><td>1.70 &middot; 10<sup>-1</sup></td><td>2.49 &middot; 10<sup>-1</sup></td><td>4.98 &middot; 10<sup>-1</sup></td><td>6.79 &middot; 10<sup>-1</sup></td></tr><tr><td>Model B</td><td>1.79 &middot; 10<sup>0</sup></td><td>8.08 &middot; 10<sup>-3</sup></td><td>1.68 &middot; 10<sup>-1</sup></td><td>2.49 &middot; 10<sup>-1</sup></td><td>4.98 &middot; 10<sup>-1</sup></td><td>6.72 &middot; 10<sup>-1</sup></td></tr><tr><td>Model C</td><td>4.35 &middot; 10<sup>-1</sup></td><td>8.17 &middot; 10<sup>-3</sup></td><td>3.00 &middot; 10<sup>-2</sup></td><td>3.82 &middot; 10<sup>-3</sup></td><td>6.17 &middot; 10<sup>-2</sup></td><td>6.61 &middot; 10<sup>-2</sup></td></tr><tr><td>Model D</td><td><strong>3.56 &middot; 10<sup>-1</sup></strong></td><td><strong>5.61 &middot; 10<sup>-3</sup></strong></td><td><strong>2.11 &middot; 10<sup>-2</sup></strong></td><td><strong>2.17 &middot; 10<sup>-3</sup></strong></td><td><strong>4.65 &middot; 10<sup>-2</sup></strong></td><td><strong>4.80 &middot; 10<sup>-2</sup></strong></td></tr></tbody></table>

<figure><img src="image_11.png" /><figcaption>**Figure 10.** The figure shows the heat maps of the MAE for various atmospheric layers and specified molecules using Models A and D. The x-axis is composed of molecules, and the y-axis represents altitude layers. Colour represents MAE between ground truth and network-generated output in the test data set. The left plot illustrates the MAE for Model A, while the right plot details that for Model D. The data shown are normalized. Notably, both models exhibit an increase in error as atmospheric pressure decreases, underscoring the challenges of accurate prediction at lower pressure levels. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

Fig. 10 shows the mean absolute error per layer and molecule for all atmospheric samples. Model D consistently demonstrates lower MAE across most molecules and layers. The decrease in standard deviation and number of outliers identified by 95th percentile MAE

is also observable in Fig. 12, particularly for atmospheres within the temperature range around 1400 K.

Table 5 shows the proportion of samples with mean absolute percentage error (MAPE) exceeding 10 per cent and 50 per cent.
<table><caption>**Table 5.** Evaluation of network predictions with high percentage errors: The table lists the number of samples and corresponding percentages of MAPE exceeding 10 per cent and 50 per cent for each neural network architecture. These thresholds highlight instances with significant prediction errors. Bold figures denote the network with the lowest number of samples and percentages surpassing the specified MAPE thresholds, indicating a relatively superior performance in terms of prediction accuracy.</caption><thead><tr><th>Model Name</th><th>Samples above 10 per cent MAPE</th><th>Percentage above 10 per cent MAPE</th><th>Samples above 50 per cent MAPE</th><th>Percentage above 50 per cent MAPE</th></tr></thead><tbody><tr><td>Model A</td><td>2497</td><td>48.50</td><td>637</td><td>12.37</td></tr><tr><td>Model B</td><td>2442</td><td>47.44</td><td>633</td><td>12.30</td></tr><tr><td>Model C</td><td>375</td><td>7.28</td><td>38</td><td>0.74</td></tr><tr><td>Model D</td><td><b>223</b></td><td><b>4.33</b></td><td><b>25</b></td><td><b>0.49</b></td></tr></tbody></table>

<table><caption>**Table 6.** Statistical analysis of model predictions using Kullback-Leibler (KL) Divergence: This table displays the mean, highest, lowest, and number of outlier values for each network model's KL divergence.</caption><thead><tr><th>Model Name</th><th>Mean KL div</th><th>Highest KL div</th><th>Lowest KL div</th><th>KL div outliers</th></tr></thead><tbody><tr><td>Model A</td><td>3.37 · 10<sup>-4</sup></td><td>1.36 · 10<sup>-2</sup></td><td>3.68 · 10<sup>-7</sup></td><td>794</td></tr><tr><td>Model B</td><td>3.37 · 10<sup>-4</sup></td><td>1.35 · 10<sup>-2</sup></td><td>2.08 · 10<sup>-7</sup></td><td>813</td></tr><tr><td>Model C</td><td>6.58 · 10<sup>-6</sup></td><td>1.17 · 10<sup>-3</sup></td><td>1.81 · 10<sup>-7</sup></td><td><b>528</b></td></tr><tr><td>Model D</td><td><b>3.67 · 10<sup>-6</sup></b></td><td><b>8.92 · 10<sup>-4</sup></b></td><td><b>7.40 · 10<sup>-8</sup></b></td><td>580</td></tr></tbody></table>

<figure><img src="image_12.png" /><figcaption>**Figure 11.** Box plot comparison of Kullback-Leibler divergence across network configurations. This figure illustrates the variability and distribution of KL divergence values for each network, highlighting differences in performance related to information loss when approximating the ground truth distribution. The box plot's central line (red line) shows the median KL divergence for each network, while the box indicates the interquartile range (IQR; 25th–75th percentile), and the whiskers extend to the furthest non-outlier data points. Points beyond the whiskers represent outliers.</figcaption></figure>

<table><caption>**Table 7.** The selection of atmospheric samples from the test data set and their initial parameters. The selected atmospheres are used throughout this study as examples. Three atmospheres were selected randomly (A2, A4, A6); one represents the worst (A1), the best (A5), and the average-performing atmosphere (A3) for the best-performing model. The numerical results for atmospheric selection are presented in Table 8.</caption><thead><tr><th>Atmosphere</th><th>Index</th><th>C/O ratio</th><th>Temperature (K)</th><th>Metallicity</th></tr></thead><tbody><tr><td>A1</td><td>3423</td><td>1.182</td><td>1348</td><td>7.687</td></tr><tr><td>A2</td><td>32</td><td>0.959</td><td>1364</td><td>31.134</td></tr><tr><td>A3</td><td>3102</td><td>1.381</td><td>1295</td><td>54.080</td></tr><tr><td>A4</td><td>1291</td><td>0.626</td><td>1282</td><td>32.890</td></tr><tr><td>A5</td><td>2730</td><td>1.258</td><td>1892</td><td>56.103</td></tr><tr><td>A6</td><td>3037</td><td>0.610</td><td>1302</td><td>31.934</td></tr></tbody></table>

Model D significantly reduces the number of high-error predictions, indicating improved reliability and accuracy, even in challenging parameter spaces.

KL divergence measures the difference between two probability distributions over the same variable. It is crucial to understand that KL divergence is inherently a comparative measure; it quantifies the difference between two probability distributions, specifically comparing the disequilibrium data distribution predicted by the network to the ground truth data distribution. Table 6 presents a KL divergence comparison between trained models, including outliers identified by this method. The hyperparameter tuning reduced the number of outliers by more than a hundred out of 5100 atmospheric samples. The same trend is observable in Fig. 11. Model C's outcomes demonstrate notable enhancements compared to the basic Model A. The reduction in outliers through hyperparameter tuning demonstrates the model's enhanced robustness and improved generalization across varying atmospheric conditions.

While the overall performance metrics highlight the general effectiveness of Model D, a detailed evaluation of individual atmospheric samples provides deeper insights into the model's behaviour under specific conditions. The analysis of selected test atmospheres (Table 7) reveals how the model handles a range of challenging scenarios. Table 8 compares the performance of Models A and D across selected atmospheric samples using MAE, MSE, STD, and KL divergence. Model D consistently outperforms Model A, with significant reductions in MAE and MSE. For worst-case prediction, Atmosphere A1, Model D decreases MAE from 0.1366 to 0.0356, indicating improved predictive accuracy.

Expanding on our analysis, we employ various visualization techniques to systematically evaluate atmospheric samples, offering distinct insights into model performance. One such technique involves generating plots that compare ground truth versus predicted values separately for all molecules, enabling a detailed, model-specific assessment. Figs 12 and 13 show predictions for CH₄ in randomly selected atmospheres (labelled A2 and A4) with different initial compositions – one with a temperature of 1364 K and the other with 1283 K. The reconstructed values are, in both cases, better for Model D. Additionally, Fig. 14 shows a reconstruction of CH₄ across
<table><caption>**Table 8.** The results of the selected atmospheric samples provide deeper insight beyond the average values presented in Tables 4 and 6. The initial four values pertain to Model A, while the remaining five correspond to Model D. Column named Note is relevant to Model D. The first subtable presents values calculated on normalized log (VMR), scaled from negative one to one, while the second subtable shows values in log (VMR). The conversion to dex in log (VMR) follows a linear transformation with a factor of 10, except for MSE calculations, where the factor is 100.</caption><thead><tr><th rowspan="2">Atmosphere</th><th colspan="4">Model A</th><th colspan="4">Model D</th><th rowspan="2">Note</th></tr><tr><th>MAE</th><th>MSE</th><th>STD</th><th>KL div</th><th>MAE</th><th>MSE</th><th>STD</th><th>KL div</th></tr></thead><tbody><tr><td>A1</td><td>1.37 · 10<sup>-1</sup></td><td>4.92 · 10<sup>-2</sup></td><td>1.87 · 10<sup>-1</sup></td><td>7.45 · 10<sup>-3</sup></td><td>3.56 · 10<sup>-2</sup></td><td>5.17 · 10<sup>-3</sup></td><td>6.34 · 10<sup>-2</sup></td><td>8.92 · 10<sup>-4</sup></td><td>Highest MAE</td></tr><tr><td>A2</td><td>3.28 · 10<sup>-2</sup></td><td>2.04 · 10<sup>-3</sup></td><td>3.71 · 10<sup>-2</sup></td><td>2.17 · 10<sup>-4</sup></td><td>2.57 · 10<sup>-3</sup></td><td>1.30 · 10<sup>-5</sup></td><td>3.41 · 10<sup>-3</sup></td><td>2.02 · 10<sup>-6</sup></td><td>Random sample</td></tr><tr><td>A3</td><td>1.83 · 10<sup>-2</sup></td><td>8.03 · 10<sup>-4</sup></td><td>2.83 · 10<sup>-2</sup></td><td>1.28 · 10<sup>-4</sup></td><td>2.10 · 10<sup>-3</sup></td><td>8.40 · 10<sup>-6</sup></td><td>2.89 · 10<sup>-3</sup></td><td>1.41 · 10<sup>-6</sup></td><td>Average MAE</td></tr><tr><td>A4</td><td>2.39 · 10<sup>-2</sup></td><td>1.16 · 10<sup>-3</sup></td><td>3.37 · 10<sup>-2</sup></td><td>1.78 · 10<sup>-4</sup></td><td>1.62 · 10<sup>-3</sup></td><td>4.93 · 10<sup>-6</sup></td><td>2.22 · 10<sup>-3</sup></td><td>8.61 · 10<sup>-7</sup></td><td>Random sample</td></tr><tr><td>A5</td><td>1.02 · 10<sup>-3</sup></td><td>2.33 · 10<sup>-6</sup></td><td>1.52 · 10<sup>-3</sup></td><td>4.19 · 10<sup>-7</sup></td><td>5.61 · 10<sup>-4</sup></td><td>6.30 · 10<sup>-7</sup></td><td>7.93 · 10<sup>-4</sup></td><td>1.31 · 10<sup>-7</sup></td><td>Lowest MAE</td></tr><tr><td>A6</td><td>1.79 · 10<sup>-1</sup></td><td>1.00 · 10<sup>-1</sup></td><td>2.65 · 10<sup>-1</sup></td><td>1.36 · 10<sup>-2</sup></td><td>1.59 · 10<sup>-2</sup></td><td>6.91 · 10<sup>-4</sup></td><td>2.23 · 10<sup>-2</sup></td><td>1.15 · 10<sup>-4</sup></td><td>Random sample</td></tr><tr><td></td><td colspan="10">log(VMR)</td></tr><tr><td>A1</td><td>1.37 · 10<sup>0</sup></td><td>4.92 · 10<sup>0</sup></td><td>1.87 · 10<sup>0</sup></td><td>-</td><td>3.56 · 10<sup>-1</sup></td><td>5.17 · 10<sup>-1</sup></td><td>6.34 · 10<sup>-1</sup></td><td>-</td><td>Highest MAE</td></tr><tr><td>A2</td><td>3.28 · 10<sup>-1</sup></td><td>2.04 · 10<sup>-1</sup></td><td>3.71 · 10<sup>-1</sup></td><td>-</td><td>2.57 · 10<sup>-2</sup></td><td>1.30 · 10<sup>-3</sup></td><td>3.41 · 10<sup>-2</sup></td><td>-</td><td>Random sample</td></tr><tr><td>A3</td><td>1.83 · 10<sup>-1</sup></td><td>8.03 · 10<sup>-2</sup></td><td>2.83 · 10<sup>-1</sup></td><td>-</td><td>2.10 · 10<sup>-2</sup></td><td>8.40 · 10<sup>-4</sup></td><td>2.89 · 10<sup>-2</sup></td><td>-</td><td>Average MAE</td></tr><tr><td>A4</td><td>2.39 · 10<sup>-1</sup></td><td>1.16 · 10<sup>-1</sup></td><td>3.37 · 10<sup>-1</sup></td><td>-</td><td>1.62 · 10<sup>-2</sup></td><td>4.93 · 10<sup>-4</sup></td><td>2.22 · 10<sup>-2</sup></td><td>-</td><td>Random sample</td></tr><tr><td>A5</td><td>1.02 · 10<sup>-2</sup></td><td>2.33 · 10<sup>-4</sup></td><td>1.52 · 10<sup>-2</sup></td><td>-</td><td>5.61 · 10<sup>-3</sup></td><td>6.30 · 10<sup>-5</sup></td><td>7.93 · 10<sup>-3</sup></td><td>-</td><td>Lowest MAE</td></tr><tr><td>A6</td><td>1.79 · 10<sup>0</sup></td><td>1.00 · 10<sup>1</sup></td><td>2.65 · 10<sup>0</sup></td><td>-</td><td>1.59 · 10<sup>-1</sup></td><td>6.91 · 10<sup>-2</sup></td><td>2.23 · 10<sup>-1</sup></td><td>-</td><td>Random sample</td></tr></tbody></table>

<figure><img src="image_13.png" /><figcaption>**Figure 12.** Comparison of reconstruction on molecule CH₄ for Models A and D in random atmosphere (Atmosphere A2). The orange triangle, blue circle, and yellow star represent network input or equilibrium data, the output of Model A and Model D, respectively.</figcaption></figure>

This atmosphere has a temperature of 1364 K, identified as a region of an initial planet parameter space with the highest reconstruction error. The bottom plot shows sorted residuals. Model D demonstrates a significant reduction in the MSE and MAE alongside a lower standard deviation of residuals. Thus, we can determine that Model D's performance exceeds Model A's. All values presented in the figure are calculated on normalized log(VMR) values, scaled from negative one to one. Conversion to dex in log(VMR) follows a linear transformation with a factor of 10.

all samples in the test data set. The above-mentioned types of plots are variations of commonly used visualizations for similar data, such as those shown in Fig. 15, which displays the logarithm of the volume mixing ratio on the x-axis and the logarithm of pressure in bar on

<figure><img src="image_14.png" /><figcaption>**Figure 13.** Example of CH₄ molecule reconstruction in a random atmosphere (Atmosphere A4) by Models A and D. In the figure, the orange triangle denotes network input or equilibrium data, the blue circle represents the output of Model A, and the yellow star signifies the output of Model D. Compared with the atmosphere shown in Fig. 12, the temperature of this sample is 1283 K, which corresponds with parameter space with a lower reconstruction error. This comparison illustrates a notable enhancement in predictive precision by Model D, showcasing its improved reconstruction accuracy in atmospheres with inherently lower error propensities. All values presented in the figure are calculated on normalized log(VMR) values, scaled from negative one to one. Conversion to dex in log(VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

the y-axis for various species, comparing ground truth with network-predicted outputs. Similarly, Fig. 16 provides a comprehensive view of Model D's performance, including its worst, best, average, and random predictions.
<figure><img src="image_15.png" /><figcaption>**Figure 14.** Example of CH₄, CO₂, H₂O, C₂H₅ molecule reconstruction in test atmospheres (approximately 5150 samples) by Models A and D. The orange triangle, blue circle, and yellow star represent network input or equilibrium data, Model A's output, and Model D's output, respectively. This plot illustrates the disparity between the networks and the improvement done by hyperparameter tuning. It is apparent in the left lower region in the first plot that the reconstruction is better if the input data do not have high dispersion. Quantitative comparison is available in Tables B1 and B2. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

To deepen our understanding of model performance, we generated heat maps for select atmospheres to advance our evaluation process further, utilizing absolute error values to show discrepancies. As illustrated in Figs 17 and 18, the results show the error values per

layer and molecule for atmospheres A1 and A6. These visualizations highlight the differences in predictive performance for atmospheres A1 and A6. Model D shows a noticeable reduction in error magnitude across most layers and molecules compared to Model A. Notably,
<figure><img src="image_16.png" /><figcaption>**Figure 15.** Comparative visualization of predicted versus actual atmospheric compositions across different models. For atmospheres A2 and A4, the plots illustrate volume mixing ratios against pressure. Dashed lines represent neural network predictions, while solid lines indicate ground truth data. Below the main plots are residual plots between ground truth and network output, emphasizing the variability in model performance. This multi-type visualization approach highlights the challenges in discerning model distinction and underscores the necessity for diverse analytical perspectives.</figcaption></figure>

the error tends to increase with atmospheric layer altitude, reflecting the increased predictive complexity in upper layers.

## 3.2 Trends and patterns

Fig. 10 illustrates the mean absolute error distribution across atmospheric layers and molecules, providing an overview of model performance variations for different chemical species. Distinct error patterns for each molecule across the networks suggest that variations in model sensitivity to specific molecules may be affecting predictive accuracy. For specific molecules (e.g. N₂O, HNO, NO), we find

consistently higher errors throughout the data set. The increase in error amplitudes with decreasing atmospheric pressure indicates that the model struggles with low-pressure conditions, possibly due to the influence of outliers that exacerbate uncertainties in these specific molecules. This pattern, observed consistently across multiple error plots (similar to Fig. 10), highlights the impact of outliers on the overall error distribution and suggests that refining the model's handling of low-pressure conditions could improve its predictive accuracy.

Building on this observation, dicyanogen (C₂N₂) consistently emerged as the molecule with the highest error, whereas helium (He)
<figure><img src="image_17.png" /><figcaption>**Figure 16.** Model performance across varying MAE levels for the same network. The plot showcases volume mixing ratios versus atmospheric pressure for four distinct cases: the atmosphere with the highest MAE A1, the lowest MAE A5, an atmosphere with the average MAE A3, and one randomly selected atmosphere. Dashed lines represent network predictions, while solid lines indicate the ground truth.</figcaption></figure>

and hydrogen ($H_2$) exhibited the lowest. A detailed comparison of error metrics for each molecule in Models A and D are available in Tables B1 and B2. The variance in error rates among molecules may reflect their relative abundances in atmospheric layers or the complexity of their chemical interactions. It is plausible that molecules with narrower distributions tend to have lower errors, indicating that abundance distribution significantly impacts model performance. However, further analysis indicates that molecular distribution alone does not fully explain the observed error patterns.

To further understand these discrepancies, we conducted an outlier analysis to determine if certain atmospheric conditions consistently produce higher errors. An examination of outliers identified through

three different methods – KL divergence outliers, the 95th percentile of MAE (Table 4), and samples with MAPE above 50 per cent (Table 5) – revealed a consistent pattern: samples with temperature near 1400 K tend to have higher error. The observed trend can be further examined through a scatter plot, where temperature is plotted on the x-axis and mean absolute error (MAE) on the y-axis, as illustrated in Fig. 19. The observed correlation suggests that specific atmospheric reactions occurring near this temperature range may significantly impact predictive accuracy, underscoring the critical role of temperature on model performance. One possible interpretation of this phenomenon is the transition from CO to $CH_4$, occurring around 1400 K (Visscher, Lodders & Fegley 2010).
<figure><img src="image_18.png" /><figcaption>**Figure 17.** Heatmap of absolute error values across various molecular concentrations and atmospheric layers for atmosphere A1 utilizing Model A and D. Atmosphere A1, as shown in Table 8 is the worst performing sample for Model D. The colour gradient represents the error magnitude, with a noticeable trend of increasing error values at higher atmospheric layers. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

Given the strong influence of temperature on prediction errors, we also investigated how other planetary properties, such as mass and C/O ratio, correlate with error distribution, as shown in Fig. 20 for Model A and D. Model A, in particular, exhibits a strong correlation between errors and both planet mass and C/O ratio. Within the critical temperature range around 1400 K, the error rate is notably higher for low-mass planets, as illustrated in the middle and right-hand panels of Fig. 20. The MAE distribution in the middle panel shows a distinct bifurcation, revealing two distinct peaks that correspond to different planet masses.

In Model A, for higher C/O ratios at higher temperatures, we observe a more rapid decrease in error compared to lower C/O ratios. In Model D, errors are significantly smaller across all C/O ratios, suggesting that including C/O as an additional parameter effectively addresses this issue.

Although these trends, observed in Figs 20 and 19 are reduced by adding additional information into the network, as in the case of Model D, a small peak remains for low planet mass at lower temperatures, accompanied by an increase in error as the pressure decreases. This reduction is unsurprising, given that parameters such as C/O ratio, temperature, and metallicity are closely linked to the atmosphere's chemical composition.

To better understand these complex interactions, we visualized a pair plot (Figs 21 and 22) for initial planetary parameters, colour-coded using MAE. The pair of corner plots illustrate higher dimensional correlations of errors with parameters. As expected,

the Temperature versus C/O ratio correlations show a peak around a C/O ratio of one, which indicates the transitional regime carbon or oxygen-dominated atmosphere. In the Temperature versus Planet Mass plot, we observe a non-linear increase in MAE for planet masses between zero and four $M_J$, which peaks at higher masses and temperatures of 1400 K. Correspondingly, a trend is also observed in the Temperature versus Metallicity plot around 1400 K, supporting the earlier findings. These observations confirm the significant role of initial parameters in determining error patterns, and Model D shows clear improvements in mitigating these influences compared to Model A.

Moreover, when Figs 21 and 22 are compared, Model D shows a marked reduction in error patterns for Planet Mass versus C/O ratio and Planet Mass versus Metallicity. This improvement suggests that Model D is better at capturing the complex interactions between planetary parameters and atmospheric chemistry, leading to more accurate predictions across diverse planetary environments.

## 3.3 Results: WASP-39b

To evaluate the robustness of Model D under previously unseen conditions, we simulated *JWST* observations of WASP-39b, as described in Section 2. The tests address two key challenges: the use of a non-isothermal temperature-pressure (T-P) profile and variations in temperature. We compared an isothermal T-P profile (atmospheric sample WASP-39b I) with a non-isothermal N-point
<figure><img src="image_19.png" /><figcaption>**Figure 18.** Heatmap indicating the absolute errors for different molecules across the atmospheric layers for atmosphere A6 using Models A and D. This visualization highlights the error's correlation with atmospheric layer altitude, aligning with the trend observed in other atmospheres. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

profile (atmospheric sample WASP-39b II) to assess how these
differences affect Model D's performance. An additional source
of error arises from the use of data with mass and radius values
that do not adhere to the relationship defined by Chen & Kipping
(2017), which may contribute to increased uncertainty in Model D's
performance.

In the first case (atmospheric sample Wasp-39 I), we used the
same methodology as our training data set to generate the example,
using an isothermal temperature of 1670 K and the previously defined
planetary parameters of WASP-39b (3). In the case of atmospheric
sample WASP-39b II, the T-P profile is an N-point profile with a
shape similar to that of WASP-39b but falls within the temperature
range of the training sample. Table 9 summarizes the numerical
results.

The results of the atmospheric sample Wasp-39b I align with
the pattern we observed when evaluating Model D. In the case of
the atmospheric sample Wasp-39b II, we find evaluation metrics
slightly lower than Model D's worst prediction. The error is highest
in pressure below $1 \times 10^{-2}$ bar, as illustrated in the right upper plot of
Fig. 23. Although the results of numerical metrics are slightly lower
than those for the worst-predicted sample by Model D, they are still
significantly lower than errors typically found for Models A and B.
Lastly, for our target N-point T-P profile of the atmospheric sample
WASP-39b, the evaluation metrics are slightly higher than those
observed in the previous case. Fig. 23 illustrates the absolute error
per molecule, and Fig. 24 shows the log (VMR) against pressure.

Although the results for atmospheric sample Wasp39-b are worse
than Model D's worst prediction, this is unsurprising given that the
temperature is outside the training range and the model has not
been trained on non-isothermal T-P profiles. However, the results
do suggest a good degree of robustness to non-isothermal T-P
profiles. This indicates that transfer learning could potentially be
used to retrain the model on different T-P profiles without requiring
a large amount of additional data. Recent retrievals of the hot Jupiter
WASP-39 b typically report $log_{10}$(VMR) uncertainties on the order
of $\pm 0.3$ to $\pm 0.5$ dex for molecules such as H$_2$O, CO$_2$, and SO$_2$ (e.g.
Constantinou, Madhusudhan & Gandhi 2023; Niraula et al. 2023).
These uncertainties reflect the precision achievable with JWST-
quality data. In comparison, Model D exhibits a mean absolute
error (MAE) of 0.47 dex in $log_{10}$(VMR) on our synthetic WASP-
39 b data. Specifically, for pressures below $10^{-2}$ bar, the molecule-
specific MAEs for H$_2$O and CO$_2$ are approximately 0.2 dex each. At
$10^{-4}$ bar, the MAEs increase to 0.27 and 0.38 dex for H$_2$O and CO$_2$,
respectively.

4 DISCUSSION

This study demonstrated that a modified U-Net neural network, CHEXANET, can accurately reproduce the disequilibrium steady-state of exoplanetary atmospheres. Calculating the disequilibrium steady-state with CHEXANET takes just a few seconds on a desktop, significantly reducing the computational time
<figure><img src="image_20.png" /><figcaption>**Figure 19.** Scatter plots illustrating the relationship between MAE and temperature for Models A, C, and D. Each plot reveals a pronounced peak in MAE at approximately 1400 K for Model A, which suggests a significant trend in prediction accuracy that correlates with temperature variations. The y-axis is not shared; all three plots have their ranges. In contrast, Model D (light colour stars) is overplotted over Model A (dark colour circles) to highlight performance differences. All values presented in the figure are calculated on normalized log(VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

<figure><img src="image_21.png" /><figcaption>**Figure 20.** Comparative visualization of MAE against isothermal temperature for Models A and D, with the left plot coloured by carbon-to-oxygen (C/O) Ratio and the right plot by planet mass. The left plot suggests a trend of decreasing MAE with increasing C/O ratios, particularly noticeable at temperatures above 1500 K for Model A. Conversely, the middle plot reveals a complex relationship between MAE, planet mass, and temperature for Model A. For temperatures below 1500 K, a higher planet mass corresponds to a lower MAE. In contrast, for temperatures above 1500 K, a higher planet mass results in a higher MAE with decreasing pressure. For planets below $4M_J$ and temperatures below 1500 K, the error increases noticeably for pressures below 1 bar. The right-side plot shows the temperature-MAE-planet mass relationship for Model D. Although some trend for higher error for low-mass planets is preserved, we can observe overall improvement in predictive strength and base removal. These findings underscore the importance of adding planetary parameters as an additional input in influencing the model's predictive performance. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

required for such complex simulations. FRECKLL calculation takes two to four minutes per atmospheric sample. Although there is still room for improvement in the neural network model regarding predictive accuracy, the computational efficiency of CHEX-ANET outperforms traditional methods, offering a promising alternative for large-scale exoplanetary atmosphere studies in the future. This leap in computational speed opens the possibility of integrating more complex atmospheric phenomena in future simulations.

As in many machine learning applications, data pre-processing is crucial. We found that normalizing the data from -1 to 1 was more effective than the standard 0 to 1 range, likely due to the wide range of values and the need to capture minor differences more accurately. This choice improved model performance and accuracy.

After optimizing the data through pre-processing, we evaluated the model's performance by testing various architectural adjustments. In terms of different models, we examined the influence of attention layers, an additional network with initial parameters as input,
<figure><img src="image_22.png" /><figcaption>**Figure 21.** Pair plots illustrating the distribution of initial exoplanetary parameters: C/O ratio, Temperature (K), Metallicity, and Planet Mass (*M*<sub>J</sub>), colour-coded by the MAE for Model A. Notable is the cluster formation and a peak in MAE around a C/O ratio of 1.0 in the Temperature versus C/O ratio plot, indicating a possible region of increased predictive difficulty. Moreover, the error for the higher C/O ratio decreases, as seen in the upper right corner of the plot. Temperature versus planet Mass plot reveals a non-linear pattern. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

and batch sizes. The U-Net architecture has skipped connections integrated by default. Including attention layers did not significantly improve performance. Nevertheless, we decided to keep them in the architecture as they might enhance the model's interpretability by highlighting which features are most relevant for predicting disequilibrium states (Vojtekova et al., in preparation).

Adding an additional network, which takes input from initial exoplanetary parameters such as C/O ratio, metallicity, temperature, planet radius, and mass, proved crucial, leading to the most considerable improvement. This is not surprising and expected.

Our decision to lower the batch size from 16 to 8 was driven by the realization that atmospheric chemistry is a sensitive system, and small-scale variations are significant. Thus, averaging over larger batch sizes might obscure critical details. We considered lowering the

batch size to two, but this approach was computationally challenging due to the gradient calculations becoming too noisy. A potential strategy to circumvent this limitation could be to train the model with a higher batch size and, in the last 10 epochs, lower the batch size to two, thus fine-tuning more granular details towards the end of training. Following these architectural adjustments, a thorough evaluation was conducted to assess their impact on the model's overall performance and accuracy.

Our numerical evaluation clearly showed that these changes improve the overall predictive strength of the model. One recurrent problem, shown in Fig. 10, occurs in every trained network: distinct error patterns and underestimation of certain species. This observation suggests that the complexity of exoplanetary atmospheric chemistry might not be fully captured by the current model structure,
<figure><img src="image_23.png" /><figcaption>**Figure 22.** Pair plots illustrating the distribution of initial exoplanetary parameters: C/O ratio, Temperature (K), Metallicity, and Planet Mass ($M_J$), colour-coded by the MAE for Model D. Although the prediction performance has improved compared to Fig. 21, it is still noticeable that there is a cluster formation and a peak in MAE around a C/O ratio of 1.0 in the Temperature versus C/O ratio plot, which suggests a possible area of increased predictive difficulty. A slight pattern is also observed for temperatures around 1500 K and planets with a small radius under the temperature of 1500 K, which is also present in right Fig. 20. All values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

<table><caption>**Table 9.** Summary of numerical results for WASP-39b atmospheric models. Comparison of different atmospheric models of WASP-39b, highlighting variations in T-P profiles and pressure ranges and their effects on prediction accuracy. All values presented in the table are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log(VMR) follows a linear transformation with a factor of 10, except for MSE calculations, where the factor is 100.</caption><thead><tr><th>Atmospheric Sample</th><th>T-P Profile</th><th>Temperature</th><th>MAE</th><th>MSE</th><th>STD</th><th>KL div</th><th>MAPE</th></tr></thead><tbody><tr><td>Wasp-39b I</td><td>Isothermal</td><td>Within training range</td><td>9.29 · 10<sup>-3</sup></td><td>3.52 · 10<sup>-4</sup></td><td>1.78 · 10<sup>-2</sup></td><td>7.10 · 10<sup>-5</sup></td><td>6.11</td></tr><tr><td>Wasp-39b II</td><td>N-point</td><td>Within training range</td><td>3.24 · 10<sup>-2</sup></td><td>4.19 · 10<sup>-3</sup></td><td>5.96 · 10<sup>-2</sup></td><td>9.60 · 10<sup>-4</sup></td><td>8.19</td></tr><tr><td>Wasp-39b</td><td>N-point</td><td>Outside training range</td><td>4.70 · 10<sup>-2</sup></td><td>9.85 · 10<sup>-3</sup></td><td>9.29 · 10<sup>-2</sup></td><td>2.15 · 10<sup>-3</sup></td><td>10.56</td></tr></tbody></table>

indicating the need for further refinement to address these discrepancies. We note that only a detailed error analysis, such as the one presented here, allows us to identify these areas of uncertainty in the first place.

Analysis of outliers revealed a notable trend in Models A and B: regardless of the method used for identification (95th percentile of MAE, MAPE errors exceeding 50 per cent, or outliers in KL
<figure><img src="image_24.png" /><figcaption>**Figure 23.** Heatmaps compare the absolute error distributions across different atmospheric layers and molecular species for the WASP-39b I model (upper right), the WASP-39b II model (upper left), and the WASP-39b model (lower right). All values presented in the figure are calculated on normalized log(VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>

divergence), there is a consistent convergence towards increased outlier frequency for atmospheres with temperature around 1400 K. This phenomenon is possibly caused by the CO to CH₄ transition, happening near 1400 K for pressures below 1 bar, as seen in Fig. B3 (Visscher et al. 2010). At lower temperatures, methane (CH₄) is the preferred stable form of carbon in a hydrogen-rich environment. As the temperature increases, the chemical equilibrium shifts, favouring the formation of carbon monoxide (CO) due to its higher thermal stability at these temperatures. This shift begins to happen significantly at the transition temperature around 1400 K. However, the exact temperature can vary depending on other parameters, such as

the pressure and other chemical species in the atmosphere. Moreover, Figs 20, 21 and 22 demonstrate that the neural networks can detect correlations between temperature, C/O ratio, planet mass, and the prediction error. This indicates that the networks can learn how the initial exoplanet parameters relate to the resulting abundances of various species. While the C/O ratio more directly influences the chemistry of the atmosphere, we postulate that the planet mass plays a measurable role through its gravitational impact on vertical mixing and atmospheric stratification (Tsai et al. 2021; Al-Refaie et al. 2024).

Applying machine learning to astrophysical problems inevitably involves certain limitations. One of the primary challenges we faced
<figure><img src="image_25.png" /><figcaption>**Figure 24.** The plot displays volume mixing ratios against pressure for Wasp-39b with an N-point T–P profile. Dashed lines represent neural network predictions, while solid lines indicate ground truth data.</figcaption></figure>

was selecting the appropriate data set, fully aware that the model could not contain the entire parameter space. In this study, we decided to cover hot Neptunes and warm Jupiters. Our analysis considered isothermal temperature-pressure profiles and did not include clouds or photochemistry at this stage. To ensure a methodical development, we adopted a step-by-step approach, starting with a more constrained parameter space to closely monitor and understand the behaviour of our selected neural network, the U-Net. The phased plan, which remains adaptable as the project progresses, is outlined as follows:

(i) Initial testing: Implement the U-Net architecture using the data set defined in this study to establish a baseline performance (this study).

(ii) Evaluation techniques: Develop and apply robust evaluation techniques to assess model accuracy and reliability (this study).

(iii) Data pre-processing and exploration: summarize in-depth exploration of data (Vojtekova et al., in preparation).

(iv) Network behaviour analysis: Investigate the internal workings of the neural network, which will inform further refinements (Vojtekova et al., in preparation).

(v) Data set Expansion: Gradually expand the data set to include more diverse temperature-pressure profiles, a broader range of eddy diffusion coefficients ($K_{zz}$), additional pressure layers, and clouds, thereby increasing the model's applicability.

(vi) Application of insights: Apply the insights gained from the above-mentioned steps to refine the model, ensuring it can handle more complex atmospheric data.

(vii) Advanced model training: The approach for further training and refinement of the model is still under discussion as we evaluate the outcomes of the initial phases.

While this study has limitations, particularly regarding its applicability to real-world data, a structured, incremental approach is essential for adequately developing this tool. With this model, we successfully demonstrated that U-Nets are viable deep learning architectures to model disequilibrium chemistry. By carefully evalu-

ating each step and gradually building complexity, we aim to create a robust model that can eventually be applied to broader astrophysical contexts.

However, one of the ideas is to leverage its computational efficiency to narrow down the parameter space during retrievals and then use established tools like FRECKLL for final calculations when only a few options remain. Although the use of the trained model will remain limited, transfer learning could be applied to adapt the network to different sets of parameters. With transfer learning, we can adjust the model to new parameter ranges using a smaller data set, effectively fine-tuning the existing model. Additionally, we can modify the network by removing the initial layer, altering the input size to include more molecules, and retraining the model to adapt to these changes.

## 5 CONCLUSIONS

In conclusion, this study has successfully demonstrated the potential of using a modified U-Net neural network, CHEXANET, for fast and accurate simulations of exoplanetary atmospheric disequilibrium steady-states. Furthermore, CHEXANET's computational efficiency represents a substantial advance over traditional methods, reducing the prediction time for atmospheric disequilibrium states to just one second per atmosphere on a standard personal computer, compared to the two to four minutes required by established models. Our rigorous evaluation, employing a wide range of metrics and visualizations, establishes a benchmark for the field and highlights areas for further refinement. While the significant speed improvements offered by CHEXANET are crucial, the rigorous evaluation presents a greater challenge, necessitating the development of sophisticated evaluation tools to ensure the model's reliability and accuracy.

This work aims to contribute to the ongoing discussion in the exoplanet research community about the methodologies for conducting and evaluating studies in this rapidly evolving field. By sharing our findings and the detailed evaluation framework, we encourage further advancements and collaborative efforts in the field of exoplanetary atmospheres.

The study's implications include enhancing our understanding of how machine learning is learning from our data set, addressing a critical gap in exoplanetary research by reducing the computational time of determining the chemical composition at disequilibrium steady-state, and, therefore, making it feasible to conduct large-scale simulations that were previously too computationally demanding. This study is also an example of the benefits of cross-disciplinary research, combining expertise from computational science, machine learning, and astronomy.

Despite the above mentioned advancements in the application of machine learning in the field of exoplanets, it is necessary to remember this approach's downfalls. The problematic areas include low-abundance species, abundance in parts of the atmospheres with lower pressure and samples with initial isothermal temperatures around 1400 K. Moreover, depending on the training data set, the trained model works in a restricted parameter space. Lastly, interpreting neural networks is a complicated field of study; therefore, it is hard to fully understand the processes happening within. We aim to gain a deeper understanding of inner processes in the pending study by Vojtekova et al. (in preparation). To mitigate these challenges, future work will focus on refining the Model's architecture to improve its interpretability and expand its parameter space, ensuring broader applicability to diverse exoplanetary environments.

In the future, we would like to refine CHEXANET further by training it with a more diverse data set (e.g. with varying Eddy diffusion
parameters or different temperature-pressure profiles), employing diffusion training or modifying the loss according to findings from the current study. Moreover, insight from the field of chemistry could further improve the Model's predictive accuracy. Ideally, we would like to test the network on atmospheric retrievals on real-time data sets from space missions like Ariel or JWST. Lastly, releasing the data sets used for training and testing CHEXANET, including any simulated atmospheric data, and including the Model in a simple Python package is crucial to enhancing research transparency and reproducibility.

In summary, CHEXANET represents a significant leap forward in integrating machine learning with studying exoplanetary atmospheres. It offers computational efficiency and accuracy in simulating disequilibrium chemical steady-states and deep insights into sources of error that occur when using it.

## ACKNOWLEDGEMENTS

AV acknowledges funding from ESA via the OSIP platform with contract number 4000139240. OV acknowledges funding from the ANR project 'EXACT' (ANR-21-CE49-0008-01) and from the Centre National d'Études Spatiales (CNES). This research received funding from the European Research Council (ERC) under the European Union's Horizon 2020 research and innovation programme (grant agreement no. 758892/ExoAI), and from the Science and Technology Facilities Council (STFC; grant no. ST/W00254X/1).

## DATA AVAILABILITY

The code developed and the data sets utilized in this study will be made fully open access upon final acceptance of the manuscript.

## REFERENCES

* Agúndez M., Venot O., Iro N., Selsis F., Hersant F., Hébrard E., Dobrijevic M., 2012, A&A, 548, A73

* Agúndez M., Parmentier V., Venot O., Hersant F., Selsis F., 2014, A&A, 564, A73

* Agúndez M., Martínez J. I., de Andres P. L., Cernicharo J., Martín-Gago J. A., 2020, A&A, 637, A59

* Ahrer E.-M. et al., 2023, Nature, 614, 653

* Al-Refaie A., Changeat Q., Waldmann I., Tinetti G., 2019, ApJ, 917, 23

* Al-Refaie A., Changeat Q., Venot O., Waldmann I., Tinetti G., 2020, Europlanet Science Congress, Vol. 14. Europlanet Science Congress. p. EPSC2020-669

* Al-Refaie A. F., Changeat Q., Waldmann I. P., Tinetti G., 2021, ApJ, 917, 37

* Al-Refaie A. F., Changeat Q., Venot O., Waldmann I. P., Tinetti G., 2022, ApJ, 932, 123

* Al-Refaie A. F., Venot O., Changeat Q., Edwards B., 2024, ApJ, 967, 132

* Ardévol Martínez F., Min M., Huppenkothen D., Kamp I., Palmer P. I., 2024, A&A, 681, L14

* Baeyens R., Decin L., Carone L., Venot O., Agúndez M., Mollière P., 2021, in European Planetary Science Congress, Vol. 15. p. EPSC2021-46

* Brogi M., Line M. R., 2019, AJ, 157, 114

* Brown P. N., Byrne G. D., Hindmarsh A. C., 1989, SIAM J. Sci. Stat. Comput., 10, 1038

* Changeat Q., Yip K. H., 2023, RAS Tech. Instrum., 2, 45

* Changeat Q., Al-Refaie A., Mugnai L. V., Edwards B., Waldmann I. P., Pascale E., Tinetti G., 2020, AJ, 160, 80

* Changeat Q. et al., 2022, ApJS, 260, 3

* Changeat Q. et al., 2024, ApJS, 270, 34

* Chen J., Kipping D., 2017, ApJ, 834, 17

* Constantinou S., Madhusudhan N., Gandhi S., 2023, ApJ, 943, L10

* D'Alessandro V. I., De Palma L., Attivissimo F., Di Nisio A., Lanzolla A. M. L., 2023, IEEE International Symposium on Medical Measurements and Applications. IEEE, 1-5

* Domínguez Sánchez H., Huertas-Company M., Bernardi M., Tuccillo D., Fischer J. L., 2018, MNRAS, 476, 3661

* Drummond B. et al., 2020, A&A, 636, A68

* Edwards B. et al., 2023, ApJS, 269, 31

* Faedi F. et al., 2011, A&A, 531, A40

* Goodfellow I., Bengio Y., Courville A., 2016, Deep Learning, MIT Press, Cambridge, MA

* Haynes K., Mandell A. M., Madhusudhan N., Deming D., Knutson H., 2015, ApJ, 806, 146

* Hendrix J. L. A. M., Louca A. J., Miguel Y., 2023, MNRAS, 524, 643

* Heng K., Tsai S.-M., 2016, ApJ, 829, 104

* Hindmarsh A. C., 1983, in Stepleman R. S., others, eds, IMACS Transactions on Scientific Computation Vol. 1, Scientific Computing. North-Holland, Amsterdam, p. 55

* Humphrey A., Kuberski W., Bialek J., Perrakis N., Cools W., Nuyttens N., Elakhrass H., Cunha P. A. C., 2022, MNRAS, 517, L116

* Khorshid N., Min M., Polman J., Waters L. B. F. M., 2024, A&A, 685, A64

* Kingma D. P., Ba J., 2014, Adam: A Method for Stochastic Optimization. ICLR, San Diego

* Knutson H. A. et al., 2012, ApJ, 754, 22

* Kruk S. et al., 2023, Nat. Astron., 7, 262

* Lintott C. et al., 2008, MNRAS, 389, 1179

* Hendrix Julius, Louca Amy Miguel Yamila, 2023, MNRAS, 524, 643

* Maas A. L., Hannun A. Y., Ng A. Y., 2013, in Proc. icml. Vol. 30. p. 3

* Madhusudhan N., Agúndez M., Moses J. I., Hu Y., 2016, Space Sci. Rev., 205, 285

* Matcheva K. I., Castillo T. A., 2022, Planet. Sci. J., 3, 205

* Mikal-Evans T., Sing D. K., Kataria T., Wakeford H. R., Mayne N. J., Lewis N. K., Barstow J. K., Spake J. J., 2020, MNRAS, 496, 1638

* Moses J. I., Tremblin P., Venot O., Miguel Y., 2022, Exp. Astron., 53, 279

* Müller A., Guido S., 2016, Introduction to Machine Learning with Python: A Guide for Data Scientists, O'Reilly Media, Sebastopol, CA

* Nair V., Hinton G. E., 2010, in Fürnkranz J., Joachims T., eds, Proceedings of the 27th International Conference on Machine Learning. Omnipress, Madison. p. 807

* Nikolov N. et al., 2018, Nature, 557, 526

* Niraula P., de Wit J., Gordon I. E., Hargreaves R. J., Sousa-Silva C., 2023, ApJ, 950, L17

* Oktay O. et al., 2018, Attention U-Net: Learning Where to Look for the Pancreas. MIDL, Amsterdam, 03999

* Pasquet J., Bertin E., Treyer M., Arnouts S., Fouchez D., 2019, A&A, 621, A26

* Polman J., Waters L. B. F. M., Min M., Miguel Y., Khorshid N., 2023, A&A, 670, A161

* Ronneberger O., Fischer P., Brox T., 2015, U-Net: Convolutional Networks for Biomedical Image Segmentation. Springer International Publishing, Cham, 234

* Roudier G. M., Swain M. R., Gudipati M. S., West R. A., Estrela R., Zellem R. T., 2021a, AJ, 162, 37

* Roudier G., Swain M., Gudipati M., West R., Estrela R., Zellem R., 2021b, AJ, 162, 161

* Rustamkulov Z. e. a., 2023, Nature, 614, 659

* Simoes L. F., Casale P., Felismino M., Yip K. H., Waldmann I. P., Tinetti G., Lueftinger T., 2024, Proceedings of SPAICE2024: The First Joint European Space Agency / IAA Conference on AI in and for Space, 186

* Stock J. W., Kitzmann D., Patzer A. B. C., Sedlmayr E., 2018, MNRAS, 479, 865

* Swain M. R. et al., 2021, AJ, 161, 213

* Tinetti G. et al., 2007, Nature, 448, 169

* Tsai S.-M., Malik M., Kitzmann D., Lyons J. R., Fateev A., Lee E., Heng K., 2021, ApJ, 923, 264

* Tsai S.-M., Lee E. K. H., Powell D., et al., 2023, Nature, 617, 483

* Tsai S.-M. et al., 2024, ApJ, 963, 41
Vasist M., Rozet F., Absil O., Mollière P., Nasedkin E., Louppe G., 2023, A&A, 672, A147

Vaswani A., Shazeer N., Parmar N., Uszkoreit J., Jones L., Gomez A. N., Kaiser L., Polosukhin I., 2017, Attention Is All You Need, Advances in Neural Information Processing Systems, 30.

Veillet R., Venot O., Sirjean B., Bounaceur R., Glaude P. A., Al-Refaie A., Hébrard E., 2024, A&A, 682, A52

Venot O., Drummond B., Miguel Y., Waldmann I. P., Pascale E., Zingales T., 2018, Exp. Astron., 46, 101

Venot O., Cavalié T., Bounaceur R., Tremblin P., Brouillard L., Lhoussaine Ben Brahim R., 2020, A&A, 634, A78

Visscher C., Lodders K., Begley Bruce J., 2010, ApJ, 716, 1060

Vojtekova A., Lieu M., Valtchanov I., Altieri B., Old L., Chen Q., Hroch F., 2020, MNRAS, 503, 3204

Woitke P., Helling C., Hunter G. H., Millard J. D., Turner G. E., Worters M., et al., 2018, A&A, 614, A1

Yip K. H. et al., 2020, in Brefeld U., Fromont E., Hotho A., Knobbe A., Maathuis M., Robardet C., eds, Machine Learning and Knowledge Discovery in Databases. Springer International Publishing, Cham, p. 322

Yip K. H., Changeat Q., Nikolaou N., Morvan M., Edwards B., Waldmann I. P., Tinetti G., 2021, AJ, 162, 195

Yip K. H. et al., 2022, in Ciccone M., Stolovitzky G., Albrecht J., eds, Proceedings of the NeurIPS 2022 Competitions Track. Vol. 220, Proceedings of Machine Learning Research PMLR, p. 1, https://proceedings.mlr.org/s/v220/yip23a.html

Yip K. H., Changeat Q., Al-Refaie A., Waldmann I. P., 2024, ApJ, 961, 30

Younisse R., Ghnemat R., Al Saraireh J., 2023, Int. J. Power Electron. Drive Syst., 13, 5406

Zamyatina M. et al., 2023, MNRAS, 519, 3129

Zhang K., Bloom J. S., 2020, ApJ, 889, 24

Zingales T., Waldmann I. P., 2018, AJ, 156, 268

Žliobaitė I., Pechenizkiy M., Gama J., 2016, Big Data Analysis: New Algorithms for a New Society, Springer International Publishing, Cham, 91

CNNs contain learnable parameters, or weights, organized into two-dimensional filters. Applying these filters to the input layer and subsequent hidden layers generates a set of feature maps. Here, hidden layers refer to any layer that is not the input or output layer. For instance, if a network contains three convolutional layers and two max-pooling layers, then it has five hidden layers.

A feature map is the output of a convolutional layer after filters are applied to the input. The number of filters specified by the user will result in an equivalent number of feature maps. In our figures, blue boxes labelled “Conv 3×3” or “Conv 1×1” (Figs 3 and 6) represent convolutional layers, where the numbers indicate the filter size – either 3×3 pixels or 1×1 pixels.

The convolutional layer can also reduce the spatial dimensions (width and height) of the input, depending on the stride chosen. The stride defines how many pixels the filter moves over the input image with each step. A stride of one preserves the input’s spatial dimensions throughout the network, except in the convolutional layer processing input X in the attention model, where the stride is set to two, producing output feature maps at half the input’s spatial size. For this study, we use the Conv2D layer from the Keras library.

**Dense layer:**

A dense or fully connected layer is a key component in many neural networks, including the Network for Additional Information (Fig. 7). The primary unit of a dense layer is the neuron. In a dense layer, each neuron is connected to all neurons in the previous layer, allowing for complex transformations of the input data. Each neuron $n$ connects through a weighted connection $w_n$, and the input $x_n$ is multiplied by this weight. The weighted inputs are then summed up:

$$z = \sum_{n} w_{n}x_{n} + b_{n}, \qquad (A1)$$

where $b_n$ is a bias term that improves the network’s performance. Once this summation is computed, the result $z$ is passed through an activation function, introducing non-linearity into the model as previously described. This study uses the Dense layer function from the Keras library. The Network for Additional Information, as discussed in Section 2.2.3 and illustrated in Fig. 7, consists of three dense layers with varying filter sizes and a LeakyReLU activation function.

**Training process:** After defining the network and pre-processing the data, the next step is to train the network, which involves setting several hyperparameters, including the loss function, batch size, and number of epochs. The loss function quantifies the difference between the network’s predicted output and the ground truth. In our work, the ground truth refers to the disequilibrium steady state of an atmosphere, serving as a benchmark for evaluating the accuracy of the network’s predictions during training.

Batch size is the number of training samples processed at once. In this context, it represents the number of atmospheres in an equilibrium state used as input to the network. Commonly, powers of two, like 4, 8, or 16, are chosen for batch size, balancing computational efficiency with model performance. Although a larger batch size can accelerate convergence, it may be computationally expensive and varies in optimal value depending on the data and model used.

Lastly, an epoch is one complete pass through the entire training data set, during which the model processes every batch once to update its weights based on the calculated loss. Alternatively, we can define a smaller subset of the data or specify how many batch-sized groups of samples will be processed in each epoch, allowing for more control over the training process.

## APPENDIX A: MACHINE LEARNING INSIGHTS

This section provides an overview of machine-learning concepts relevant to this study. As machine learning and neural networks are vast fields, we recommend the following resources for a deeper understanding: Introduction to Machine Learning with Python: A Guide for Data Scientists (Müller & Guido 2016) and Deep Learning (Goodfellow et al. 2016).

**Network input:**

The input to a neural network (NN) can be a vector or a matrix, depending on whether we are using a convolutional neural network (CNN) or a fully connected neural network (FNN). For matrix inputs, typical dimensions include width, height, and the number of channels. In the case of colour image data, for instance, the input matrix has three channels representing the red, green, and blue colour components.

In this work, the primary network, U-net, receives an input matrix representing the atmospheric equilibrium state. This matrix has a width of 64 (corresponding to 64 molecules), a height of 96 (corresponding to pressure layers), and a single channel. The input for the Additional Information Network is a vector with dimensions 5×1, representing five initial parameters. The input size is fixed for these neural network architectures, and all input data is uniformly normalized.

**Convolution layer:**

Convolutional neural networks (CNNs) are designed to process matrix-like data, typically represented as a three-dimensional matrix. These networks use a specialized mathematical operation called convolution.
The process of training a neural network can be summarized as an algorithm:

(i) Randomly initialize weights

(ii) Start of the epoch:

(a) Select a random batch of samples

(b) Feed the batch into the network (process through all network layers),

(c) Calculate the loss function,

(d) Update weights using backpropagation,

(e) Repeat steps (ii)(a)–(ii)(d) over the defined data subset.

(iii) Repeat step (ii) until the loss converges or the selected number of epochs is reached.

(iv) Evaluate the network.

Backpropagation [Step (ii)(d)] calculates the gradient of the loss function with respect to weights, which the optimization algorithm (e.g. stochastic gradient descent) uses to minimize loss and improve model performance. The training process stops either when epochs reach a preset number or convergence criteria are met. Keras callbacks manage convergence in the training process. The EarlyStopping monitors validation loss and stops training if there is no improvement after a defined number of epochs, preventing overfitting. The ReduceLROnPlateau reduces the learning rate by a specified factor (e.g. 0.1) when the validation loss plateaus for several consecutive epochs, helping in smoother convergence. Based on experimental results as determined by these convergence criteria, the training was set to stop after 60 epochs.

**Optimization algorithm:** In neural networks, an optimization algorithm iteratively adjusts the model's weights to minimize the loss function and enhance the model's performance during training. We selected Adaptive Moment Estimation (Adam) (Kingma & Ba 2014) as an optimization algorithm. Two optimization methods, RMSprop and Momentum, are combined to efficiently update model weights

during training. The optimizer employs an algorithm that calculates the first and second moments of the gradients, which is then used to adjust the learning rate for each parameter. This adaptive learning rate facilitates quicker convergence and improves the handling of noisy gradients. The algorithm of the Adam optimizer is as follows:

$$m_t = \beta_1 \cdot m_{t-1} + (1 - \beta_1) \cdot g_t \quad (A2)$$

$$v_t = \beta_2 \cdot v_{t-1} + (1 - \beta_2) \cdot g_t^2 \quad (A3)$$

$$\hat{m}_t = \frac{m_t}{1 - \beta_1^t} \quad (A4)$$

$$\hat{v}_t = \frac{v_t}{1 - \beta_2^t} \quad (A5)$$

$$\theta_t = \theta_{t-1} - \frac{\alpha}{\sqrt{\hat{v}_t} + \epsilon} \cdot \hat{m}_t \quad (A6)$$

Here, $m_t$ and $v_t$ represent the first and second moments of the gradients, respectively, and $g_t$ is the gradient at time $t$. The terms $\beta_1$ and $\beta_2$ are exponential decay rates, $\alpha$ is the predefined learning rate. $\theta_t$ represents the vector of parameters (weights) of the neural network model at the $t$th iteration. A small constant $\epsilon$ is added in the denominator of equation (A6) to prevent division by zero.

## APPENDIX B: THE ROLE OF ATMOSPHERIC MOLECULES IN ERROR ANALYSIS

This section provides the performance metrics of different molecules. Tables B2 and B2 detail the evaluation metrics – MAE, MSE, STD, MAPE, and KL Divergence – for Models A and D, offering a direct comparison of their predictive accuracy. Following these tables, Figs B1 and B2 further elaborate on the MAE for each molecule across Models A, C, and D. These plots are strategically ordered to highlight the molecules’ performance, from those with the highest average MAE (indicating lower predictive accuracy) to the lowest.
<table><caption>**Table B1.** Detailed evaluation statistics for Model A (baseline model) and Model D (the best performing model) across selected molecules. The table displays the evaluation metrics for 32 molecules as both Models A and D assessed, allowing for a comprehensive molecule-by-molecule performance analysis.</caption><thead><tr><th rowspan="2">Molecule</th><th colspan="5">Model A</th><th colspan="5">Model D</th></tr><tr><th>MAE</th><th>MSE</th><th>STD</th><th>MAPE</th><th>KL Div</th><th>MAE</th><th>MSE</th><th>STD</th><th>MAPE</th><th>KL Div</th></tr></thead><tbody><tr><td>HCNO</td><td>0.02776</td><td>0.00377</td><td>0.13942</td><td>3.52463</td><td>0.00136</td><td>0.00311</td><td>3.24e-05</td><td>0.14298</td><td>0.38077</td><td>1.22e-05</td></tr><tr><td>N2O</td><td>0.03189</td><td>0.00655</td><td>0.2521</td><td>10.55312</td><td>0.00176</td><td>0.00325</td><td>4.99e-05</td><td>0.25177</td><td>1.18217</td><td>1.39e-05</td></tr><tr><td>CH2CHO</td><td>0.01653</td><td>0.00196</td><td>0.25232</td><td>3.07313</td><td>5.48e-04</td><td>0.00228</td><td>2.01e-05</td><td>0.25447</td><td>0.57631</td><td>5.74e-06</td></tr><tr><td>CH3CO</td><td>0.01487</td><td>0.00183</td><td>0.27588</td><td>5.81994</td><td>4.72e-04</td><td>0.00217</td><td>1.77e-05</td><td>0.27766</td><td>4.83386</td><td>4.62e-06</td></tr><tr><td>NCO</td><td>0.01977</td><td>0.00163</td><td>0.15568</td><td>4.27801</td><td>3.59e-04</td><td>0.0024</td><td>1.95e-05</td><td>0.15549</td><td>0.49728</td><td>4.66e-06</td></tr><tr><td>CH3O</td><td>0.02205</td><td>0.00348</td><td>0.19782</td><td>3.4948</td><td>9.66e-04</td><td>0.0025</td><td>2.51e-05</td><td>0.20044</td><td>0.41192</td><td>7.17e-06</td></tr><tr><td>O2</td><td>0.01635</td><td>0.00242</td><td>0.33085</td><td>5.79237</td><td>7.39e-04</td><td>0.00249</td><td>2.99e-05</td><td>0.33262</td><td>1.84473</td><td>9.34e-06</td></tr><tr><td>CH3CHO</td><td>0.01585</td><td>0.00197</td><td>0.34773</td><td>20.79624</td><td>4.41e-04</td><td>0.00249</td><td>2.09e-05</td><td>0.34963</td><td>6.76706</td><td>4.67e-06</td></tr><tr><td>HNO</td><td>0.02772</td><td>0.0051</td><td>0.23385</td><td>7.04478</td><td>0.00124</td><td>0.00286</td><td>3.72e-05</td><td>0.23526</td><td>0.70763</td><td>9.42e-06</td></tr><tr><td>C</td><td>0.01529</td><td>0.00201</td><td>0.32318</td><td>3.89765</td><td>5.64e-04</td><td>0.00267</td><td>2.41e-05</td><td>0.32544</td><td>1.18404</td><td>7.44e-06</td></tr><tr><td>CHCO</td><td>0.02049</td><td>0.00305</td><td>0.24895</td><td>7.65297</td><td>6.38e-04</td><td>0.00244</td><td>2.53e-05</td><td>0.25189</td><td>2.1565</td><td>5.59e-06</td></tr><tr><td>CO2H</td><td>0.01377</td><td>0.00164</td><td>0.23955</td><td>3.49308</td><td>4.05e-04</td><td>0.00231</td><td>2.02e-05</td><td>0.24027</td><td>1.05292</td><td>5.05e-06</td></tr><tr><td>HOCN</td><td>0.01966</td><td>0.00171</td><td>0.14226</td><td>10.62551</td><td>2.89e-04</td><td>0.00236</td><td>1.89e-05</td><td>0.14411</td><td>2.55236</td><td>3.26e-06</td></tr><tr><td>C2H5</td><td>0.0231</td><td>0.00568</td><td>0.38329</td><td>71.4295</td><td>0.00114</td><td>0.00267</td><td>4.07e-05</td><td>0.38571</td><td>14.88022</td><td>8.52e-06</td></tr><tr><td>C2H</td><td>0.02744</td><td>0.00712</td><td>0.41229</td><td>41.73325</td><td>0.00149</td><td>0.00302</td><td>4.93e-05</td><td>0.41662</td><td>5.13871</td><td>1.08e-05</td></tr><tr><td>CH2OH</td><td>0.02162</td><td>0.00344</td><td>0.19975</td><td>6.27893</td><td>7.33e-04</td><td>0.00239</td><td>2.50e-05</td><td>0.2031</td><td>1.31175</td><td>5.42e-06</td></tr><tr><td>CH</td><td>0.0163</td><td>0.00223</td><td>0.29721</td><td>5.46459</td><td>5.47e-04</td><td>0.00234</td><td>2.04e-05</td><td>0.30005</td><td>1.53629</td><td>5.47e-06</td></tr><tr><td>C2H6</td><td>0.02656</td><td>0.0071</td><td>0.48282</td><td>54.61452</td><td>0.00119</td><td>0.0028</td><td>4.38e-05</td><td>0.48454</td><td>4.30949</td><td>7.50e-06</td></tr><tr><td>C2H3</td><td>0.0254</td><td>0.00692</td><td>0.38599</td><td>103.0582</td><td>0.00127</td><td>0.00263</td><td>4.24e-05</td><td>0.38951</td><td>16.21903</td><td>7.87e-06</td></tr><tr><td>CH2CO</td><td>0.021</td><td>0.00324</td><td>0.28273</td><td>70.90344</td><td>4.85e-04</td><td>0.00232</td><td>2.27e-05</td><td>0.28543</td><td>9.00265</td><td>3.39e-06</td></tr><tr><td>NNH</td><td>0.00219</td><td>1.49e-05</td><td>0.17622</td><td>1.52893</td><td>3.16e-06</td><td>0.00118</td><td>2.84e-06</td><td>0.1758</td><td>0.74425</td><td>5.78e-07</td></tr><tr><td>H2CN</td><td>0.01851</td><td>0.00115</td><td>0.2278</td><td>26.68494</td><td>2.10e-04</td><td>0.00222</td><td>1.42e-05</td><td>0.22796</td><td>5.57452</td><td>2.61e-06</td></tr><tr><td>CH3OH</td><td>0.02669</td><td>0.00451</td><td>0.31511</td><td>106.34728</td><td>7.18e-04</td><td>0.00255</td><td>3.11e-05</td><td>0.31647</td><td>9.70387</td><td>5.07e-06</td></tr><tr><td>N4S</td><td>0.01746</td><td>0.00141</td><td>0.26003</td><td>85.56737</td><td>2.29e-04</td><td>0.0022</td><td>1.79e-05</td><td>0.26069</td><td>8.81525</td><td>3.23e-06</td></tr><tr><td>N2D</td><td>0.01706</td><td>0.00137</td><td>0.28972</td><td>3.78749</td><td>3.54e-04</td><td>0.00244</td><td>2.04e-05</td><td>0.29086</td><td>0.50871</td><td>6.05e-06</td></tr><tr><td>CN</td><td>0.01837</td><td>0.00113</td><td>0.28193</td><td>43.11976</td><td>1.84e-04</td><td>0.00215</td><td>1.31e-05</td><td>0.28013</td><td>6.29934</td><td>2.21e-06</td></tr><tr><td>1CH2</td><td>0.01607</td><td>0.0023</td><td>0.23659</td><td>15.5487</td><td>4.48e-04</td><td>0.00188</td><td>1.59e-05</td><td>0.2402</td><td>5.06562</td><td>3.13e-06</td></tr><tr><td>HNCO</td><td>0.01945</td><td>0.00163</td><td>0.15177</td><td>87.10821</td><td>1.99e-04</td><td>0.00232</td><td>1.85e-05</td><td>0.15403</td><td>9.09737</td><td>2.28e-06</td></tr><tr><td>NO</td><td>0.02777</td><td>0.00524</td><td>0.25015</td><td>114.42283</td><td>7.90e-04</td><td>0.00271</td><td>4.54e-05</td><td>0.25299</td><td>10.37119</td><td>7.09e-06</td></tr><tr><td>O3P</td><td>0.01299</td><td>0.00169</td><td>0.29248</td><td>6.0511</td><td>3.44e-04</td><td>0.0019</td><td>1.70e-05</td><td>0.29526</td><td>1.89073</td><td>3.57e-06</td></tr><tr><td>O1D</td><td>0.00343</td><td>8.19e-05</td><td>0.22875</td><td>0.60008</td><td>3.24e-05</td><td>0.00203</td><td>1.22e-05</td><td>0.2284</td><td>0.38494</td><td>5.03e-06</td></tr><tr><td>C2H4</td><td>0.02693</td><td>0.00758</td><td>0.46928</td><td>106.79477</td><td>9.89e-04</td><td>0.00282</td><td>4.96e-05</td><td>0.47294</td><td>16.55994</td><td>6.53e-06</td></tr></tbody></table>
<table><caption>**Table B2.** Detailed evaluation statistics for Model A (Baseline Model) and Model D (Highest performing model) across selected molecules. The table displays the evaluation metrics for 32 molecules as both Model A and Model D assessed, allowing for a comprehensive molecule-by-molecule performance analysis. The molecules include different excited states of atoms, such as N2D, which refers to the first excited state of nitrogen; O3P, which refers to the fundamental state of oxygen atomic; and O1D, which refers to the first excited state of oxygen atomic.</caption><thead><tr><th rowspan="2">Molecule</th><th colspan="6">Model A</th><th colspan="6">Model D</th></tr><tr><th>MAE</th><th>MSE</th><th>STD</th><th>MAPE</th><th>KL Div</th><th>MAE</th><th>MSE</th><th>STD</th><th>MAPE</th><th>KL Div</th></tr></thead><tbody><tr><td>NH</td><td>0.01705</td><td>0.00144</td><td>0.16347</td><td>91.01526</td><td>2.05e-04</td><td>0.00202</td><td>1.58e-05</td><td>0.1646</td><td>10.20361</td><td>2.36e-06</td></tr><tr><td>3CH2</td><td>0.01605</td><td>0.0023</td><td>0.22651</td><td>62.62723</td><td>3.64e-04</td><td>0.00184</td><td>1.57e-05</td><td>0.23002</td><td>7.97111</td><td>2.49e-06</td></tr><tr><td>HCO</td><td>0.0051</td><td>1.47e-04</td><td>0.15953</td><td>22.92054</td><td>1.94e-05</td><td>0.00131</td><td>4.31e-06</td><td>0.15726</td><td>7.88401</td><td>5.75e-07</td></tr><tr><td>C2H2</td><td>0.02791</td><td>0.0077</td><td>0.41959</td><td>42.21477</td><td>8.12e-04</td><td>0.0026</td><td>4.58e-05</td><td>0.42336</td><td>3.93354</td><td>4.81e-06</td></tr><tr><td>H2CO</td><td>0.00837</td><td>4.82e-04</td><td>0.15102</td><td>33.71023</td><td>5.45e-05</td><td>0.00156</td><td>7.55e-06</td><td>0.15214</td><td>6.42954</td><td>8.54e-07</td></tr><tr><td>NH2</td><td>0.01695</td><td>0.00149</td><td>0.08379</td><td>105.40955</td><td>1.73e-04</td><td>0.00193</td><td>1.52e-05</td><td>0.08669</td><td>16.83793</td><td>1.81e-06</td></tr><tr><td>CO2</td><td>0.01373</td><td>0.00183</td><td>0.21701</td><td>50.81062</td><td>1.61e-04</td><td>0.00191</td><td>1.91e-05</td><td>0.21807</td><td>4.10611</td><td>1.67e-06</td></tr><tr><td>OH</td><td>0.01301</td><td>0.00177</td><td>0.1979</td><td>65.33476</td><td>2.20e-04</td><td>0.0017</td><td>1.63e-05</td><td>0.20101</td><td>10.00527</td><td>2.05e-06</td></tr><tr><td>CH3</td><td>0.01675</td><td>0.00242</td><td>0.22207</td><td>45.23229</td><td>2.54e-04</td><td>0.00172</td><td>1.58e-05</td><td>0.22495</td><td>5.60728</td><td>1.66e-06</td></tr><tr><td>HCN</td><td>0.01777</td><td>0.00109</td><td>0.17347</td><td>11.84062</td><td>9.51e-05</td><td>0.00194</td><td>1.24e-05</td><td>0.17538</td><td>1.21374</td><td>1.08e-06</td></tr><tr><td>NH3</td><td>0.01711</td><td>0.00156</td><td>0.17668</td><td>6.41899</td><td>1.30e-04</td><td>0.0018</td><td>1.48e-05</td><td>0.17797</td><td>0.81138</td><td>1.26e-06</td></tr><tr><td>CH4</td><td>0.01743</td><td>0.00252</td><td>0.31617</td><td>27.47914</td><td>2.10e-04</td><td>0.00176</td><td>1.78e-05</td><td>0.31764</td><td>1.69159</td><td>1.48e-06</td></tr><tr><td>N2</td><td>0.00112</td><td>2.60e-06</td><td>0.03992</td><td>0.15571</td><td>1.71e-07</td><td>7.77e-04</td><td>1.23e-06</td><td>0.03995</td><td>0.10766</td><td>8.29e-08</td></tr><tr><td>CO</td><td>0.00483</td><td>1.59e-04</td><td>0.07341</td><td>0.63087</td><td>1.01e-05</td><td>0.001</td><td>3.07e-06</td><td>0.07195</td><td>0.13</td><td>2.00e-07</td></tr><tr><td>H2O</td><td>0.01322</td><td>0.00186</td><td>0.21619</td><td>3.41078</td><td>1.39e-04</td><td>0.00154</td><td>1.59e-05</td><td>0.21715</td><td>0.34026</td><td>1.19e-06</td></tr><tr><td>H</td><td>0.00177</td><td>1.26e-05</td><td>0.15854</td><td>0.35331</td><td>9.49e-07</td><td>8.58e-04</td><td>1.59e-06</td><td>0.1585</td><td>0.17757</td><td>1.23e-07</td></tr><tr><td>He</td><td>0.00109</td><td>2.00e-06</td><td>0.00178</td><td>0.11862</td><td>8.28e-08</td><td>6.68e-04</td><td>9.09e-07</td><td>0.00172</td><td>0.07297</td><td>5.34e-08</td></tr><tr><td>H2</td><td>9.25e-04</td><td>1.65e-06</td><td>0.00241</td><td>0.09346</td><td>8.24e-08</td><td>6.84e-04</td><td>9.27e-07</td><td>0.00229</td><td>0.06904</td><td>5.16e-08</td></tr><tr><td>N2H2</td><td>0.02752</td><td>0.00506</td><td>0.17001</td><td>6.12761</td><td>0.00114</td><td>0.00252</td><td>4.23e-05</td><td>0.1737</td><td>0.58747</td><td>1.02e-05</td></tr><tr><td>N2H3</td><td>0.03014</td><td>0.0054</td><td>0.20215</td><td>5.11303</td><td>0.00144</td><td>0.00291</td><td>4.50e-05</td><td>0.20579</td><td>0.50594</td><td>1.27e-05</td></tr><tr><td>HNOH</td><td>0.02072</td><td>0.00225</td><td>0.1686</td><td>2.88127</td><td>7.32e-04</td><td>0.00228</td><td>2.22e-05</td><td>0.16995</td><td>0.30905</td><td>7.46e-06</td></tr><tr><td>NH2OH</td><td>0.02083</td><td>0.00236</td><td>0.18368</td><td>3.01083</td><td>7.28e-04</td><td>0.00245</td><td>2.49e-05</td><td>0.1852</td><td>0.35657</td><td>7.96e-06</td></tr><tr><td>H2NO</td><td>0.02207</td><td>0.00279</td><td>0.18213</td><td>3.32962</td><td>8.49e-04</td><td>0.00246</td><td>2.70e-05</td><td>0.18316</td><td>0.36436</td><td>8.45e-06</td></tr><tr><td>C2N2</td><td>0.03411</td><td>0.004</td><td>0.33786</td><td>50.54947</td><td>8.39e-04</td><td>0.00373</td><td>4.56e-05</td><td>0.34035</td><td>5.49715</td><td>1.02e-05</td></tr><tr><td>HCNH</td><td>0.0186</td><td>0.00116</td><td>0.2368</td><td>22.51091</td><td>2.26e-04</td><td>0.00216</td><td>1.37e-05</td><td>0.2366</td><td>3.44796</td><td>2.73e-06</td></tr><tr><td>HNC</td><td>0.00204</td><td>1.06e-05</td><td>0.14816</td><td>1.60722</td><td>9.73e-07</td><td>0.00127</td><td>3.48e-06</td><td>0.1479</td><td>1.08876</td><td>3.41e-07</td></tr><tr><td>NCN</td><td>0.03194</td><td>0.00346</td><td>0.22498</td><td>4.94461</td><td>0.00101</td><td>0.00364</td><td>4.46e-05</td><td>0.22586</td><td>0.55207</td><td>1.38e-05</td></tr><tr><td>HCOH</td><td>0.022</td><td>0.00338</td><td>0.15741</td><td>7.58491</td><td>6.44e-04</td><td>0.00227</td><td>2.52e-05</td><td>0.1609</td><td>1.024</td><td>4.94e-06</td></tr><tr><td>HOCHO</td><td>0.01361</td><td>0.00174</td><td>0.28976</td><td>75.42917</td><td>3.52e-04</td><td>0.00217</td><td>2.17e-05</td><td>0.2908</td><td>17.64686</td><td>4.40e-06</td></tr><tr><td>H2Oc</td><td>9.43e-04</td><td>1.75e-06</td><td>0.00165</td><td>0.16689</td><td>4.26e-07</td><td>6.92e-04</td><td>9.42e-07</td><td>0.00139</td><td>0.12253</td><td>2.25e-07</td></tr><tr><td>CH4c</td><td>8.03e-04</td><td>1.22e-06</td><td>0.00141</td><td>0.14217</td><td>2.85e-07</td><td>5.02e-04</td><td>4.96e-07</td><td>0.00123</td><td>0.08893</td><td>1.20e-07</td></tr><tr><td>NH3c</td><td>6.93e-04</td><td>1.10e-06</td><td>0.00137</td><td>0.12272</td><td>2.60e-07</td><td>4.98e-04</td><td>4.98e-07</td><td>0.00122</td><td>0.08811</td><td>1.19e-07</td></tr></tbody></table>
<figure><img src="image_26.png" /><figcaption>**Figure B1.** The predictive abilities (MAE) of Models A, C, and D for molecules. The molecules are sorted from the worst-performing to the best. Tables B1 and B2 provide detailed evaluations of the molecules. Error values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>
<figure><img src="image_27.png" /><figcaption>**Figure B2.** The predictive abilities (MAE) of Models A, C, and D for molecules. The molecules are sorted from the worst-performing to the best. Tables B1 and B2 provide detailed evaluations of the molecules. Error values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.</figcaption></figure>
<figure><img src="image_28.png" /></figure>

**Figure B3.** Variation of mean absolute error with pressure and isothermal temperature in Model A. The heatmap depicts how the mean absolute error changes across a range of pressures and isothermal temperatures. Error-values indicate increased prediction challenges around temperatures of 1400 K and pressures below 1 bar. Error values presented in the figure are calculated on normalized log (VMR) values, scaled from negative one to one. Conversion to dex in log (VMR) follows a linear transformation with a factor of 10.

This paper has been typeset from a *TEX/LATEX* file prepared by the author.