# Comparison of Neural Network and Gradient-Boosted Decision Tree Surrogates of Linear Gyrokinetic Simulations for Turbulent Transport Modelling in Tokamak Plasmas

M Lanzarone, E Fransson, Y Camenen, G Fuhr

► To cite this version:

M Lanzarone, E Fransson, Y Camenen, G Fuhr. Comparison of Neural Network and Gradient-Boosted Decision Tree Surrogates of Linear Gyrokinetic Simulations for Turbulent Transport Modelling in Tokamak Plasmas. 2025. [hal-05280765](https://hal.science/hal-05280765v1)

HAL Id: hal-05280765

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

Preprint submitted on 6 Oct 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
# Comparison of Neural Network and Gradient-Boosted Decision Tree Surrogates of Linear Gyrokinetic Simulations for Turbulent Transport Modelling in Tokamak Plasmas

Preprint of Paper to be submitted for publication in Physics of Plasmas

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

This work has been carried out within the framework of the EUROfusion Consortium, funded by the European Union via the Euratom Research and Training Programme (Grant Agreement No 101052200 — EUROfusion). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Commission. Neither the European Union nor the European Commission can be held responsible for them.
This document is intended for publication in the open literature. It is made available on the clear understanding that it may not be further circulated and extracts or references may not be published prior to publication of the original when applicable, or without the consent of the Publications Officer, EUROfusion Programme Management Unit, Boltzmannstr. 2, 85748 Garching, Germany or e-mail publications.officer@euro-fusion.org

Enquiries about Copyright and reproduction should be addressed to the Publications Officer, EUROfusion Programme Management Unit, Boltzmannstr. 2, 85748 Garching, Germany or e-mail publications.officer@euro-fusion.org
# Comparison of Neural Network and Gradient-Boosted Decision Tree Surrogates of Linear Gyrokinetic Simulations for Turbulent Transport Modelling in Tokamak Plasmas

M. Lanzarone,¹, a) E. Fransson,¹ Y. Camenen,¹ and G. Fuhr¹

¹) Aix-Marseille Université, CNRS, PIIM, UMR 7345 Marseille, France

(Dated: 6 October 2025)

A principal computational bottleneck in the integrated modelling of tokamak plasmas is the evaluation of turbulent transport fluxes, with even reduced quasi-linear gyrokinetic models remaining orders of magnitudes too slow for large scale predictive simulations and real-time applications. Machine-learning surrogate models offer a route to accelerate integrated modelling by replacing existing turbulent transport models with fast predictive surrogates. We present a systematic comparison of two widely used algorithms: Neural Networks (NNs), and Gradient-Boosted Decision Trees (XGBoost) to predict linear stability and unstable-mode growth rates, trained on linear QuaLiKiz simulations spanning a 22-dimensional input space covering the experimental domain of the JET tokamak. Model performance is assessed in a heavily optimised comparison study using $10^5$ training points and a scaling study covering $10^2$ to $10^7$ training points. NNs across both studies were significantly more sensitive requiring more tuning and diagnostic effort to match the predictive capability of the XGBoost models. Both models achieved comparable performance at $10^5$ data points when heavily optimised but XGBoost pulls ahead in the scaling study with reduced optimisation and demonstrates improved robustness of model hyperparameters. XGBoost additionally delivers ~50× faster inference and ~35× faster training and hyperparameter-optimisation times across the board. The reduced training times and increased robustness of XG-Boost are beneficial for faster and cheaper model development, especially in active learning based pipelines, while its faster inference speeds are ideal for rapid integrated modelling applications. The primary drawback to XGBoost is its step function predictions with discontinuous derivatives which may be problematic in certain applications.

## I. INTRODUCTION

Turbulent transport models are a crucial component of integrated modelling frameworks for tokamak plasmas¹, such as JINTRAC² or RAPTOR³, which couple multiple physics codes to self-consistently simulate plasma discharge evolution. Integrated modelling has key applications in studying experimental results, optimising plasma scenarios and designing future reactors. However, these workflows are computationally intensive, motivating the development of faster, accurate integrated modelling pipelines, which are of significant interest to the fusion community and critical for ITER plasma scenario development⁴. Moreover, a sufficiently rapid integrated modelling suite could enable real-time plasma control applications.

Individual codes within integrated modelling predict quantities evolving on different time scales, hence some modules are called more frequently than others. Turbulent transport models, for calculating particle and heat fluxes, are generally called on the order of $10^2 - 10^3$ times per second of simulated plasma and at multiple radial locations to adequately resolve both the temporal and spatial dynamics. This corresponds to approximately $10^4 - 10^5$ calls to turbulent transport solvers per second of integrated modelling simulation, which can become prohibitive when targeting confinement times of current and future devices, on the order of seconds.

A variety of physics-based models for turbulent transport exist, generally trading off computation time in return for

higher fidelity physics. High-fidelity non-linear flux-tube $\delta f$ gyrokinetic codes such as GKW⁵, GS²,⁶,⁷, and CGYRO⁸ are computationally prohibitive for integrated modelling, ranging into hundreds of thousands of CPU hours per call. Gyrokinetic codes can also provide the linear response, which combined with an ad-hoc model approximating the non-linear saturation, provides an estimate of the transport fluxes in what is called a quasi-linear model⁹. Quasi-linear models gain a large increase in speed due to simulating at significantly fewer wavenumbers and neglecting the non-linear coupling interactions between them. Despite this, they remain too slow for routine use in integrated modelling. Faster but reduced quasi-linear models have been developed, sacrificing fidelity for speed through various constraints and assumptions, with run times on the order of 1-10 CPU seconds. Some examples are: the gyrokinetic QuaLiKiz¹⁰ (QLK) which reduces run times by assuming only electrostatic perturbations, strongly ballooned eigenfunctions and $s-\alpha$ magnetic equilibria (circular, concentric magnetic flux surfaces); and TGLF¹¹,¹² which uses a gyrofluid model for increased speed. Nevertheless, even with fast reduced models such a QLK, turbulent transport simulations account for over 99% of the total integrated modelling runtime. A faster model resolving this bottleneck has the potential to reduce integrated modelling run times by a factor of over 100 and is highly sought after, provided it maintains comparable predictive performance.

Machine learning (ML) models have emerged as promising alternatives to direct computation of turbulent transport, offering surrogate models, trained on data from simulations or experiments, that reproduce the same input-output relationships while achieving orders-of-magnitude speedups compared to traditional models¹³⁻¹⁵. However, as machine learning mod-

<footnote>a) matisse.lanzarone@univ-amu.fr</footnote>
els approximate the statistical patterns of their training data, their accuracy is constrained by that data's quality and coverage. As a result, reliable predictions are typically limited to the training domain, and performance will not surpass that of the physics-based models from which this data is sourced. Studies have looked to offset this by using larger amounts of lower fidelity training data to capture general trends and fine tuning the resulting model on smaller amounts of higher fidelity data¹⁶ with the aim of improving prediction accuracy without incurring the increased cost of generating large high-fidelity datasets. A similar idea is behind the Gyrokinetics Database (GKDB)¹⁷ project which aims to tackle the issues related to generating large amounts of simulation data by storing turbulent transport simulations in an open-source database for various applications, including ML tasks.

Neural Networks (NNs) have been the predominant ML approach for turbulent transport surrogates, successfully predicting turbulent transport fluxes of: QLK (QLKNN)¹³,¹⁴,¹⁸, TGLF¹⁹,²⁰ and GENE¹⁵, as well as the eigenvalue solutions (growth rate and real frequency) of QLK²¹. Gaussian Processes (GPs) are another algorithm that have successfully been used, showing excellent performance with small training set sizes such as the surrogate of the GS2 code which predicts both eigenvalues and electron heat fluxes²². However, GP models suffer from severe restrictions on the amount of data used for training, scaling in complexity with $O(N^3)$ due to matrix inversion though newer techniques²³ have been shown to reduce this scaling. Some ML models have been applied recently in integrated modelling pipelines, with QLKNN being used for studies on the JET current ramp-up²⁴, DTT design²⁵ and for the ITER hybrid scenario²⁶, while the TGLF surrogate has been used for core-edge coupled ITER simulations²⁰.

To the authors' knowledge, no published works have explored the use of Gradient Boosted Decision Tree (GB-DTs) algorithms for turbulent transport simulation surrogates. These have been widely used in structured data tasks on both classification and regression where they have historically outperformed NN methods²⁷, evident in various Kaggle²⁸,²⁹ and KDD Cup competitions³⁰, although recent, more intricate implementations of NNs have been shown to close the gap³¹ and in certain cases to outperform GBDTs³². GBDTs may also have advantages when predicting sharp transitions²⁷ such as the threshold for the onset of mode instability in plasmas. Here, we present a comparative analysis of NNs, implemented in Keras³³, and a widely used GBDT algorithm, XGBoost³⁴, as surrogates for predicting linear gyrokinetic growth rates (regression) and classifying mode stability in tokamak plasmas trained on QuaLiKiz simulations. We perform a focused study with extensive hyperparameter optimisation on 10⁵ data points, and a scaling study, using lower levels of optimisation for tractability, across dataset sizes from 10² up to 10⁷ points.

The focus on growth rates as opposed to quasi-linear fluxes is motivated by the flexibility of decoupling outputs from a fixed saturation rule, enabling use of different existing models depending on the parameter regime (e.g. high beta or electrostatic) as well as development of new and arbitrary quasi-linear models. This ensures the continued applicability of the model no matter the saturation rule used for flux

prediction, eliminating the need for retraining. Sharp discontinuities in eigenvalues are expected as QLK clips negative growth rates for stable modes to zero, and with frequency predictions when the most unstable mode transitions from one type to another. This can be problematic as machine learning models, especially neural networks, tend to overly smooth these transitions which would ideally remain sharp. An initial stability classification model followed by a growth rate prediction when unstable has been suggested as beneficial in this particular application²¹, hence the inclusion of a classifier in this study. Both models are trained on data sourced from the same dataset of QLK simulations covering the experimental domain of the JET tokamak and previously used to train NNs to predict fluxes¹⁴ and linear eigenvalues²¹. For this work, the dataset is rescaled to conform to the Integrated Modelling and Analysis Suite³⁵ (IMAS) normalisation standard for gyrokinetic simulations³⁶ using conversion codes developed for use in the GKDB project.

