A transformer-FSL-GAN framework for multiclass water quality classification under limited data conditions

Citation

Rahman, Ashikur and Chung, Gwo Chin and Ng, Yin Hoe (2026) A transformer-FSL-GAN framework for multiclass water quality classification under limited data conditions. Frontiers in Water, 8. ISSN 2624-9375

[img] Text
frwa-8-1902772.pdf - Published Version
Restricted to Repository staff only

Download (4MB)

Abstract

Inroduction: Precise water quality classification is critical for sustainable water quality monitoring and water resource management. However, most machine learning models depend on large amounts of labelled data, which are often difficult to obtain in real-world applications where water quality observations are sparse and likely class imbalance. In this study, a transformer-based few shot learning enhanced with the Wasserstein Generative Adversarial Network (Transformer-FSL-GAN) framework is proposed for the multiclass water quality classification problem when the number of data samples is limited.Methods: The proposed framework includes three complementary components: (i) a Tabular Transformer encoder that learns global nonlinear relationships among physicochemical parameters with multi-head self-attention; (ii) a Prototypical Network that enables metric-based few-shot classification in a discriminative embedding space; and (iii) class-specific Wasserstein Generative Adversarial Networks with Gradient Penalty (WGAN-GP) that generate realistic synthetic support samples to reduce data scarcity and class imbalance. The Canadian Surface Water Quality dataset was used for experiments, which included 3,949 samples, 8 physicochemical parameters, and 5 quality classes from the Canadian Centre for Measurement, Evaluation and Management of Water Quality (CCME-WQI): Excellent, Good, Fair, Marginal and Poor. We have assessed the proposed framework in 1, 5 and 10 shots learning condition and benchmarked it against the traditional machine learning models, deep learning baselines, and ablation variants.Results: The Transformer-FSL-GAN achieved the highest overall accuracy of 97.33%, macro precision of 97.50%, macro recall of 97.33%, macro F1-score of 97.33%, macro ROC-AUC of 99.53% in the 10-shot configuration. The proposed framework showed large effect sizes in the paired t-tests compared with the evaluated baseline methods (p < 0.001) for all comparisons. The latent representations of the water quality classes showed good separation in the UMAP visualization. The most influential features that contribute to model predictions were also identified through SHAP and LIME analysis, which included features such as orthophosphate, ammonia, nitrate, dissolved oxygen, and biochemical oxygen demand.Discussion: The findings indicate that combining transformer-based representation learning, prototypical few-shot classification with WGAN-GP-based data augmentation is an effective approach to accurately classify water quality in severe data limitations while maintaining interpretability for the multiclass problem. The proposed framework offers a promising approach for data-scarce environmental monitoring applications and could support the development of intelligent and AIoT-enabled water quality management systems.

Item Type: Article
Uncontrolled Keywords: water quality classification, transformer, few-shot learning, WGAN-GP, explainable artificial intelligence, data scarcity, smart water quality monitoring, water resource management
Subjects: T Technology > T Technology (General)
Divisions: Faculty of Artificial Intelligence & Engineering (FAIE)
Depositing User: Ms Suzilawati Abu Samah
Date Deposited: 01 Oct 2026 07:15
Last Modified: 01 Oct 2026 07:15
URII: http://shdl.mmu.edu.my/id/eprint/16790

Downloads

Downloads per month over past year

View ItemEdit (login required)