# Robust Confinement State Classification with Uncertainty Quantification through Ensembled Data-Driven Methods

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [XGBoost: A Scalable Tree Boosting System](https://mlsanity.com/paper/1603.02754/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [Optuna: A Next-generation Hyperparameter Optimization Framework](https://mlsanity.com/paper/1907.10902/) (2019)
- [A Comprehensive Survey on Transfer Learning](https://mlsanity.com/paper/1911.02685/) (2019)
- [Deep Neural Networks and Tabular Data: A Survey](https://mlsanity.com/paper/2110.01889/) (2021)
- [xLSTM: Extended Long Short-Term Memory](https://mlsanity.com/paper/2405.04517/) (2024)
- [U-FNO -- An enhanced Fourier neural operator-based deep-learning model for multiphase flow](https://mlsanity.com/paper/2109.03697/) (2021)
- [FourCastNet: Accelerating Global High-Resolution Weather Forecasting using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2208.05419/) (2022)
- [The Road Less Scheduled](https://mlsanity.com/paper/2405.15682/) (2024)
- [Uncertainty Quantification and Deep Ensembles](https://mlsanity.com/paper/2007.08792/) (2020)
- [Multivariate Confidence Calibration for Object Detection](https://mlsanity.com/paper/2004.13546/) (2020)
- [Improving Adversarial Robustness of Ensembles with Diversity Training](https://mlsanity.com/paper/1901.09981/) (2019)
- [Ensemble Methods as a Defense to Adversarial Perturbations Against Deep Neural Networks](https://mlsanity.com/paper/1709.03423/) (2017)
- [Plasma Surrogate Modelling using Fourier Neural Operators](https://mlsanity.com/paper/2311.05967/) (2023)
- [Hybrid deep learning architecture for general disruption prediction across tokamaks](https://mlsanity.com/paper/2007.01401/) (2020)
- [Fast Dynamic 1D Simulation of Divertor Plasmas with Neural PDE Surrogates](https://mlsanity.com/paper/2305.18944/) (2023)
- [Scalings for Tokamak Energy Confinement](https://mlsanity.com/paper/1910.02381/) (2019)

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

- [Principled Approaches for Extending Neural Architectures to Function Spaces for Operator Learning](https://mlsanity.com/paper/2506.10973/) (2025)
- [Plasma State Monitoring and Disruption Characterization using Multimodal VAEs](https://mlsanity.com/paper/2504.17710/) (2025)

## Abstract (arXiv)

Maximizing fusion performance in tokamaks relies on high energy confinement, often achieved through distinct operating regimes. The automated labeling of these confinement states is crucial to enable large-scale analyses or for real-time control applications. While this task becomes difficult to automate near state transitions or in marginal scenarios, much success has been achieved with data-driven models. However, these methods generally provide predictions as point estimates, and cannot adequately deal with missing and/or broken input signals. To enable wide-range applicability, we develop methods for confinement state classification with uncertainty quantification and model robustness. We focus on off-line analysis for TCV discharges, distinguishing L-mode, H-mode, and an in-between dithering phase (D). We propose ensembling data-driven methods on two axes: model formulations and feature sets. The former considers a dynamic formulation based on a recurrent Fourier Neural Operator-architecture and a static formulation based on gradient-boosted decision trees. These models are trained using multiple feature groupings categorized by diagnostic system or physical quantity. A dataset of 302 TCV discharges is fully labeled, and will be publicly released. We evaluate our method quantitatively using Cohen's kappa coefficient for predictive performance and the Expected Calibration Error for the uncertainty calibration. Furthermore, we discuss performance using a variety of common and alternative scenarios, the performance of individual components, out-of-distribution performance, cases of broken or missing signals, and evaluate conditionally-averaged behavior around different state transitions. Overall, the proposed method can distinguish L, D and H-mode with high performance, can cope with missing or broken signals, and provides meaningful uncertainty estimates.