The remainder of this paper is structured as follows. Section II introduces the training datasets and preprocessing, while Section III outlines the machine learning algorithms together with the training and evaluation methodology. The results and performance of the models are presented in Section IV, and Section V concludes with a summary and an outlook for future work.

## II. TRAINING DATASET

Machine learning models' predictions are a reflection of the data they are trained on and so it is important to understand the scope, quality, and limitations of that data. The training datasets used in this study are derived from the publicly available QLK-JET dataset³⁷, which contains simulations based on the JET experimental domain from the local, delta-f, electrostatic QuaLiKiz which uses s-α geometry. Each data point consists of 22 IMAS-normalised plasma parameters as inputs, and either mode stability (classification) or mode growth rate (regression) as ouptuts. The QLK-JET dataset is first filtered to remove outliers, converted to IMAS normalisation standards and rescaled for machine learning, resulting in the QLK-JET-IMAS dataset. From QLK-JET-IMAS, two large intermediate pools are generated via random sampling: the Classification Pool containing an even mix of unstable and stable points, and the Regression Pool containing only unstable points. Training datasets of various sizes, ranging from 10² to 10⁷ are then randomly sampled from these pools, on which both algorithms are trained. For the scaling study, the validation and test splits are sampled separately and remain constant throughout to ensure fair comparisons of testing metrics (See Sec. III A). An overview of this processing pipeline is shown in Fig. 1, with each step described in more detail below.

### A. QLK-JET Dataset

The QLK-JET dataset was originally generated to train an NN version of QuaLiKiz¹⁴. These simulations are based on
<figure><img src="image_3.png" /><figcaption>FIG. 1. Process used to obtain the training datasets from the original QLK-JET dataset³⁷. The original dataset is composed of 16 million QLK simulations with 18 wavevectors for total of 288 million data points.</figcaption></figure>

2135 experimental JET discharges, across 12,328 time windows, covering both quasi-steady-state and transient scenarios. The dataset spans the experimental JET domain without hyper-cube sampling, which from our results is beneficial for XGBoost, but can lead to concentrated clusters in the input space that may reduce training efficiency³⁸.

## B. Filtering and IMAS conversion

The QLK-JET dataset is processed into QLK-JET-IMAS as follows:

1. Only wavenumbers ranging in normalised QLK values from 0.1 to 1, corresponding to ion scales, are kept to focus on the larger scales of most interest for transport.

2. Each wavenumber is taken as an input, with it's corresponding growth rate as an output.

3. QLK is an eigenvalue solver and can produce the growth rates of multiple modes, however only the most unstable mode is kept for regression training.

4. Outliers are filtered out by removing the top and bottom 1% of values of the parameters listed in Table I removing approximately 12% of points, primarily due to numerical issues with QLK.

5. The data is normalised to the IMAS gyrokinetic standard³⁶ using a previously developed QLK-IMAS conversion script³⁹. This is not a simple linear rescaling; for example, the transformation for wavenumbers depends on the main ion mass (H, D, or T), resulting in splitting of the original nine QLK wavenumbers into multiple IMAS-normalised values ($k_\theta \rho_{\text{ref}}$, Fig. 2)

6. For faster training and improved model performance, the data in each dimension is rescaled to have a mean of 0 and a standard deviation of 1 using a standard linear transformation.

This results in a total of 256 × 10⁶ data points in the full QLK-JET-IMAS dataset. The distribution of inputs and outputs of the QLK-JET-IMAS dataset is shown in Fig. 2. Not shown in these distribution plots are main ion mass, secondary ion mass and secondary ion charge, which are still used as model inputs. The main ion is always a hydrogen

<table><caption>TABLE I. Simulation and surrogate model inputs in IMAS normalisation. Densities are given at $\theta = 0$. Further details on these conversions can be found in the IMAS Gyrokinetics paper³⁶.</caption><thead><tr><th>Normalised Input Parameter</th><th>Definition</th><th>Filtered</th></tr></thead><tbody><tr><td>Electron Density Gradient</td><td>R/L<sub>n<sub>e</sub></sub></td><td>×</td></tr><tr><td>Primary Ion Density Gradient</td><td>R/L<sub>n<sub>i0</sub></sub></td><td>×</td></tr><tr><td>Impurity Ion Density Gradient</td><td>R/L<sub>n<sub>iZ</sub></sub></td><td>×</td></tr><tr><td>Electron Temperature Gradient</td><td>R/L<sub>T<sub>e</sub></sub></td><td>×</td></tr><tr><td>Primary Ion Temperature Gradient</td><td>R/L<sub>T<sub>i0</sub></sub></td><td>×</td></tr><tr><td>Impurity Ion Temperature Gradient</td><td>R/L<sub>T<sub>iZ</sub></sub></td><td>×</td></tr><tr><td>Toroidal Velocity</td><td>u<sub>N</sub></td><td></td></tr><tr><td>Toroidal Velocity Gradient</td><td>u'<sub>N</sub></td><td></td></tr><tr><td>Perpendicular E × B Shear</td><td>γ<sub>E<sub>N</sub></sub></td><td>×</td></tr><tr><td>Magnetic Shear</td><td>Φ</td><td>×</td></tr><tr><td>Pressure Gradient</td><td>P'<sub>N</sub></td><td></td></tr><tr><td>Minor Radius</td><td>r<sub>0</sub></td><td></td></tr><tr><td>Safety Factor</td><td>q</td><td></td></tr><tr><td>Species Temperature</td><td>T<sub>sN</sub></td><td></td></tr><tr><td>Secondary Ion Charge</td><td>Z<sub>i1</sub></td><td></td></tr><tr><td>Primary Ion Density</td><td>n<sub>i0N</sub></td><td></td></tr><tr><td>Secondary Ion Density</td><td>n<sub>i1N</sub></td><td></td></tr><tr><td>Tertiary Ion Density</td><td>n<sub>i2N</sub></td><td></td></tr><tr><td>Primary Ion Mass</td><td>m<sub>i0N</sub></td><td></td></tr><tr><td>Secondary Ion Mass</td><td>m<sub>i1N</sub></td><td></td></tr><tr><td>Collisionality</td><td>V<sub>N</sub><sup>a/b</sup></td><td></td></tr><tr><td>Binormal Wavenumber</td><td>k<sub>θ</sub>ρ<sub>ref</sub></td><td></td></tr></tbody></table>

isotope, predominantly deuterium (76%), followed by hydrogen (22%) and tritium (2%) and these ratios remain nearly unchanged when considering only unstable points. The secondary ion represents an aggregated light impurity species with atomic masses between 2–10 and charge numbers between 3–20, while the tertiary ion represents a heavier impurity with atomic mass 58 and charge number 28. Impurity ion density and temperature gradients are the same for secondary and tertiary ions and labelled as $R/L_{nZ}$ and $R/L_{TZ}$, respectively.

## C. Classification and Regression Pools

As the QLK-JET-IMAS dataset was too large to fit in local memory, two smaller intermediate pools are sampled from it, one for each of classification and regression:

* Classification Pool: 10⁸ points, evenly split between linearly stable and unstable cases to avoid bias in the stability prediction models, with instability of the mode (0 for stable, 1 for unstable) as the target output.

* Regression Pool: 10⁸ unstable points only, with the growth rate as the target output. Stable points are excluded from regression because QuaLiKiz clips negative growth rates to zero, producing a sharp discontinuity that degrades regression performance due to the smoothing tendencies of neural networks and regularization. Attempts to mitigate this using a ReLU output activation on a mixed stable/unstable training data did
<figure><img src="image_4.png" /><figcaption>FIG. 2. IMAS Normalised Input-Output distribution for the full QLK-JET-IMAS dataset after filtering (originally sourced from³⁷). The black lines represent the distribution for all modes while the lighter salmon represents only unstable modes. All the quantities are inputs aside from γ (growth rate) and ω (real frequency) represent which are the outputs. Unstable modes account for approximately 32% of the filtered dataset.</figcaption></figure>

not improve results. The classification stage therefore acts as a stability cutoff, as done by Fransson et al.²¹.

### D. Training Sets

Training datasets for regression and classification are obtained by random sampling from their respective pools. All comparisons between NNs and XGBoost of a given training size are done using the same training, validation and test sets to ensure consistency.

For the focused study, 10⁵ points are sampled from the pools and split into 80%/10%/10% training/validation/test sets. This amount of points reflects what could be generated with approximately 5-10 million CPU hours of high-fidelity linear gyrokinetic simulations, the goal of a follow-up study.

For the scaling study, fixed validation and test sets of 10⁴ are first sampled from the pool and remain constant across the study. This is to ensure that the metrics used to evaluate performance remain directly comparable across the scaling study. Training sets ranging from 10² to 10⁷ are then sampled from the points remaining in the pool after sampling of the validation and test sets.

## III. METHODOLOGY AND MACHINE LEARNING ALGORITHMS

With the training datasets now defined, we turn to the methodology used to train, evaluate, and compare the performance of NNs and XGBoost for two tasks: instability classification and growth rate regression. In the classification task, the model predicts whether a given set of inputs is stable or unstable, whereas in the regression task, it predicts the value
of the growth rate given those inputs. Although ML models typically benefit from large training datasets, producing such datasets is often unfeasible for high-fidelity turbulent transport simulations, where each data point requires significant computational resources. We therefore also examine how each algorithm's performance scales with the number of available training points, allowing us to assess which method is more effective across different dataset sizes and computational constraints.

## A. Loss Functions and Performance Metrics

Machine learning algorithms work by minimising a metric called a loss function, summed or averaged across a set of training points, by automatic tuning of model's parameters during training. Both XGBoost's and NNs' parameters are tuned by performing gradient descent, moving in a direction given by the derivative of this loss function with respect to the model's parameters.

The classification algorithms are trained minimising logistic loss, often called "binary crossentropy" for Keras NNs and "binary:logistic" in XGBoost, defined as:

$$ \text{logistic loss} = - \frac{1}{N} \sum_{i=1}^{N} (y_i \log(p_i) + (1 - y_i) \log(1 - p_i)) \quad (1) $$

where $y_i$ are the individual true labels (taking values of 1 for unstable points or 0 for stable points), $p_i$ are the predicted probabilities of mode instability and N is the total number of points used for training. A perfect model has a logistic loss of 0.

