# Efficient dataset construction using active learning and uncertainty-aware neural networks for plasma turbulent transport surrogate models

arXiv: https://arxiv.org/abs/2507.15976 · HTML: https://mlsanity.com/paper/2507.15976/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 78 references that are in the map.
- First posted on arXiv 21 Jul 2025.

## Builds on (most cited first, 18 of 18)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Deep Learning in Neural Networks: An Overview](https://mlsanity.com/paper/1404.7828/) (2014)
- [Evidential Deep Learning to Quantify Classification Uncertainty](https://mlsanity.com/paper/1806.01768/) (2018)
- [Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds](https://mlsanity.com/paper/1906.03671/) (2019)
- [A Survey of Deep Active Learning](https://mlsanity.com/paper/2009.00236/) (2020)
- [BatchBALD: Efficient and Diverse Batch Acquisition for Deep Bayesian Active Learning](https://mlsanity.com/paper/1906.08158/) (2019)
- [Simple and Principled Uncertainty Estimation with Deterministic Deep Learning via Distance Awareness](https://mlsanity.com/paper/2006.10108/) (2020)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [Why gradient clipping accelerates training: A theoretical justification for adaptivity](https://mlsanity.com/paper/1905.11881/) (2019)
- [Maximizing acquisition functions for Bayesian optimization](https://mlsanity.com/paper/1805.10196/) (2018)
- [A Review of Uncertainty Estimation and its Application in Medical Imaging](https://mlsanity.com/paper/2302.08119/) (2023)
- [A Simple Approach to Improve Single-Model Deep Uncertainty via Distance-Awareness](https://mlsanity.com/paper/2205.00403/) (2022)
- [Active Learning and Bayesian Optimization: a Unified Perspective to Learn with a Goal](https://mlsanity.com/paper/2303.01560/) (2023)
- [Multidimensional Uncertainty-Aware Evidential Neural Networks](https://mlsanity.com/paper/2012.13676/) (2020)
- [Addressing Deep Learning Model Calibration Using Evidential Neural Networks and Uncertainty-Aware Training](https://mlsanity.com/paper/2301.13296/) (2023)
- [Enhancing predictive capabilities in fusion burning plasmas through surrogate-based optimization in core transport solvers](https://mlsanity.com/paper/2312.12610/) (2023)
- [EuroPED-NN: Uncertainty aware surrogate model](https://mlsanity.com/paper/2402.00760/) (2024)

## Built on by (most cited first, 2 of 2)

- [TGLF-WINN: Data-Efficient Deep Learning Surrogate for Turbulent Transport Modeling in Fusion](https://mlsanity.com/paper/2509.07024/) (2025)
- [Machine Learning for Electron-Scale Turbulence Modeling in W7-X](https://mlsanity.com/paper/2511.04567/) (2025)

## Abstract (arXiv)

This work demonstrates a proof-of-principle for using uncertainty-aware architectures, in combination with active learning techniques and an in-the-loop physics simulation code as a data labeller, to construct efficient datasets for data-driven surrogate model generation. Building off of a previous proof-of-principle successfully demonstrating training set reduction on static pre-labelled datasets, using the ADEPT framework, this strategy was applied again to the plasma turbulent transport problem within tokamak fusion plasmas, specifically the QuaLiKiz quasilinear electrostatic gyrokinetic turbulent transport code. While QuaLiKiz provides relatively fast evaluations, this study specifically targeted small datasets to serve as a proxy for more expensive codes, such as CGYRO or GENE. The newly implemented algorithm uses the SNGP architecture for the classification component of the problem and the BNN-NCP architecture for the regression component, training models for all turbulent modes (ITG, TEM, ETG) and all transport fluxes (Qe, Qi, Gammae, Gammai, and Pii) described by the general QuaLiKiz output. With 45 active learning iterations, moving from a small initial training set of 102 to a final set of 104, the resulting models reached a F1 classification performance of ~0.8 and a R2 regression performance of ~0.75 on an independent test set across all outputs. This extrapolates to reaching the same performance and efficiency as the previous ADEPT pipeline, although on a problem with 1 extra input dimension. While the improvement rate achieved in this implementation diminishes faster than expected, the overall technique is formulated with components that can be upgraded and generalized to many surrogate modeling applications beyond plasma turbulent transport predictions.