Classification tasks are often compared on 3 metrics: precision, recall and F1. These scores all fall within the interval [0, 1] with higher values being better. For the purposes of classifying instability, a True Positive is an unstable mode correctly labelled as unstable while a False Positive is a stable mode incorrectly labelled as unstable. Similarly, a True Negative is a stable mode correctly labelled as stable while a False Negative is an unstable mode incorrectly labelled as stable.

Precision is the fraction of positive predictions that are correct. A model with higher precision will generate less False Positives. For classifying instability, a higher precision results in the model being more conservative when labelling points as unstable, biasing the classifier towards predicting more stable points. In the case of growth rates predictions that feed into quasi-linear saturation rules to output fluxes, wrongly predicted unstable modes tend to be the most problematic resulting in spurious fluxes where there would otherwise be none. The opposite case would also reduce accuracy, resulting in smaller fluxes than may be expected but these are not as critical for accurate integrated modelling as spurious fluxes in stable regions. Thus a model with higher precision would minimise false positives (stable points labelled as unstable in this case), potentially leading to higher performance down stream.

$$ \text{Precision} = \frac{\text{True Positives}}{\text{True Positives} + \text{False Positives}} \quad (2) $$

Recall is a measure of how many of the positive labels were correctly identified. A model with higher recall will minimise the number of positive cases that are misslabled. For instability classification, a higher recall corresponds to the model predicting more unstable points.

$$ \text{Recall} = \frac{\text{True Positives}}{\text{True Positives} + \text{False Negatives}} \quad (3) $$

F1 score is defined as the harmonic mean of recall and precision, and is commonly chosen as a good indicator of model performance. This will thus be the primary metric used to compare the two algorithms for instability classification.

$$ \text{F1} = \frac{2 \cdot \text{Recall} \cdot \text{Precision}}{\text{Recall} + \text{Precision}} \quad (4) $$

A confusion matrix provides a compact visual summary of classification outcomes and displays the counts of true positives, false positives, true negatives and false negatives from which precision, recall and F1 can be estimated at a glance or directly derived if desired.

For the regressions on growth rates, both XGBoost and NN models are trained to minimise the mean squared error (MSE) which is defined as:

$$ \text{MSE} = \frac{1}{N} \sum_{i=1}^{N} (p_i - y_i)^2 \quad (5) $$

Where $y_i$ are the true growth rates and $p_i$ are the predicted growth rates.

The evaluation of model performance is done primarily using the $R^2$ and mean absolute percentage error (MAPE) metrics, though the root mean square error (RMSE) is also reported. These are chosen rather than the MSE (on which models are trained) to provide a metric that is more intuitive to interpret on a figure than MSE. Additionally RMSE and $R^2$ are heavily correlated to MSE. These metrics are calculated as follows:

$$ R^2 = 1 - \frac{\sum_i (p_i - y_i)^2}{\sum_i (p_i - \bar{y}_i)^2} \quad (6) $$

where $\bar{y}_i$ is the mean of the true growth rates. $R^2$ has a maximum value of 1 for a perfect fit while particularly bad fits can result in large negative numbers. It should be noted that $R^2$ values are specific to the set upon which it is calculated, and as such should only be used to compare models tested on the same set of data.

$$ \text{MAPE} = \frac{1}{N} \sum_{i=1}^{N} \left| \frac{p_i - y_i}{y_i} \right| \quad (7) $$

$$ \text{RMSE} = \sqrt{\text{MSE}} = \sqrt{\frac{1}{N} \sum_{i=1}^{N} (p_i - y_i)^2} \quad (8) $$

MAPE and RMSE values can range from 0 to infinitely large, with values closer to zero being better.
## B. XGBoost Models

Decision trees are conceptually a flowchart-like algorithm where at every branch, the input set of points is split in two based on the value of one of the model inputs, eventually arriving at a "leaf node" which outputs the predicted quantities. The parameters of the model, tuned by the algorithm during training, are the input on which to make a split and the value at which the split is made. A tree that is too shallow can underfit data, failing to capture important patterns in the data, while a tree that is too large may overfit, learning noise in the training set both of which impact the models ability to make accurate predictions on unseen data (generalisation). Ensemble methods such as bagging<sup>40</sup> and boosting<sup>41</sup> mitigate these issues by combining multiple trees, improving generalisation compared to a single large tree. XGBoost is a gradient boosted decision tree (GBDT) algorithm that utilises boosting, building a series of decision trees sequentially with each tree correcting the errors of its predecessors by optimising a differentiable loss function through gradient descent. One of the major differences in the outputs of NNs and tree-based methods is that NNs provide smooth interpolation within the range of the training data, whereas tree based methods approximate interpolation in a piecewise manner, dividing the input space into discrete regions with constant outputs within each region. Using more complex trees or ensemble methods such as XGBoost creates more partitions in the input space, reducing average step sizes and, producing a closer approximation to smooth interpolation.

The XGBoost classification trees are trained on datasets with a balanced 50/50 mix of stable/unstable points, using "binary:logistic" loss, equivalent to the loss defined in Eq. (1).

The XGBoost regression trees are trained on datasets with only unstable modes, using MSE loss, defined in Eq. (5).

Aside from these differences, the training of both XGBoost classifiers and regressors proceeds identically. To optimise the hyperparameters of the model (Table II), a Bayesian search is done on the validation loss using the Hyperopt<sup>42</sup> python library. This is run for 3000 trials of different hyperparameter combinations over the ranges shown in the table, resulting in 3000 potential models. A total of 3000 trials for this initial Bayesian optimisation was deemed sufficient based on diminishing returns of improvement with increasing number of trials (see Fig. 4) and matches the total number of trials done for NNs. The ranges of hyperparameters initially covered standard values and were extended when best model performance occurred at the boundaries of the search range.

The 50 best performing models (in $R^2$ for regression and F1 for classification) on the validation data are then retrained using 5 different initialisations for a total of 250 models from which the best performer on validation data is chosen as the final model. This extra training step, where initialisation is varied is done to reduce the impact of initialisation and favour the best performing set of hyperparameters. The best model is retrained with 50 different initialisations to provide statistics of the model performance variance due to initialisation. The best performing model on validation data is then used to generate evaluation metrics on the unseen test data.

<table><caption>TABLE II. Hyperparameters varied during XGBoost optimisation and the values selected for the best performing XGBoost classification and regression models. learning_rate controls how strongly each new tree influences the overall prediction. gamma sets the minimum improvement in the loss function for a split to occur. reg_alpha (L1 regularisation) encourages sparsity by penalising large tree outputs, while reg lambda (L2 regularisation) penalises large tree outputs quadratically, discouraging overly strong contributions from individual leaves and improving stability. min_child_weight specifies the minimum total weight of data points in a child node for a split to be made. colsample_by sets the fraction of features randomly chosen for each node, tree, or level on which a split can be made, while subsample specifies the fraction of training data used to build each tree. max_depth limits how many successive splits are allowed within a tree, and n_estimators is the total number of successive trees (boosting rounds) in the model</caption><thead><tr><th>XGBoost<br>Parameter</th><th>Search<br>Range</th><th>Best Classifier<br>Values</th><th>Best Regressor<br>Values</th></tr></thead><tbody><tr><td>learning_rate</td><td>10<sup>-6</sup> - 10<sup>2</sup></td><td>0.0857</td><td>0.100</td></tr><tr><td>gamma</td><td>10<sup>-8</sup> - 10<sup>2</sup></td><td>3.25 × 10<sup>-4</sup></td><td>5.05</td></tr><tr><td>reg_alpha</td><td>10<sup>-6</sup> - 10<sup>5</sup></td><td>3.65</td><td>3.01</td></tr><tr><td>reglambda</td><td>10<sup>-6</sup> - 10<sup>4</sup></td><td>7.15</td><td>0.144</td></tr><tr><td>min_child_weight</td><td>1 - 100</td><td>7.71</td><td>39.8</td></tr><tr><td>colsample_bytree</td><td>0.75 - 1</td><td>0.835</td><td>0.924</td></tr><tr><td>colsample_bylevel</td><td>0.5 - 1</td><td>0.874</td><td>0.847</td></tr><tr><td>colsample_bynode</td><td>0.75 - 1</td><td>0.879</td><td>0.954</td></tr><tr><td>subsample</td><td>0.5 - 1</td><td>0.849</td><td>0.798</td></tr><tr><td>max_depth</td><td>5 - 20</td><td>13</td><td>17</td></tr><tr><td>n_estimators</td><td>1 - 500</td><td>220</td><td>246</td></tr></tbody></table>

## C. Neural Network Models

Neural networks<sup>43</sup> process a set of input values by propagating them through multiple layers of interconnected nodes. A diagram of a single example node is shown in Fig. 3, with a bias term b, inputs $x_1, x_2, ..., x_i$, and associated weights $w_1, w_2, ..., w_i$. At each node, the weighted sum $b + \sum_i x_i w_i$ is computed and then passed through the node's activation function f to produce the node's output $f(b + \sum_i x_i w_i)$. These outputs serve as inputs to the nodes of the next layer, and this process is repeated through successive layers until the final layer, where the outputs are the model's predictions. The weights and biases are the parameters tuned via gradient descent during training to reproduce the behaviour of the data it is trained on. As with trees, complex NNs with many nodes, layers or excessive training can overfit, memorising training examples while models that are too small can underfit reducing generalisation. The neural networks used in this paper are fully connected feed-forward network, i.e. each neuron in a layer receives input from all neurons in the previous layer, and information flows strictly from input to output, without any feedback loops or connections back to earlier layers.

The NN classifiers are trained using "binary crossentropy" as the objective (Eq.1), the same as the "binary:logistic" used in XGBoost classifiers. The single output neuron predicting stability is activated by the standard sigmoid function used for binary classification tasks. The NN classifiers are trained on the same datasets as the XGBoost classifiers, using the same
<figure><img src="image_5.png" /><figcaption>FIG. 3. Example of a neural network node, with a bias term *b*, inputs $x_1, x_2, ..., x_i$, and associated weights $w_1, w_2, ..., w_i$. At each node, the weighted sum $b + \sum_i x_i w_i$ is computed and then passed through the node's activation function *f* to produce the node's output. Parameters of the node, tuned during model training, are shown in blue.</figcaption></figure>

training, validation and test splits.

For the NN regressors, the objective function is MSE (Eq.5) as in the XGBoost models. The single output neuron predicting growth rate is in this case is activated by using a linear function. A linear activation was chosen over ReLU<sup>44</sup> (linear activation with negatives clipped to zero) for the following reasons: a ReLU activation will penalise negative outputs less and in testing led to worse model performance as all the training data is positive, and this method generalises better to future work following a similar methodology on regressions of growth rates from gyrokinetic simulations which can output negative values. The NN regressors are also trained on the same datasets as the XGBoost regressors, again using the same training, validation and test splits.

Aside from this change, the training of both the NN regressor and classifiers then follows the same path. Model shape is pushed towards a tapering cone by having a $1/N$ decrease in max layer size. This is similar to what was used by Ho et al.<sup>14</sup> and which in internal testing was generally found to perform better than multiple layers of the same size used by Citrin et al.<sup>18</sup> and Fransson et al.<sup>21</sup> who added dropout layers to reduce overfitting. The models are trained using an Adam optimiser and a batch size of 4096. A higher batch size was found to significantly speed up training while having no measurable impact on model performance and with 4096 being near the limit possible on local hardware. Early stopping was applied to accelerate training, with more aggressive criteria used during the larger hyperparameter searches. Early stopping halts training when the performance on the validation set stops improving, avoiding unnecessary computation and potential overfitting.

Hyperparameter optimisation is performed using a Bayesian search with the Keras Tuner to minimise the validation loss, over the parameter ranges shown in Table III. The Bayesian search uses the same data splits as for XGBoost, with 80% for training and 10% for validation. A total of 1000 different hyperparameter combinations are tested for up to 1000 epochs. To account for the higher sensitivity of NNs to random weight initialisations (compared to XGBoost), each of these hyperparameter combinations is trained with three random initialisations resulting in 3000 different models with only the best performing run of the three being kept for a final pool of 1000 models.

<table><caption>TABLE III. Hyperparameters explored during neural network optimisation and their values for the best performing NN classification and regression models. *Learning Rate* sets the step size for updating model weights during training. *Hidden Layers* and *Layer Widths* control the number of layers between the input and output nodes and their widths, respectively, defining model architecture. *L1 Regularisation* penalises large weights while *L2 Regularisation* penalises weights by their square, promoting sparsity and smoother weight distributions, respectively, to prevent overfit. *Activation Funcs.* specify the output functions used in the hidden layers to apply nonlinearities.</caption><thead><tr><th>NN Parameter</th><th>Search Range</th><th>Best Classifier Value</th><th>Best Regressor Value</th></tr></thead><tbody><tr><td>Learning Rate</td><td>10<sup>-5</sup> - 1</td><td>2.51 &times; 10<sup>-3</sup></td><td>6.31 &times; 10<sup>-3</sup></td></tr><tr><td>Hidden Layers</td><td>1 - 5</td><td>4</td><td>5</td></tr><tr><td>Layer 1 width</td><td>2 - 250</td><td>187</td><td>105</td></tr><tr><td>Layer 2 width</td><td>2 - 125</td><td>76</td><td>55</td></tr><tr><td>Layer 3 width</td><td>2 - 83</td><td>35</td><td>42</td></tr><tr><td>Layer 4 width</td><td>2 - 62</td><td>42</td><td>62</td></tr><tr><td>Layer 5 width</td><td>2 - 50</td><td>-</td><td>19</td></tr><tr><td>L1 Regularisation</td><td>10<sup>-8</sup> - 1</td><td>1 &times; 10<sup>-4</sup></td><td>1 &times; 10<sup>-4</sup></td></tr><tr><td>L2 Regularisation</td><td>10<sup>-8</sup> - 1</td><td>1 &times; 10<sup>-6</sup></td><td>1 &times; 10<sup>-6</sup></td></tr><tr><td>Activation Func.</td><td>swish, tahn, relu, sigmoid</td><td>tahn</td><td>swish</td></tr></tbody></table>

The best 50 hyperparameter combinations from the Bayesian search are then re-trained using 5 different initialisations yielding 250 models. Of these 250, the best performer on validation data is chosen as the final model. This final model is additionally retrained using 50 different initialisations for variance statistics as with the XGBoost model. Final model evaluations are done using the unseen test set.

## D. Scaling

For the scaling of algorithm performance with number of training samples, performance of XGBoost and NNs are compared across a large range of training points from $10^2$ to $10^7$. Once again, identical set are used for both methods at each dataset size. A Bayesian search of 150 trials is used to optimise hyperparameters and was shortened relative to the focused study for tractability based on the results from Fig. 4, where a short search was enough to find a well performing set of hyperparameters. The best set of hyperparameters from this search was then given 5 different initialisations, with the best performer on validation data being chosen as the final model.

# IV. RESULTS

## A. Hyperparameter Search

Here we discuss the results of the hyperparameter searches for each combination of NN/XGBoost and classifier/regressor in the focused study. In terms of training times, NNs take significantly longer to train and optimise hyperparameters. The Bayesian search of 3000 models took 95 hours on local hardware for the NNs (1.9 minutes each) and only 2 hours and 40
<figure><img src="image_6.png" /></figure>

<figure><img src="image_7.png" /><figcaption>FIG. 4. Classifier F1 scores (top) and regressor $R^2$ scores (bottom) over the course of the initial Bayesian hyperparameter search with NNs in blue and XGBoost in green. *Trials are the total number of trainings done for each model during the search: 3000 individual hyperparameter combinations for XGBoost and 1000 sets of hyper-parameter combinations with 3 initialisations each for the NNs.</figcaption></figure>

minutes for XGBoost, a factor of 38 faster.

Fig. 4 shows the best performance achieved for all models, as a function of the number of models tested during the initial Bayesian searches. Across all cases, performances improves sharply in the first few trials, with subsequent gains beyond this point being greatly reduced. These trends are consistent with those reported by Honda and Narita⁴⁵ when training turbulent transport surrogates, in which rapid early improvements followed by diminishing returns were observed when optimising NN hyperparameters with Hyperopt. Differences in the smoothness of improvement are likely due to a combination of factors. The first are differences in the NN and XG-Boost algorithms themselves, and the Keras Tuner and Hyperopt libraries used for the hyperparameter search in each case. Another difference is in the size of the hyperparameter search. The XGBoost search covers 3000 hyperparameter combinations, each evaluated once. In contrast, the neural network search explores only 1000 different hyperparameter combinations, but each is trained with three random initialisations to address the higher variance from initialisation discussed in Sec. IV B.

To determine which hyperparameters have an important effect on the predictive capability of the surrogate models, and to justify the search ranges used during training, hyperparameter scan results were plotted against model performance for all the hyperparameters of the model, some of which are shown in Figs. 5-7. Due to the Bayesian search algorithm used, conclusions on hyperparameter importance and correlation with performance should be interpreted with caution as the algorithm tends to cluster in what it determines are optimal ranges resulting in sparse data outside of these. However, that clustering occurs does imply that optimising hyperparameters still has value. Further investigation on hyperparameter importance and optimisation are left for future work.

In general, most of the hyperparameters are weakly correlated with performance. Some however do present a minimum threshold below which performance degrades significantly. For all the models, the scaling with learning rate, shown in Fig. 5, has the strongest trend, resulting in a band of optimal values in the range of $10^{-2}$ to 1 for the XGBoost models and $10^{-3}$ to $10^{-1}$ for the NN models. This strong correlation means that it is important to correctly optimise for the learning rate to extract the best performance possible.

Examples of the impact of model complexity on regressor performance over the Bayesian search are shown in Fig. 6. As expected, increased model complexity leads to generally higher performance across the board, though the scaling inevitably flattens past a certain point. This also held true for the classifiers though they are not shown here. The scaling of XGBoost predictions with maximum depth and number of sequential trees are in line with previous work⁴⁶ where stability predictions were made using XGBoost classifiers trained on an 8D hyperrectangle, covering a larger input space. Of note with the NNs is the sharp jump in performance from one to two layers, also reflected though to a lesser extent in the classifier performance. We therefore recommend at least 2 hidden layers for NNs in these applications.

Not shown on the plots but relevant are a large initial hidden layer size being important for the NNs though succeeding layer sizes were not strongly correlated with performance. Increasing model complexity beyond the ranges listed in Tables II and III was not done as diminishing returns in performance similar to those in Fig. 6. Overfitting also becomes more likely as the parameter counts get large relative to the number of training points and the performance improvements were visibly reduced towards higher parameter counts. In hindsight however, with the speed of XGBoost training, the number of XGBoost estimators could have been increased for only a small cost in computation resources for a potential gain in predictive capabilities.

For the XGBoost regression models, there was a weak correlation of alpha and lambda regularisation (also called L1 and L2 regularisation) with model performance. In the case of alpha regularisation shown in Fig. 7, the maximum appears to be at a boundary, however the Bayesian search was allowed to explore much higher than this but never did so. Looking at regularisation values for best performing XGBoost models in Table II, these values, aside from "gamma" for classifiers and "reglambda" are all considered high levels of regularisation
<figure><img src="image_8.png" /></figure>

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

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

<figure><img src="image_11.png" /><figcaption>FIG. 5. Model performance across the Bayesian hyperparameter search as a function of learning rate. The black line (left axis) denotes the highest score obtained at each learning rate while the coloured histogram (right axis) shows the number of times each learning rate range was evaluated over the Bayesian search. Note that a learning rate of $10^0$ is already extremely large for NNs in the Keras framework and the search never explores the higher allowed ranges.</figcaption></figure>

which push the model towards simpler trees. In comparison, the best performing NNs in Table III have low to intermediate levels of regularisation.

## B. Best Model Choice

Each of the final Bayesian hyperparameter searches result in 250 models (50 hyperparameter combinations × 5 initialisations) with varying performance in each of the evaluation metrics. The choice of the best performing model out of these 250 could be done using just the $R^2$ value for regressors and the F1 score for classifiers, however looking at multiple metrics in tandem is likely to lead to better performing models. Shown in Fig. 8 are plots of these 250 model's validation performance across two important metrics: $R^2$ and mean absolute percentage error (MAPE) for the regressors, and F1 and precision scores for the classifiers. The choice of precision over recall here is to reduce spurious fluxes, as explained in Sec. III A. The best model candidates are the darker points forming the upper left/right edge of each cluster where the two metrics are at their best values. As multiple options exist, this allows some small optimisation of increased performance

in one of the two metrics at the cost of the other. For the regressors, the choice is made to use the leftmost models, accepting a minor trade-off in $R^2$ for a significant reduction in the MAPE. This should, in theory, increase the model performance at low growth rates which are generally present around critical thresholds and at small wavenumbers. For the classifiers, the trade-off in F1 and precision is roughly equal in both cases and a higher F1 score is chosen.

The difference in performance of the models in Fig. 8 was most likely due to variations in hyperparameter and randomised model initialisation. To quantify the impacts of changing hyperparameters, the performance of the best 50 hyperparameter combinations (from the Bayesian search) was averaged across the five random models initialisations. This is shown in Fig. 9 where the spread of results even among the best 50 hyperparameter combinations is quite large for the NNs relative to XGBoost. Similarly, the impacts of random initialisation a quantified by comparing the performance of 50 random initialisations of the best model, shown in Fig. 10. Again, the NN has a noticeably larger spread in performance than for XGBoost. A summary of these spreads is given in Table IV showing the XGBoost points are generally 1.5 to 10 times more grouped and highlighting the robustness of the
<figure><img src="image_12.png" /></figure>

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

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

<figure><img src="image_15.png" /><figcaption>FIG. 6. Regression model performance across the Bayesian hyperparameter search as a function of parameters that describe complexity. On the left in blue are the NNs with total trainable parameters (top) and number of hidden layers (bottom) and on the right in green are the XGBoost models with maximum tree depth (top) and number of sequential tress (bottom). The black line (left axis) denotes the highest score obtained at each "complexity" value while the coloured histogram (right axis) shows the number of times each of these values was evaluated over the Bayesian search.</figcaption></figure>

<table><thead><tr><th rowspan="2">Model</th><th colspan="2">Hyperparameter Combination</th><th colspan="2">Initialisation</th></tr><tr><th>R<sup>2</sup></th><th>MAPE</th><th>R<sup>2</sup></th><th>MAPE</th></tr></thead><tbody><tr><td>NN Reg</td><td>~ 0.035</td><td>~ 10%</td><td>~ 0.015</td><td>~ 6%</td></tr><tr><td>XGB Reg</td><td>~ 0.005</td><td>~ 1%</td><td>~ 0.005</td><td>~ 1%</td></tr><tr><td>XGB / NN</td><td>~ 1/7</td><td>~ 1/10</td><td>~ 1/3</td><td>~ 1/6</td></tr></tbody></table>

<table><caption>TABLE IV. Range of model performance among the best 50 hyper-parameter combinations and due to 50 different initialisations of the best hyperparameter combination.</caption><thead><tr><th rowspan="2">Model</th><th colspan="2">Hyperparameter Combination</th><th colspan="2">Initialisation</th></tr><tr><th>F1</th><th>Precision</th><th>F1</th><th>Precision</th></tr></thead><tbody><tr><td>NN Class</td><td>~ 0.02</td><td>~ 0.03</td><td>~ 0.075</td><td>~ 0.0175</td></tr><tr><td>XGB Class</td><td>~ 0.005</td><td>~ 0.006</td><td>~ 0.05</td><td>~ 0.006</td></tr><tr><td>XGB / NN</td><td>~ 1/4</td><td>~ 1/6</td><td>~ 2/3</td><td>~ 1/3</td></tr></tbody></table>

algorithm to changes in hyperparameter combinations. This generally means that XGBoost models are likely to still perform well even with smaller hyperparameter searches and a smaller number of random initialisations.

<table><caption>TABLE V. Classifier performance metrics for the best NN and XG-Boost models on the test dataset. Metrics are defined in Sec. III A with the prediction time being the mean per-sample inference time, calculated from the total runtime of 10,000 predictions.</caption><thead><tr><th>Model</th><th>F1</th><th>Precision</th><th>Recall</th><th>Prediction Time (s)</th></tr></thead><tbody><tr><td>NN</td><td>0.924</td><td>0.913</td><td>0.935</td><td>1.51 × 10<sup>-4</sup></td></tr><tr><td>XGBoost</td><td>0.924</td><td>0.905</td><td>0.944</td><td>2.76 × 10<sup>-6</sup></td></tr></tbody></table>

### C. Classification

Having chosen the best models, we move onto evaluating their performance on the test sets, starting with the classifiers used to predict mode stability. The performance of the best classifiers using NNs and XGBoost is summarised in the confusion matrices in Fig. 11 with the associated classification performance metrics listed in Tab.V. Both algorithms achieve essentially identical performance in classification metrics, however the XGBoost models retain a significant edge over NNs with inference speeds on the order of 50 times faster.

In this binary classification problem, models output a value
<figure><img src="image_16.png" /><figcaption>FIG. 7. XGBoost Regressor performance across the Bayesian hyperparameter search as a function of Alpha (L1) regularisation. The black line (left axis) denotes the highest score obtained for each Alpha value while the coloured histogram (right axis) shows the number of times each value was evaluated over the Bayesian search.</figcaption></figure>

ranging from 0 (stable) to 1 (unstable) and by default these are rounded to the nearest integer (i.e. 0.5 threshold) which gives the model prediction. Shown in Fig. 12 are the effects of adjusting this rounding threshold from 0.1 to 0.9 for both the NN and XGBoost classifiers. Doing this allows the model to be biased towards either precision or recall which is useful in situations where one of these metrics is preferable. As mentioned in Sec.III A, our application favours precision over recall, since false positives can result in spurious fluxes that are more problematic than missed instabilities. Both algorithms show similar behaviour here and it can be seen that increasing the threshold slightly to around 0.6 results in improved precision without heavy loss in the F1 score and thus overall performance. Moving too far from a central threshold causes a sharp drop in F1 score which is undesirable. Unfortunately this drop off occurs more quickly when pushing the model towards higher precision, the reverse of the ideal situation in this application.

## D. Regression

We now move onto the performance of the regressors used to predict growth rates of unstable modes. This is summarised for XGBoost and the NNs in the prediction heatmaps in Fig. 13 and the associated regression performance metrics, listed in Table VI. Both algorithms achieve essentially identical performance in terms of regression metrics. The XGBoost model may have a small edge in terms of MAPE which is particularly evident in the tighter spread at low growth rate predictions and interesting for accurate predictions around instability thresholds. As with the classifiers, XGBoost retains an edge in inference speeds which are once again approximately 50 times faster. An interesting feature of the XGBoost algorithm is its apparent avoidance of negative predictions, in contrast with the NN model’s continuous prediction function

<table><caption>TABLE VI. Regressor performance metrics of the best NN and XG-Boost models on the test dataset. Metrics are defined in Sec. III A with the prediction time being the mean per-sample inference time, calculated from the total runtime of 10,000 predictions.</caption><thead><tr><th>Model</th><th>R<sup>2</sup></th><th>RMSE</th><th>MAPE</th><th>Prediction Time (s)</th></tr></thead><tbody><tr><td>NN</td><td>0.944</td><td>0.0744</td><td>25.1%</td><td>1.51 × 10<sup>-4</sup></td></tr><tr><td>XGBoost</td><td>0.946</td><td>0.0728</td><td>22.8%</td><td>3 × 10<sup>-6</sup></td></tr></tbody></table>

which can produce in negative values. This arises from XG-Boost’s step-like predictions: with no negative training values, very few leaf nodes output negative values. While this is unlikely to occur when negative growth rates are included in the training data, it should still be monitored to ensure proper model behaviour in such cases.

## E. Growth Rate Spectrum and R/L<sub>T</sub> Scan

It is helpful to examine actual predictions of the surrogate models, not just statistical metrics. Fig. 14 shows a representative prediction of a growth rate spectrum around the peaks of the input distribution (see Fig. 2), using a classifier, followed by a regressor to predict growth rate if the mode is predicted as unstable. Predictions are shown for: NN (blue), XG-Boost (green) and an average across an ensemble of 10 different XGBoost regressors (red). Predictions in regions outside the training domain of $0.115 < k_\theta \rho_{\text{ref}} < 2$ are unreliable extrapolation and the $k_\theta \rho_{\text{ref}} \ge 1.41$ region was sparsely sampled in training and less accurate predictions in this region are expected. Dotted blue and green lines represent the predictions of only the NN and XGBoost regressors across the whole domain, respectively. These demonstrate how the regressor performance is unreliable when extrapolating in untrained (stable) regions, hence the need to first classify the mode before predicting growth rates.

Overall, predictions are fairly accurate, consistent with the evaluation metrics. Both ML models often struggle to accurately capture rapid changes due to regularisation dampening their sensitivity to local variations. This can be an issue when predicting sharp transitions, such as between unstable modes as the transition point will tend to be smoothed over. It can however also be beneficial as in Fig. 14, where both models appear robust to an anomalously high peak growth rate in the QLK data that does not align with the trend of the other QLK points and is a likely numerical error that made it through filtering. The use of a classifier generally covers the downsides of this smoothing around the sharp onsets of instability, but often results in cutoffs that are too abrupt as seen in the figure on the left and right edges of the growth rate peak. This is due to the regressors not being trained on any negative growth rate data, hence the cutoff points determined by the classifiers do not always closely coincide with regressor predictions close to zero, resulting in large drops in growth rate. Training regressors on negative growth rate data, if available, is likely to improve the regression models around these stability transition regions and may make the classification step unnecessary.

A problematic scenario for the use of such a classifier-
Best Model Choices

<figure><img src="image_17.png" /><figcaption>FIG. 8. Performance of the 250 models obtained in the second optimisation step (50 model architectures with 5 separate initialisations) plotted as a function of R² and MAPE for regressors, and F1 and precision for the classifiers. The best candidates are highlighted in darker colour and the chosen "best" models as dark crosses. Regressors are shown on the left and classifiers on the right. Neural network models are shown as blue triangles and XGBoost models as green circles.</figcaption></figure>

Effect of Hyperparameters

<figure><img src="image_18.png" /><figcaption>FIG. 9. Scatter plots of showing the performance, averaged across each set of 5 initialisations, of the best 50 models for each task. Regressors are shown on the left and classifiers on the right, with the NN models as blue triangles and XGBoost models as green dots.</figcaption></figure>
Best Model Initilisation Variance

<figure><img src="image_19.png" /><figcaption>FIG. 10. Scatter plots of showing the performance of the best model with 50 different random initialisations for each task. Regressors are shown on the left and classifiers on the right, with the NN models as blue triangles and XGBoost models as green dots.</figcaption></figure>

regressor model is visible in Fig. 14 for the NN prediction at small wavenumbers. As $k_\theta \rho_{\text{ref}}$ approaches zero and exits the regressor training domain, the NN growth rate predictions begin to grow again rather than approach 0. In this region, the classifier is expected to "turn off" the regressor but gets the threshold slightly incorrect resulting in a set of higher than expected growth rates. Despite XGBoost showing better performance in this region of the figure, similar behaviour to the NN was regularly observed for the XGBoost models in other spectra not shown here. This problem is likely due to the lack of negative growth rate data in the training set which would increase the accurate prediction domain of the regressor into the stable region and reduce the need for classification to identify the stability transition point. Both classifiers also occasionally predicted instability, often in tandem, on otherwise stable QLK scans, which would also lead to spurious fluxes in downstream applications.

The step-function nature of decision tree-based models is clearly visible in Fig. 14. Because decision trees partition the input space into discrete regions, XGBoost models will always produce piecewise-constant predictions. The impact of these steps however depends heavily on the prediction grid. When evaluated at lower resolutions more typical of quasi-linear models, especially if predictions are made at the same training wavenumbers, the steps are largely hidden and have little practical effect. In contrast, at high prediction resolutions as in the figure (~1000 points), not only are the steps more visible, but may also result in nonphysical behaviours, such as spurious inversions in local concavity. An example of this is observed in Fig. 14, where the ensemble prediction of XGBoost models exhibits a reversal in local curvature around the fourth QLK data point from the right. These step-like artefacts are therefore most likely to be problematic in

applications requiring high resolution predictions relative to the training data or when gradients of the predicted quantities are needed.

These steps were most apparent on wavenumber spectra, where the model is trained on sparsely sampled input variables and the discontinuities in the XGBoost predictions align noticeably with the set of training wavenumbers shown in Fig. 2. We suspected that the large step sizes were primarily due to this sparse training distribution and that a more smoothly distributed training set would reduce prediction step sizes, provided the wavevector has a strong effect on model predictions. To test this, we first examined the model's feature importance parameter called "gain", which quantifies the contribution of each input feature to the model's predictions, measured as the average reduction in training loss achieved by tree splits using that feature. The five features with the highest gain are shown in Fig. 15. As would be expected from the physics, the model places a high importance on the wavenumber $k_\theta \rho_{\text{ref}}$ (solid outline) as well as on $R/L_{T_i}$ (dashed outline). This makes $R/L_{T_i}$ an ideal input on which to test our hypothesis as it is also one of the most smoothly distributed training inputs. Scans in $R/L_{T_i}$, one of which is shown in Fig. 16, showed significantly smoother predictions, confirming that the sparsely sampled wavenumbers were likely responsible for the larger steps in the predictions. The XGBoost ensemble approach also reduces the step sizes in the $R/L_{T_i}$ scan more significantly than for the growth rate spectra in Fig. 14 with the red ensemble line being markedly smoother than the single model in green.

These results motivate the use of smooth, randomly distributed training data for training DT algorithms when possible. This is especially true for wavenumbers which have a tendency to be fixed in simulations and therefore in the training data sourced from them. Other ways of reducing step sizes
<figure><img src="image_20.png" /></figure>

<figure><img src="image_21.png" /><figcaption>FIG. 11. Confusion matrices for the best performing Neural Network (blue, top) and XGBoost (green, bottom) classifiers on the test dataset.</figcaption></figure>

could include increasing the model complexity, such as increasing the depth, number of trees or decreasing the minimum child weights, which increase the number of leaf nodes in the tree, smoothing out individual tree predictions. However, more complex trees raise the risk of over-fitting and, like ensembles, increase inference time.

## F. Scaling with Training Samples

We now move onto the scaling study, looking at the performance of NNs and XGBoost across a large range of training

<figure><img src="image_22.png" /><figcaption>FIG. 12. Effect of varying the classification threshold on precision (dotted), recall (dashed), and F1 score (solid) for both the NN (blue) and XGBoost (green) classifiers. Increasing the threshold slightly (~0.6) improves precision without substantial loss in F1, while moving further from 0.5 leads to a sharp drop in F1, particularly when biasing the models toward higher precision.</figcaption></figure>

set sizes for both classification and regression. These models use a shorter Bayesian search detailed in Sec.III D and their performance is shown in Fig. 17. The initial set of models are shown in dark blue for the NNs and dark green for XGBoost and it can be seen that XGBoost models outperform the NNs across the board.

An unexpected drop in NN regression performance was observed around $3 \times 10^4$ training points, suggesting potential issues in the training pipeline. To investigate, changes to early stopping, training set splits, batch size, optimisation and number of epochs were individually tested, but to little success. Examination of the training loss curves revealed that NNs prior to this performance drop exhibited a second descent in training loss unlike those within the anomaly which had only one.

To specifically target this behaviour, early stopping was removed entirely, the number of epochs was increased by a factor of 15 and regularisation was strengthened to avoid overfitting. The results of this final attempt are shown on the figure as the red triangles. These modifications resulted in markedly improved NN performance above $10^4$ training points though no change was seen below this dataset size. Despite this increased effort, the NN models still did not outperform XG-Boost which required far less computational effort. It is possible that further optimisation could yield models that outperform XGBoost but this same idea could be applied to XG-Boost and is not guaranteed in either case.

The lighter coloured points in the figure are obtained by re-using the same hyperparameters as the best models from
<figure><img src="image_23.png" /></figure>

<figure><img src="image_24.png" /><figcaption>FIG. 14. Example of model predictions of growth rate spectra around modal dataset values. QLK "ground truth" values are shown in the black crosses, NN predictions are shown in blue, XGBoost predictions in green and an ensemble of 10 XGBoost regressors in red. Solid lines are with both regression and classification while dotted lines show just the regressor predictions to visualise model extrapolation.</figcaption></figure>

<figure><img src="image_25.png" /><figcaption>FIG. 13. Growth rate predictions on the test dataset for the best performing Neural Network (blue, top) and XGBoost (green, right) regressors against QLK "ground truth" values.</figcaption></figure>

the focused study, retrained and tested on the same training subsets used in the scaling study. For the XGBoost models, at large training sets above that used to train the focused study models (8 × 10⁴ points) the performance is quite similar. However, as training dataset size is lowered below that of the focused study, the performance increasingly degrades relative to the re-optimised models. For the NN models, the focused study hyperparameter models and re-optimised model

performances alternate and remain significantly worse than XGBoost at smaller training set sizes.

Overall, XGBoost predictions appear superior across most dataset sizes and the models require significantly less time for tuning, training and inference.

# V. CONCLUSIONS AND FUTURE WORK

## A. Conclusions

In this work, we have presented a systematic comparison between neural networks and gradient-boosted decision trees (XGBoost) for predicting linear growth rates and classifying mode stability, trained on QuaLiKiz simulations covering the JET experimental space and rescaled to IMAS standards. Overall, both algorithms perform similarly, when heavily optimised, in terms of accuracy predicting QLK growth rates and stability.

However, XGBoost has a number of advantages over NNs that are particularly beneficial for use as a turbulent transport surrogate. The most important advantage is the rapid inference speed of XGBoost, on the order of 1.5 × 10⁻⁶s on a local desktop for both classifier and regression, which is approximately 50 times faster than the NN models. This computational efficiency provides significant flexibility in
<figure><img src="image_26.png" /><figcaption>FIG. 15. Feature importance of the XGBoost regressor ranked by gain, which quantifies the average information gain from splits using each input across all trees in the model. High-gain features contribute most to the predictions, with the five most influential inputs shown. The inputs investigated here: wavevector $k_\theta \rho_{\text{ref}}$ and the normalised temperature gradient $R/L_{T_i}$, are highlighted with solid and dashed outlines, respectively.</figcaption></figure>

<figure><img src="image_27.png" /><figcaption>FIG. 16. Example of model predictions for an $R/L_{T_i}$ scan around modal dataset values to check for increased smoothness in XGBoost predictions when training data for that dimension is more continuously distributed. NN predictions are shown in blue, XGBoost predictions in green and an ensemble of 10 XGBoost regressors in red. Solid lines are with both regression and classification while dotted lines show just the regressor predictions. This region is well trained over and the model is not expected to be extrapolating.</figcaption></figure>

how the models can be used. Multiple models can be run in parallel and still output faster than a NN model. This allows using an ensemble of models to further reduce variance and improve overall model predictions, or to have multiple models each individually predict various quantities such as mode frequencies and quasi-linear weights, which is likely to perform better than a single model trying to predict all of these at once. Other benefits of faster inference speed are: increased spatial and temporal resolutions using the same computational budget, the ability to run in real time even on less powerful hardware such as laptops and the lower computational cost of each call. Beyond the faster inference speed, XGBoost models are simpler to implement from a coding perspective and the regressors had a slight edge in percentage error, favouring prediction accuracy at small growth rates that is likely to benefit flux predictions. XGBoost also has less hyperparameters to optimise and is more robust to both hyperparameter variations and random initialisation. This along with approximately 40 times faster training times than NNs makes its particularly good in applications where training and optimisation search times are long relative to the available computational budget. An XGBoost model can therefore be optimised in significantly less time than a NN model when on limited computational budgets. This also heavily benefits workflows that utilise active learning to guide dataset generation, as in Zainisi et al.⁴⁷, where hyperparameter optimisations, training and inference are repeated many times. Increased robustness to hyperparameter variations further benefits XGBoost here as the hyperparameter optimisation may not need to be done at every step to retain good performance. This frees up more computational budget with which more data can be generated, as opposed to spent on training to optimise which points to simulate.

While XGBoost has many advantages, it is not without downsides which favour NNs in certain applications. NNs generally have smooth, continuous outputs aside from any irregularities introduced by activation functions such as ReLU. XGBoost on the other hand splits the input space into discrete regions, resulting in step-function predictions. This is mostly problematic when training data for inputs that are important to the prediction are sparsely sampled in training such as the wavenumber in this study. This results in larger steps leading to less accurate predictions when in between training points and a difficulty of calculating gradients of output quantities, as well as occasionally producing spurious local concavity inversions. It's important to note that this issue is present only in applications where the predictions are made on a grid that has resolutions approaching or exceeding that of the input data as these steps are not visible at low resolutions. This is also less relevant for wavenumbers, as they tend to be calculated on a fixed grid that can be made to match training data.

To summarise, we generally recommend XGBoost in applications with smoothly distributed training data, particularly in the dimensions that have a large effect on growth rates and stability though this may not be required when sampling at
<figure><img src="image_28.png" /><figcaption>FIG. 17. Scaling of XGBoost (green, circles) and NNs (blue/red, triangles) for regression $R^2$ (left) and classification F1 score (right) with the number of points used to train models. The NN points in blue represent using a similar optimisation framework as for XGBoost, with a lighter version of that used in the more focused study. For both algorithms, the lighter "No Opt" points are the performance of models reusing the best hyperparameters from the focused study while the darker points have re-optimised hyperparameters. The NN points in red are significantly more optimised in attempts to bridge the gap to the XGBoost performance after the sharp drop-off in performance below $10^5$ training points, discussed in Sec.IV F. The largest classifier models are not optimised in this plot due to limitations in local workstation memory when conducting Bayesian searches on large datasets.</figcaption></figure>

the same values used in training such as is often the case with wavenumbers. In applications that require gradients of predicted quantities or extremely smooth predictions on fine grids then NNs are likely to perform well, provided the increased complexity, computational requirements and slower inference speeds are not prohibitive.

## B. Future Work

We conclude with some avenues for future research based on the findings of this study.

Firstly, the training dataset contains input variables that are heavily correlated such as ion and electron density gradients as well as main ion charge and masses. Feature importance analysis could be used here to potentially reduce the number of input dimensions, allowing for a higher density of training points in the more important input dimensions improving performance for little to no cost, similar to that done by Gillgren et al.⁴⁸. This, along with feature engineering to combine or adjust input dimensions (e.g. squaring, etc), could also help boost NN performance as this has historically been more important in NNs than in XGBoost.

Next, this study predicts only growth rates during regression, whereas a full surrogate model would also require predictions of quasi-linear weights, real frequencies, and in certain cases eigenfunctions to be used in quasi-linear saturation rules to calculate fluxes. Investigating the capabilities of XG-

Boost in reproducing these additional quantities is therefore a logical next step.

In a similar vein, it would be valuable to explore more relevant metrics and loss functions than $R^2$ and MSE for growth rate predictions. While these standard metrics reflect how accurately the model reproduces QLK growth rates on average, not all predicted growth rates contribute equally to the total flux after passing through quasi-linear saturation rules. In particular, errors at smaller wavenumbers correspond to larger transport scales and therefore have a greater impact on total flux. One way to address this is to weight the points relative to wavenumber in the loss function, for example using a quasi-linear approximation where contributions of growth rates to total flux scale as $\gamma/k^2$. Weighting errors by $1/k^2$ emphasises accuracy at low wavenumbers, and the exponent can be adjusted to find a balance between overall model accuracy and accuracy as lower wavenumbers. Alternatively, the training dataset could be populated with more points at smaller wavenumbers, applying a similar weighting strategy for data selection rather than modifying the loss function, to ensure that the model learns more effectively in the regions most critical for flux prediction.

Another recommendation is the use of datasets containing negative growth rates when possible. This is likely to reduce the mismatch of classifier and regressor seen around stability transitions in this study and potentially remove the need for a classifier altogether.
# ACKNOWLEDGMENTS

We thank Lorenzo Zainisi for insightful discussions on machine learning improvements and active learning; Karel van de Plassche, Jonathan Citrin, and Aaron Ho for their assistance with QLK and the QLK dataset; Feda Almuhisen for valuable suggestions to enhance the machine learning models, and Harry Dudding for discussions on current metrics and potential future improvements. This work was funded by the Agence Nationale de la Recherche (ANR), project ANR-22-CE30-0023. This work has been carried out within the framework of the EUROfusion Consortium, funded by the European Union via the Euratom Research and Training Programme (Grant Agreement No 101052200 — EUROfusion). Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Commission. Neither the European Union nor the European Commission can be held responsible for them. For the purposes of open access, the authors have applied a CC-BY public copyright licence to any Author Accepted Manuscript (AAM) version arising from this submission.

# AUTHOR DECLARATIONS

## Conflict of Interest

The authors have no conflicts to disclose.

## Author Contributions

**Matisse Lanzarone:** Conceptualisation (equal), Data Curation (lead), Formal Analysis (lead), Investigation (lead), Methodology (lead), Software (lead), Writing - original draft (lead), Writing - review and editing (equal)

**Emil Fransson:** Formal Analysis (supporting), Methodology (supporting), Writing - original draft (supporting), Writing - review and editing (equal)

**Yann Camenen:** Conceptualisation (equal), Formal Analysis (supporting), Funding Acquisition (lead), Methodology (supporting), Resources (lead), Writing - original draft (supporting), Supervision (equal), Writing - review and editing (equal)

**Guillaume Fuhr:** Conceptualisation (equal), Formal Analysis (supporting), Methodology (supporting), Supervision (equal), Writing - original draft (supporting), Writing - review and editing (equal)

# DATA AVAILABILITY STATEMENT

The data that support the findings of this study are available from the corresponding author upon reasonable request.

<footnote n="1">¹F. Poli, “Integrated tokamak modeling: When physics informs engineering and research planning,” Physics of Plasmas **25**, 055602 (2018).</footnote>

<footnote n="2">²M. Romanelli, G. Corrigan, V. Parail, S. Wiesen, R. Ambrosino, P. D. S. A. Belo, L. Garzotti, D. Harting, F. Köchl, T. Koskela, et al., “Jintrac: a system of codes for integrated simulation of tokamak scenarios,” Plasma and Fusion research **9**, 3403023–3403023 (2014).</footnote>

<footnote n="3">³F. Felici, O. Sauter, S. Coda, B. Duval, T. Goodman, J.-M. Moret, J. Palety, and the TCV Team, “Real-time physics-model-based simulation of the current density profile in tokamak plasmas,” Nuclear Fusion **51**, 083052 (2011).</footnote>

<footnote n="4">⁴D. Campbell, A. Loarte, D. Boilson, X. Bonnin, P. de Vries, L. Giancarli, Y. Gribov, S. Kim, M. Lehnen, T. Luce, I. Nunes, A. Polevoi, S. Pinches, R. Pitts, R. Reichle, M. Schneider, J. van der Laan, and G. Vayakis, “Iter research plan within the staged approach (level iii – final version),” Technical Report ITR-24-005 (ITER Organization, 2024) accessed: 2025-07-16. https://www.iter.org/sites/default/files/media/2024-04/itr-24-005-ok1.pdf.</footnote>

<footnote n="5">⁵A. Peeters, Y. Camenen, F. Casson, W. Hornsby, A. Snodin, D. Strintzi, and G. Szepesi, “The nonlinear gyro-kinetic flux tube code gkw,” Computer Physics Communications **180**, 2650–2672 (2009), 40 YEARS OF CPC: A celebratory issue focused on quality software for high performance, grid and novel computing architectures.</footnote>

<footnote n="6">⁶M. Kotschenreuther, G. Rewoldt, and W. Tang, “Comparison of initial value and eigenvalue codes for kinetic toroidal plasma instabilities,” Computer Physics Communications **88**, 128–140 (1995).</footnote>

<footnote n="7">⁷M. Barnes, D. Dickinson, W. Dorland, P. A. Hill, J. T. Parker, C. M. Roach, M. Giacomini, N. Mandell, R. Numata, T. Schuett, S. Biggss-Fox, N. Christen, J. Parisi, G. Wilkie, L. Anton, J. Ball, J. Baumgaertel, G. Colyer, M. Hardman, J. Hein, E. Highcock, G. Howes, A. Jackson, M. T. Kotschenreuther, J. Lee, H. Leggate, A. Mauriya, B. Patel, T. Tatsuno, and F. Van Wyk, “Gs2 v8.2.1,” https://doi.org/10.5281/zenodo.14022029 (2024), 10.5281/zenodo.14022029.</footnote>

<footnote n="8">⁸J. Candy and R. Waltz, “An eulerian gyrokinetic-maxwell solver,” Journal of Computational Physics **186**, 545–581 (2003).</footnote>

<footnote n="9">⁹G. Staebler, C. Bourdelle, J. Citrin, and R. Waltz, “Quasilinear theory and modelling of gyrokinetic turbulent transport in tokamaks,” Nuclear Fusion **64**, 103001 (2024).</footnote>

<footnote n="10">¹⁰J. Citrin, C. Bourdelle, F. J. Casson, C. Angioni, N. Bonanomi, Y. Camenen, X. Garbet, L. Garzotti, T. Görler, O. Gürcan, F. Koechl, F. Imbeaux, O. Linder, K. v. d. Plassche, P. Strand, G. Szepesi, and J. Contributors, “Tractable flux-driven temperature, density, and rotation profile evolution with the quasilinear gyrokinetic transport model qualikiz,” Plasma Physics and Controlled Fusion **59**, 124005 (2017).</footnote>

<footnote n="11">¹¹G. M. Staebler, J. E. Kinsey, and R. E. Waltz, “A theory-based transport model with comprehensive physicsa),” Physics of Plasmas **14**, 055909 (2007), https://pubs.aip.org/aip/pop/article-pdf/doi/10.1063/1.2436852/13888907/055909_1_online.pdf.</footnote>

<footnote n="12">¹²G. M. Staebler, R. E. Waltz, J. Candy, and J. E. Kinsey, “New paradigm for suppression of gyrokinetic turbulence by velocity shear,” Phys. Rev. Lett. **110**, 055003 (2013).</footnote>

<footnote n="13">¹³K. L. van de Plassche, J. Citrin, C. Bourdelle, Y. Camenen, F. J. Casson, V. I. Dagnelie, F. Felici, A. Ho, S. Van Mulders, and J. Contributors, “Fast modeling of turbulent transport in fusion plasmas using neural networks,” Physics of Plasmas **27**, 022310 (2020), https://pubs.aip.org/aip/pop/article-pdf/doi/10.1063/1.5134126/15772521/022310_1_online.pdf.</footnote>

<footnote n="14">¹⁴A. Ho, J. Citrin, C. Bourdelle, Y. Camenen, F. J. Casson, K. L. van de Plassche, H. Weisen, and J. Contributors, “Neural network surrogate of QuaLiKiz using JET experimental data to populate training space,” Physics of Plasmas **28**, 032305 (2021), https://pubs.aip.org/aip/pop/article-pdf/doi/10.1063/5.0038290/12361366/032305_1_online.pdf.</footnote>

<footnote n="15">¹⁵J. Citrin, P. Trochim, T. Goerler, D. Pfau, K. L. van de Plassche, and F. Jenko, “Fast transport simulations with higher-fidelity surrogate models for iter,” Physics of Plasmas **30**, 062501 (2023), https://pubs.aip.org/aip/pop/article-pdf/doi/10.1063/5.0136752/17925119/062501_1_5.0136752.pdf.</footnote>

<footnote n="16">¹⁶S. Maeyama, M. Honda, E. Narita, and S. Toda, “Multi-fidelity information fusion for turbulent transport modeling in magnetic fusion plasma,” Nature (2024), 10.21203/rs.3.rs-4791765/v1.</footnote>

<footnote n="17">¹⁷Y. Camenen and G. Fuhr, “Gkdb documentation,” https://gitlab.com/gkdb/gkdb/wikis/home (2025).</footnote>
¹⁸J. Citrin, S. Breton, F. Felici, F. Imbeaux, T. Aniel, J. Artaud, B. Baiocchi, C. Bourdelle, Y. Camenen, and J. Garcia, "Real-time capable first principle based modelling of tokamak turbulent transport," Nuclear Fusion 55, 092001 (2015).

¹⁹O. Meneghini, S. Smith, P. Snyder, G. Staebler, J. Candy, E. Belli, L. Lao, M. Kostuk, T. Luce, T. Luda, J. Park, and F. Poli, "Self-consistent core-pedestal transport simulations with neural network accelerated models," Nuclear Fusion 57, 086034 (2017).

²⁰O. Meneghini, G. Snoep, B. Lyons, J. McClenaghan, C. Imai, B. Grier-son, S. Smith, G. Staebler, P. Snyder, J. Candy, E. Belli, L. Lao, J. Park, J. Citrin, T. Cortemiglia, A. Biwole, and S. Mordijck, "Neural-network accelerated coupled core-pedestal simulations with self-consistent transport of impurities and compatible with iter imas," Nuclear Fusion 61 (2020), 10.1088/1741-4326/abb918.

²¹E. Fransson, A. Gillgren, A. Ho, J. Borsander, O. Lindberg, W. Rieck, M. Åqvist, and P. Strand, "A fast neural network surrogate model for the eigenvalues of QuaLiKiz," Physics of Plasmas 30, 123904 (2023), https://pubs.aip.org/aip/pop/article-pdf/doi/10.1063/5.0174643/18266035/123904_1_5.0174643.pdf.

²²W. Hornsby, A. Gray, J. Buchanan, B. Patel, D. Kennedy, F. Casson, C. Roach, M. Lykkegaard, H. Nguyen, N. Papadimas, B. Fourcin, and J. Hart, "Gaussian process regression models for the properties of micro-tearing modes in spherical tokamak," (2023), arXiv:2309.09785 [physics.plasm-ph].

²³J. Quinonero-Candela and C. Rasmussen, "A unifying view of sparse approximate gaussian process regression," Journal of Machine Learning Research, v.6, 1935-1959 (2005) 6 (2005).

²⁴A. Ho, J. Citrin, C. Challis, C. Bourdelle, F. Casson, J. Garcia, J. Hobirk, A. Kappatou, D. Keeling, D. King, F. Koechl, E. Lerche, C. Maggi, J. Mail-loux, and J. Contributors, "Predictive jet current ramp-up modelling using qualikiz-neural-network," Nuclear Fusion 63, 066014 (2023).

²⁵I. Casiraghi, P. Mantica, F. Koechl, R. Ambrosino, B. Baiocchi, A. Castaldo, J. Citrin, M. Dicorato, L. Frassinetti, A. Mariani, P. Vincenzi, P. Agostinetti, L. Aucone, L. Balbinot, S. Ceccuzzi, L. Figini, G. Granucci, P. Innocente, T. Johnson, H. Nyström, and M. Valisa, "First principle-based multi-channel integrated modelling in support of the design of the divertor tokamak test facility," Nuclear Fusion 61, 116068 (2021).

²⁶S. Van Mulders, F. Felici, O. Sauter, J. Citrin, A. Ho, M. Marin, and K. L. van de Plassche, "Rapid optimization of stationary tokamak plasmas in rap-tor: demonstration for the iter hybrid scenario with neural network surrogate transport model qlknn," Nuclear Fusion 61 (2021), 10.1088/1741-4326/ac0d12.

²⁷L. Grinsztajn, E. Oyallon, and G. Varoquaux, "Why do tree-based models still outperform deep learning on tabular data?" (2022), arXiv:2207.08815 [cs.LG].

²⁸D. Beam and M. Schramm, "Rossmann store sales," https://cs229.stanford.edu/proj2015/205_report.pdf (2015).

²⁹K. Belkhayat, A. Omar, G. Bruner, Y. Wang, and R. Wattenhofe, "Xgboost and lgbm for porto seguro's kaggle challenge: A comparison semester project," https://pub.tik.ee.ethz.ch/students/2017-HS/SA-2017-98.pdf (2018).

³⁰V. Sandulescu and M. Chiru, "Predicting the future relevance of research institutions - the winning solution of the kdd cup 2016," (2016), arXiv:1609.02728 [cs.LG].

³¹I. Shavitt and E. Segal, "Regularization learning networks: deep learning for tabular datasets," Advances in Neural Information Processing Systems 31 (2018).

³²A. Kadra, M. Lindauer, F. Hutter, and J. Grabocka, "Well-tuned simple nets excel on tabular datasets," Advances in neural information processing systems 34, 23928-23941 (2021).

³³F. Chollet, "Keras," https://keras.io (2015).

³⁴T. Chen and C. Guestrin, "XGBoost: A scalable tree boosting system," in *Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining*, KDD '16 (ACM, New York, NY, USA, 2016) pp. 785-794.

³⁵F. Imbeaux, S. Pinches, J. Lister, Y. Buravand, T. Casper, B. Duval, B. Guillerminet, M. Hosokawa, W. Houlberg, P. Huynh, S. Kim, G. Manduchi, M. Owsiaik, B. Palak, M. Plociennik, G. Rouault, O. Sauter, and P. Strand, "Design and first applications of the iter integrated modelling & analysis suite," Nuclear Fusion 55, 123006 (2015).

³⁶Y. Camenen, F. Imbeaux, G. Fuhr, and et al., "Design and first applications of the imas interface data structure for flux-tube gyrokinetic simulations," Submitted to Open Plasma Science (2024).

³⁷A. Ho, "Qualikiz-v2.6.2 linear instability spectra based on jet experimental plasma profiles," (2021).

³⁸B. J. J. Kremers, J. Citrin, A. Ho, and K. L. van de Plassche, "Two-step clustering for data reduction combining <scp>dbscan</scp> and k-means clustering," Contributions to Plasma Physics 63 (2023), 10.1002/ctpp.202200177.

³⁹M. L. E. Vergnaud, "Qlk to imas conversion software," https://gitlab.com/gkdb/imas-gk-/-/tree/main/tools/qlk (2024).

⁴⁰L. Breiman, "Bagging predictors," Machine learning 24, 123-140 (1996).

⁴¹L. Breiman, "Bias, variance, and arcing classifiers," Technical Report 460, Statistics Department, University of California (2000).

⁴²J. Bergstra, D. Yamins, and D. D. Cox, "Making a science of model search," (2012), arXiv:1209.5111 [cs.CV].

⁴³I. Goodfellow, Y. Bengio, and A. Courville, *Deep Learning* (MIT Press, 2016) http://www.deeplearningbook.org.

⁴⁴X. Glorot, A. Bordes, and Y. Bengio, "Deep sparse rectifier neural networks," in *Journal of Machine Learning Research*, Vol. 15 (2010).

⁴⁵M. HONDA and E. NARITA, "Development of a surrogate turbulent transport model and its usefulness in transport simulations," Plasma and Fusion Research 16, 2403002-2403002 (2021).

⁴⁶M. Lanzarone, Y. Camenen, G. Fuhr, F. Almuhisen, C. Bourdelle, J. Citrin, and K. L. van de Plassche, "Higher fidelity surrogate models for gyrokinetic simulations," in *IAEA Workshop on AI for Accelerating Fusion and Plasma Science*, Vienna (2023).

⁴⁷L. Zanisi, A. Ho, T. Madula, J. Barr, J. Citrin, S. Pamela, J. Buchanan, F. Casson, V. Gopakumar, and J. contributors, "Efficient training sets for surrogate models of tokamak turbulence with active deep ensembles," Nuclear Fusion 64 (2024), 10.1088/1741-4326/ad240d.

⁴⁸A. Gillgren, E. Fransson, A. Ludvig-Osipov, W. Enström, L. Flyckt, M. Green, M. Kvartsén, Y. Liljegren, E. Olsson, A. Orthag, H. Wennberg, and P. Strand, "Investigating characteristics of the growth rates from qualikiz using an interpretable surrogate model," Physics of Plasmas 32, 052306 (2025), https://pubs.aip.org/aip/pop/article/pdf/doi/10.1063/5.0261456/20528643/052306_1_5.0261456.pdf.

⁴⁹Y. Camenen, "Gkw to imas documentation," https://gitlab.com/gkdb/gkdb-/-/blob/master/doc/general/IOGKDB.pdf (2024).

⁵⁰B. Lakshminarayanan, A. Pritzel, and C. Blundell, "Simple and scalable predictive uncertainty estimation using deep ensembles," Advances in neural information processing systems 30 (2017).

⁵¹A. G. Peeters and D. Strintzi, "The effect of a uniform radial electric field on the toroidal ion temperature gradient mode," Physics of Plasmas 11, 3748-3751 (2004), https://pubs.aip.org/aip/pop/article/pdf/11/8/3748/19080468/3748_1_online.pdf.

⁵²A. G. Peeters, C. Angioni, and D. Strintzi, "Toroidal momentum pinch velocity due to the coriolis drift effect on small scale instabilities in a toroidal plasma," Phys. Rev. Lett. 98, 265003 (2007).

⁵³J. Friedman, "Greedy function approximation: A gradient boosting machine," The Annals of Statistics 29 (2000), 10.1214/aos/1013203451.

⁵⁴C. Bourdelle, J. Citrin, B. Baiocchi, A. Casati, P. Cottier, X. Garbet, F. Imbeaux, and J. Contributors, "Core turbulent transport in tokamak plasmas: bridging theory and experiment with qualikiz," Plasma Physics and Controlled Fusion 58, 014036 (2015).