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

1 Introduction

Water quality monitoring is one of the core elements of environmental sustainability, public health, agricultural productivity and ecosystem conservation. Agriculture consumes a large volume of freshwater resources on the planet, and it is important to ensure a reliable water quality evaluation as a stable way to irrigate land and secure a stable food production in the long term (Uddin et al., 2021; Kumar et al., 2024). Nevertheless, despite its critical importance, rising scarcity coupled with contamination in numerous regions remains a major global concern. Water pollution caused by industrial effluents, agricultural pollution, and urbanization has driven a high demand for precise and smart water quality evaluation systems (Chidiac et al., 2023). In addition, the clean drinking water and proper sanitation are still not accessible to millions of the global population and exposing them to serious health risks (Das, 2026). Traditional methods of laboratory-based analysis and rule-based water quality index (WQI) approaches can be reliable, but are generally time-consuming, labor-intensive and not suitable for real-time or at large-scale applications (Akhtar et al., 2021). Machine learning (ML) and deep learning (DL) methods have become increasingly popular in automated water quality prediction and classification in recent years. Support vector machines, random forests, artificial neural networks, and convolutional neural network models have proven to achieve good performance in analyzing physicochemical parameters of water. Nevertheless, the approaches usually necessitate large amounts of labelled data for stable generalization (Helaly et al., 2025; Zhu et al., 2022; Zhi et al., 2024). In practical water monitoring applications, particularly in rural areas, developing states or in newly installed systems, labeled data is often scarce, incomplete, or costly. This poses a significant limitation for the applicability of classical ML and DL models due to insufficient data availability (Campos et al., 2026; Lokman et al., 2025). Few-Shot Learning (FSL) has become a promising learning paradigm that enables learning from a limited number of labeled samples by leveraging transferability of representations and episodic training mechanisms. Metric-based FSL approaches, particularly prototypical networks, have been shown to have a high potential in low-data settings due to their simplicity and robustness. However, the current FSL methods for water quality analysis still exhibit a limited scope of representation and often fail to capture the complex interdependences among water quality indicators (Ren et al., 2018; Gu et al., 2025). Transformer-based architectures, built upon self-attention mechanisms, have recently shown enhanced effectiveness in learning long-range dependencies and feature interactions across various domains. The capability of learning contextual representations makes them efficient multivariate water quality data (Ergün, 2025). Transformer models, however, are also data-intensive models, and they tend to overfit when trained with small datasets, which limits their applicability in few-shot settings. To overcome data scarcity and data representation limits, generative models including generative adversarial networks (GANs) have been used to augment data. GANs are capable of producing realistic examples to supplement training data and enhance model robustness (Goyal and Mahmoud, 2024). Although they are effective, the majority of the existing studies on GAN-based water quality focus on binary classification tasks and they are not integrated with few-shot learning approaches. Besides that, most of the deep-learning models are black-box in nature with poor interpretability, limiting their usage in safety-critical applications. To address these challenges, an innovative Transformer-FSL-GAN framework is proposed for multiclass CCME-WQI water quality classification in limited-data scenarios. The proposed framework combines advanced feature encoding using Transformer with WGAN-GP for producing realistic class-conditioned synthetic samples to alleviate data scarcity and class imbalance, and Prototypical Networks with distance-based few-shot classification. The suggested framework is not an alternative to the analytical CCME-WQI methodology, but rather is intended to correctly classify the established classes of CCME-WQI using a relatively small number of labeled samples. Moreover, SHAP (SHapley Additive exPlanations) and Local Interpretable Model-agnostic Explanations (LIME) are included to enhance the model’s global and local interpretability, increasing transparency and enabling trustworthy environmental decisions.

This study has focused on the following main contributions:

• A novel Transformer-FSL-GAN framework is proposed for CCME-WQI multiclass water quality classification under the limited labeled data scenario.
• The episodic few-shot learning framework is embedded with WGAN-GP-based data augmentation to tackle data scarcity and class imbalance by synthesizing class-conditioned realistic samples.
• Transformer-based feature embeddings are used to learn complex, nonlinear relations and global interactions between physico-chemical water quality parameters.
• Both SHAP and LIME are integrated to give global-level and instance-level explanations, which makes the proposed model transparent and interpretable.
• The proposed framework consistently achieves superior performance compared to the traditional machine learning models, conventional few-shot learning models, Transformer-based baselines, and ablation variants through extensive experiments on the same few-shot evaluation protocols while retaining high-quality and interpretable performance on multiclass water quality classification.

2 Literature review

The evaluation of water quality is based on chemical analysis of water in laboratories and index-based evaluation programs, including the WQI. These techniques combine several physicochemical variables, such as pH, dissolved oxygen (DO), biochemical oxygen demand (BOD), turbidity, conductivity, and total dissolved solids (TDS) into a single scalar measure of overall water quality status (Patel et al., 2023; Najafzadeh et al., 2021). WQI-based methods are very popular among regulatory authorities for monitoring and decision-making due to their simplicity and interpretability (Fathi et al., 2018). However, they are extremely reliant on expert-driven thresholds, pre-set weighting and constant standards of regulation, which limit their adaptability to changing environmental conditions and geographical variations. Moreover, WQI-based approaches fail to capture nonlinear relationships among the water quality parameters and rely on labor-intensive sampling techniques, which hinder their applicability in real-time and large-scale monitoring systems.

ML models have been used more and more to predict and classify water quality automatically to overcome these limitations. Algorithms that have proven to be extremely predictive are support vector machines (SVM), random forests (RF), decision trees (DT), k-Nearest Neighbors (k-NN), logistic regression (LR), CatBoost, XGBoost, and multilayer perceptron (MLP) as they are capable of capturing nonlinear relationships in WQI-based data. Among these, ensemble-based models, particularly CatBoost and stacking models have achieved classification accuracies of up to 99.96% under controlled experimental conditions, which demonstrates the efficacy of boosting algorithms (Nasir et al., 2022). On the same note, more advanced models like LightGBM and XGBoost have reported accuracies of up to 99.65% while supporting feature importance analysis through explainable AI methods (Hridoy et al., 2025). A number of researchers also confirm the effectiveness of ML methods in related applications. ML-based classification of the irrigation water quality index (IWQI) has demonstrated high performance, with accuracies of up to 98.5% and is able to deal with class imbalance, enabling its deployment in low-cost IoT-based monitoring systems (Mustapha et al., 2025). Soft voting classifiers, among ensemble approaches, have enhanced classification reliability, and the performance of this approach has reached about 96.39% (Singh et al., 2025). Moreover, the optimized random forest models have proven to be robust in classification of groundwater quality classification, achieving up to 90% accuracy and demonstrating the ability to handle noise and feature interaction (Bhagat et al., 2026). Despite these promising results, ML models are still very reliant on large datasets that have been properly labeled. In real-world water monitoring applications, data collection is often limited by sampling cost, inadequate sensor deployment, missing values, and faulty labeling (Ningthoujam et al., 2026; Sharma et al., 2026; Jayaraman et al., 2024). As a result, ML models may suffer from overfitting and poor generalization when trained on small or imbalanced data. Moreover, traditional ML methods usually require manual feature engineering and are not adaptable to evolving environmental conditions, which restricts their applicability in dynamic environments.

Further improvements in predictive performance have been achieved through deep learning (DL) models which automatically learn hierarchical feature representations. Hybrid Convolutional neural networks (CNN) and long short-term memory (LSTM) architectures have also shown good results in real-time water quality monitoring, with reported accuracies of up to 98.40% (Chellaiah et al., 2024; Man et al., 2025). Similarly, CNN and LSTM have achieved accuracies of up to 97% in WQI prediction and lower error rates (Saroja and Dharshini, 2023). LSTM models have also been applied to contamination source identification and water quality forecasting in in semi-arid areas, achieving performance levels of about 89.07% (Wang et al., 2024). Moreover, more sophisticated architectures such as ConvLSTM have shown improved capability in modeling spatiotemporal relationships, attaining classification accuracies of up to 92% in water pollution applications (Zhu et al., 2023). However, DL models have several limitations in spite of their effectiveness. They require large amounts of labeled data to operate optimally, which is often unavailable in environmental monitoring scenarios. Furthermore, most of the existing literature is devoted to binary classification or single-parameter prediction tasks, which restrict their use to real-world multiclass water quality evaluation. In addition, their black-box nature diminishes trust and acceptance in regulatory and decision-support systems (Zhu et al., 2023; Ewuzie et al., 2022; Pandey et al., 2026).

Data augmentation techniques have been considered to overcome the problem of data scarcity and class imbalance. Conventional methods such as noise injection and resampling offer limited diversity and are likely to distort the original data distribution. In contrast, generative adversarial networks (GANs) provide a more powerful alternative by generating realistic synthetic samples that closely match real data (Robb et al., 2020). For example, hybrid frameworks combining GANs with CNN and explainable AI have achieved up to 97.86% accuracy while improving explainability (Rathinam, 2025). GAN-based augmentation has also been applied in related domains such as air quality prediction, hydrological modeling, and climate data synthesis (Radhakrishnan et al., 2026; Chen et al., 2026; Alzahem et al., 2023). Nevertheless, in water quality studies, GANs have mainly been applied in regression tasks or dataset balancing, with little attention paid to multiclass classification. Moreover, their integration with data-efficient methods, including few-shot learning, is only minimally investigated, which is a huge gap in research.

FSL is a promising technique for addressing data scarcity by enabling models to learn from a very small number of labeled samples. Prototypical networks, matching networks, and relation networks are metric-based methods that learn an embedding space in which samples of the same classes are tightly clustered and can be classified in terms of similarity measures (Rahman et al., 2026b; Snell et al., 2017). FSL has shown promising results in the environmental context. For instance, prototypical networks have achieved up to 93% accuracy in assessing the quality of groundwater under limited data conditions. Besides that, FSL with multi-objective optimization has also been applied to water resource management activities that have resulted in enhanced prediction accuracy, scalability, and efficiency (Derdour et al., 2025). Hybrid models, such as LSTM-based prototypical networks, have also demonstrated improved performance in runoff prediction under data-scarce settings (Zeng and Xiao, 2024). However, with these developments, current water quality research using FSL largely relies on shallow embedding networks, which are inadequate to capture the complex interdependencies of physicochemical parameters. Moreover, the majority of the literature is concerned with binary classification and does not cover the multiclass case or combines generative augmentation methods, which may restrict its practical applicability (Zhao et al., 2025; Yang et al., 2023).

Transformer models, built on self-attention networks, have recently proven to be highly effective in modeling long-range dependencies and complex feature interactions. These models have achieved state-of-the-art performance in various domains, such as natural language processing, computer vision, and time-series analysis (Lee et al., 2024), are increasingly applied to air pollution prediction, hydrological modeling, and water quality assessment. For instance, positional encoding-based models have been found to achieve good classification accuracy of up to 98.37% by effectively encoding feature interactions (Nejad et al., 2024). Transformer-based frameworks, including hybrid or ensemble models, have demonstrated strong predictive performance with high correlation measures, with correlation values reaching R2 up to 0.965 in wastewater monitoring (Rahman et al., 2026). CNN-Transformer architectures have enhanced dissolved oxygen prediction performance with reduced error rates (Srinath and Tuppad, 2025), while ALBERT-WPD has outperformed CNN and RNN models in water potability prediction, reaching accuracy of 96% (Xu et al., 2026). Nevertheless, transformer models remain underused in classifying the water quality under data-scarce conditions because they require large datasets and are computationally expensive (Rejini et al., 2025). Furthermore, the integration of transformer-based feature extraction with few-shot learning frameworks is an unexplored research area with strong potential for further development.

Lastly, explainable Artificial Intelligence (XAI) has become increasingly important for improving transparency and trust in AI-based water quality assessment systems. Techniques such as SHAP and LIME provide interpretable insights into the contributions of the features and model predictions, which allow domain experts to understand and confirm the findings. Although XAI has been used extensively with traditional ML models, including random forests and gradient boosting machines (Choudhary et al., 2025; Talukdar et al., 2024), little work has been done to apply it to deep learning and few-shot learning models. This lack of interpretability may hinder the practical implementation of more sophisticated AI models in regulatory and safety-sensitive water management systems (Makumbura et al., 2024; Infant et al., 2025; Rahman et al., 2026a). The critical research gaps identified from the literature review and summarized in Table 1 highlight several limitations in the field of water quality prediction. First, FSL techniques have a low level of adoption when it comes to water quality classification in environments with severe data scarcity, which is typical in the real world. Second, transformer-based representation learning has not been widely integrated into FSL systems for environmental datasets with complex feature interactions. Third, generative models like GANs have demonstrated promising performance in data augmentation, their integration into pipelines of few-shot water quality classification frameworks remain limited. Moreover, majority of the current literature concentrates on binary classification with little consideration given to multiclass water quality classification, despite its practical importance. In addition, the lack of interpretability in deep learning and FSL models constrains their use into the safety-critical applications. The study presents a novel transformer-based FSL framework integrated with GAN-based data augmentation and XAI. The limitations of the data scarcity are addressed through the proposed method, which employed self-attention mechanisms, generative modeling, and episodic learning strategies to improve feature level of representation and enable robust multiclass classification. The combination of explainability methods ensures the transparency and interpretability of the models, which is highly suitable in the real-world intelligent water quality monitoring systems.

Table 1. Comparison of existing water quality classification methods.

3 Methodology

The proposed methodology proposes a robust machine learning approach geared towards solving the problem of classification with a small amount of tabular water quality data. The pipeline consists of multiple stages, each representing a different stage of the process, as shown in Figure 1. The framework starts with a high-fidelity data engineer, which includes structured data collection, exploratory analysis, and extensive data preprocessing. The methodology addresses the issue of data scarcity and leverages the limited data to maximize model performance in this context by using a dual-stream optimization path: a typical machine learning pipeline and an optimized hybrid generative-few-shot learning pipeline. The advanced stream combines a Wasserstein Generative Adversarial Network (WGAN-GP) for targeted data augmentation with a multi-head self-attention Transformer Feature Encoder that captures complex feature dynamics, along with a Few-Shot Learning (FSL) model based on episodic sampling and distance-based metric classification. A framework aims to enhance the classification accuracy, generalization ability, and robustness in situations where labeled samples are limited, unbalanced, and heterogeneous.

A comparative experimental design was applied to rigorously assess the effectiveness of the proposed approach that involved three categories of models:

• Conventional ML classifiers, such as random forest, support vector machine, K-nearest neighbor, XGBoost, and logistic regression, which are traditional supervised learning methods that are trained on complete datasets.
• A baseline FSL-GAN, where the embedding encoder is an MLP with GAN-based synthetic data augmentation.
• A new transformer-FSL-GAN model, where the MLP encoder is substituted with a transformer-based architecture with self-attention to learn more complex nonlinear relationships between features and contextual information for water quality prediction.

3.1 Data acquisition

The surface water quality dataset used in this study was obtained from the publicly available repository entitled A Comprehensive Surface Water Quality Monitoring Dataset (1940–2023): 2.82 Million Record Resource for Empirical and Machine Learning-Based Research, which is hosted on Figshare under the persistent digital object identifier (DOI): 10.6084/m9.figshare.27800394 (Karim et al., 2025). The dataset is distributed under the Creative Commons Attribution 4.0 International (CC BY 4.0) license, permitting unrestricted use, distribution, and reproduction provided appropriate attribution is given. The complete repository contains approximately 2.82 million surface water quality measurements collected between 1940 and 2023 from five countries, namely the United States, Canada, Ireland, England, and China. Since this study aims to develop and evaluate a water quality classification framework for limited-data scenarios within a consistent environmental setting, only the Canadian subset (Canada_dataset.csv) from Version 2 of the repository, released on 23 February 2025, was selected for analysis. Restricting the study to a single country minimizes variations arising from differences in national monitoring protocols, regulatory standards, climatic conditions, and watershed characteristics, thereby providing a more homogeneous dataset for model development. The Canadian subset consists of 3,949 water quality observations collected from 2,354 environmental monitoring stations between 1968 and 2021. The observations represent four major waterbody types, namely rivers, lakes, transitional waters, and coastal waters, and include long-term physicochemical measurements obtained through routine environmental monitoring programs. Each observation contains eight physicochemical water quality parameters together with the corresponding Canadian Council of Ministers of the Environment Water Quality Index (CCME-WQI) score and water quality classification. This dataset serves as the foundation for developing and evaluating the proposed Transformer-FSL-GAN framework for multiclass water quality classification under limited labeled data conditions.

3.1.1 Data description

The extracted Canadian dataset was provided in comma-separated values (CSV) format and comprised physicochemical water quality measurements collected from multiple environmental monitoring stations over different geographical locations and sampling periods. Each record included environmental metadata (Country, Area, Waterbody Type, and Date), eight physicochemical water quality parameters, the numerical Canadian Council of Ministers of the Environment Water Quality Index (CCME_Values), and the corresponding categorical water quality class (CCME_WQI). To develop a generalized machine learning model based solely on measurable water quality indicators, only the physicochemical variables were retained as predictive input features. The selected variables included Ammonia (NH₃), Five-Day Biochemical Oxygen Demand (BOD₅), Dissolved Oxygen (DO), Orthophosphate (PO₄³⁻), Potential Hydrogen (pH), Water Temperature, Total Nitrogen (N), and Nitrate (NO₃⁻). These parameters are widely recognized as key indicators of nutrient enrichment, organic pollution, and the ecological condition of surface waters, and they are routinely monitored in national water quality assessment programs. The environmental metadata (Country, Area, Waterbody Type, and Date) were retained only for data exploration and dataset organization and were excluded from model development because they do not represent physicochemical measurements and may introduce location- or time-specific bias. Likewise, the numerical CCME_Values variable was used exclusively for descriptive statistical analysis and verification of the assigned water quality categories. Since the categorical target variable (CCME_WQI) is directly derived from the numerical CCME-WQI score, the CCME_Values variable was removed before model training to eliminate potential label leakage. Table 2 summarizes the descriptive statistics of the eight physicochemical variables together with the numerical CCME-WQI scores. The dataset contains 3,949 observations and exhibits substantial variability across several water quality indicators, reflecting a broad spectrum of environmental conditions ranging from relatively pristine to heavily polluted waters. The average pH was 7.68 ± 0.63, indicating generally neutral water conditions, while the mean water temperature was 12.42 ± 4.07 °C. Dissolved oxygen concentrations averaged 9.89 ± 1.44 mg/L, suggesting that most sampling locations maintained adequate oxygen availability for aquatic ecosystems. Pollution-related parameters, including ammonia, BOD₅, orthophosphate, total nitrogen, and nitrate, exhibited positively skewed distributions characterized by relatively low median concentrations but substantially larger maximum values. This distribution indicates that the dataset contains predominantly low-pollution observations together with a limited number of highly contaminated water samples, reflecting the natural variability of surface water quality. As shown in Table 3, the numerical CCME-WQI scores ranged from 37.65 to 100.00, with a median value of 93.11, indicating that most water samples were classified as Good or Excellent. Furthermore, the categorical class distribution was highly imbalanced, with Good (52.19%) and Excellent (35.48%) accounting for the majority of observations, whereas the Fair (5.93%), Marginal (4.81%), and Poor (1.60%) classes were considerably underrepresented. Such class imbalance poses a significant challenge for conventional supervised learning algorithms, which generally require sufficient labeled examples for each class. Consequently, this study employs WGAN-GP for minority class augmentation and few-shot learning to improve the classification performance of underrepresented water quality categories under limited labeled data conditions.

Table 2. Statistical summary of the eight physicochemical variables and CCME-WQI numerical scores.

Table 3. Distribution of water quality classes.

Bold value express the total sample.

3.1.2 CCME_WQI computation and water quality classification

The Canadian Council of Ministers of the Environment Water Quality Index (CCME_WQI) is a widely adopted composite index for evaluating the overall status of surface water quality. The index integrates multiple physicochemical parameters into a single score ranging from 0 to 100, where higher values indicate better water quality. From Equations 1–7, the CCME_WQI is computed using three statistical factors that quantify the Scope (F₁), Frequency (F₂), and Amplitude (F₃) of deviations from prescribed water quality objectives.

The first factor, Scope (F₁), measures the percentage of water quality variables that fail to satisfy their guideline objectives at least once during the monitoring period and is computed as

F 1 = ( Number of failed parameters Total number of parameters ) × 100 (1)

Where a failed variable is any parameter whose measured value exceeds (or falls below, depending on the guideline) the prescribed objective.

The second factor, Frequency (F₂), represents the percentage of individual measurements that violate the corresponding guideline objectives and is calculated as

F 2 = ( Number of failed tests Total number of tests ) × 100 (2)

The third factor, Amplitude (F₃), measures the extent by which failed observations exceed the allowable limits. The excursion for each failed observation is calculated differently depending on whether the guideline specifies an upper or lower objective.

For parameters with a maximum allowable objective, the excursion for a failed observation is calculated as:

Excursion i = Failed test value i Objective valu e j − 1 (3)

For parameters with a minimum allowable objective, the excursion is calculated as:

Excursion i = Objective value j Failed test value i − 1 (4)

The normalized sum of excursions (NSE) is computed as:

NSE = ( ∑ i = 1 n Excursion i Total number of tests ) (5)

Finally,

F 3 = NSE 0.01 NSE + 0.01 (6)

The overall CCME_WQI is computed as:

CCMEWQI = 100 − F 1 2 + F 2 2 + F 3 2 1.732 (7)

Where 1.732 is the square root of three, which normalizes the index to the range of 0–100.

The resulting numerical score is converted into one of five water quality categories according to the official CCME guideline shown in Table 4.

Table 4. CCME-WQI classification categories.

In the present study, the downloaded Canadian dataset already contains two CCME-WQI variables: CCME_Values, representing the numerical water quality index, and CCME_WQI, representing the corresponding categorical class. These variables were generated by the original dataset authors following the official CCME-WQI methodology described above. Therefore, the CCME-WQI scores were not recalculated in this study. Instead, CCME_Values were used solely for descriptive statistical analysis and verification of the assigned water quality categories, whereas the categorical CCME_WQI labels served as the ground-truth target variable for multiclass classification. To prevent target leakage, both CCME_Values and all metadata fields (Country, Area, Waterbody Type, and Date) were excluded before model training, leaving only the eight physicochemical parameters as predictive input features.

3.2 Exploratory data analysis

The heatmap in Figure 2 shows the mean chemical contents in five groups of water quality, with the classification of the rest one showing much higher levels of ammonia, biochemical oxygen demand (BOD), orthophosphate and nitrate. Although the concentrations in the water sources of Excellent and Good are much lower than the average (green shading), the extreme positive deviations of the Poor category take up extreme values like 17.89 mg/L of orthophosphate and 24.18 mg/L of nitrate. Conversely, such parameters as pH and dissolved oxygen do not change greatly among the classes, which means that nutrient enrichment and organic loading are the most influential in distinguishing water quality in this dataset.

Figure 3 shows the kernel density distributions of physicochemical parameters important to the quality of water: ammonia, BOD, dissolved oxygen (DO), orthophosphate, pH, temperature, nitrogen, and nitrate in five CCME-WQI classes (Excellent, Good, Fair, Marginal, and Poor). The distributions of most of the variables are skewed towards the right, especially in ammonia, orthophosphate, nitrogen, and nitrate, meaning that the variables have extremely high-valued observations. Log-scale transformations also indicate that the parameters related to pollutants (ammonia, BOD, orthophosphate, and nitrate) move to higher levels with the worsening of water quality between Excellent and Poor. The opposite is observed with dissolved oxygen, with Excellent and Good classes having higher concentrations, whereas lower values are more prevalent in the marginal and poor conditions. The distributions of the pHs are relatively short and concentrated at neutral (approximately 7.8), which implies little variation and low discriminatory ability between the classes. Changes in temperature are uniform across classes, and small rises have been noted in lower-quality classes of water.

Boxplots of individual parameters in relation to water quality class are shown in Figure 4, which show the central tendencies and outliers. It has a clear monotonic rise in median values of ammonia, BOD, orthophosphate, and nitrate between the Excellent and Poor classes. These parameters also have significant variability and many high-value outliers in the degraded water scenario, which shows that there are occasional instances of pollution or sources of contamination. On the other hand, the median trend of dissolved oxygen shows a decreasing pattern as the level of pollution increases, which supports the idea of dissolved oxygen being an important indicator of the health of aquatic ecosystems. The occurrence of the outliers in all classes indicates a natural variability in the environment, but the dispersion in Marginal and Poor is much greater. Nitrogen exhibits an irregular distribution with extreme values. It can be attributed to measurement error or inhomogeneous pollution contributions

Figure 5 displays the Pearson correlation heatmap between all the features and the values of the CCME-WQI index. It is found that there are strong negative relationships between CCME-WQI and ammonia (r = −0.64) and orthophosphate (r = −0.55), meaning that a rise in the concentration of the said pollutants severely harms the quality of water. The water quality also has moderate negative correlations with nitrate (r = −0.30) and BOD (r = −0.26). On the contrary, dissolved oxygen has a weak positive relationship (r = 0.12) with CCME-WQI, which proves its positive contribution to the quality of water. Inter-feature correlations indicate that there is a significant positive relationship between ammonia and orthophosphate (r = 0.55), suggesting that the sources of pollution are similar, i.e., agricultural runoff or wastewater discharge. The other pairs of features exhibit weak associations, suggesting that they have rather independent effects on the overall water quality index

The Principal Component Analysis (PCA) plot in Figure 6 shows the distribution of water quality samples based on the CCME_WQI index, which represents a combined 26.0% of all the data variance. In the visualization, there is a clear horizontal gradient in Principal Component 1 (16.3%), where samples move to be either “Excellent” (dark purple) on the left or “Poor” (light green) on the right. Although there is a lot of overlap between the classes with higher quality (Excellent-Fair), the samples with poor water quality have a significantly greater dispersion and include several distinct clusters with higher PC1 values. This implies that although water of high quality is chemically characterized by a similar profile, the poor quality of water can be caused by different kinds of extreme environmental conditions or pollutants.

3.3 Data preprocessing

Prior to model development, a systematic preprocessing pipeline was applied to ensure data consistency while preventing information leakage. Initially, variables that were not intended for prediction were excluded from the learning process. These included the environmental metadata (Country, Area, Waterbody Type, and Date) and the numerical CCME_Values, since the latter is directly associated with the target class (CCME_WQI) and could introduce label leakage if retained during model training. Consequently, only the eight physicochemical parameters were used as predictor variables throughout all experiments. The final feature set consisted of Ammonia (NH₃), Five-Day Biochemical Oxygen Demand (BOD₅), Dissolved Oxygen (DO), Orthophosphate (PO₄³⁻), Potential Hydrogen (pH), Temperature, Total Nitrogen (N), and Nitrate (NO₃⁻). The categorical target variable (CCME_WQI) was encoded into five numerical labels corresponding to the official CCME-WQI classification categories including Excellent = 0, Good = 1, Fair = 2, Marginal = 3, Poor = 4. Although the selected Canadian dataset contained no missing observations, a median imputation strategy was incorporated into the preprocessing pipeline to ensure reproducibility and compatibility with future datasets that may contain incomplete measurements. Median imputation was selected because of its robustness against skewed environmental variables and extreme outliers. To improve feature distribution and stabilize model optimization, the predictor variables were transformed using a Quantile Transformer, which maps the original feature distributions to an approximately Gaussian distribution. Compared with conventional min-max normalization or z-score standardization, quantile transformation is less sensitive to highly skewed environmental variables and effectively compresses extreme observations while preserving their relative ordering. After transformation, the predictor variables were standardized using z-score normalization to produce zero-centered features with unit variance. Standardization is defined as Equation 8

z i = x i − μ σ (8)

Where x i represents the original feature value, μ is the mean, and σ is the standard deviation.

To eliminate data leakage, all preprocessing operations, including median imputation, quantile transformation, and feature standardization, were fitted exclusively using the training data and subsequently applied to the evaluation data using the learned transformation parameters. No information from the evaluation subset was utilized during preprocessing.

3.3.1 Data splitting

To ensure a fair and leakage-free evaluation under limited-data conditions, the preprocessed dataset was partitioned into two mutually exclusive subsets prior to model development. A stratified random split was performed according to the CCME-WQI class labels to preserve the original class distribution in each subset, particularly for the minority Poor class. The complete dataset consisted of 3,949 water quality observations. Approximately 70% (2,762 samples) were assigned to the Base Subset, while the remaining 30% (1,187 samples) formed the Evaluation Pool. The class distribution after splitting is summarized in Table 5.

Table 5. Class distribution after splitting dataset.

Bold value indicate the total distribution sample.

The Base Subset was used exclusively for representation learning and pre-training of the MLP and Transformer feature extractors. Importantly, the Base Subset was not used to train the WGAN-GP. No samples from the Evaluation Pool were accessed during this stage. The Evaluation Pool was reserved exclusively for few-shot evaluation. Fixed episodic tasks were generated from this subset using a balanced 5-way K-shot protocol. Three support configurations were investigated: 1-shot, 5-shot, and 10-shot, corresponding to one, five, and ten labeled support samples per class, respectively. To ensure identical evaluation conditions across all competing models, the query set size was fixed at 15 samples per class, resulting in 75 query samples for every episode regardless of the support size. Because the minority Poor class contained only 19 samples in the Evaluation Pool, the maximum support size was limited to 10-shot to prevent overlap between support and query sets. The support and query samples were selected without replacement whenever sufficient samples were available. A deterministic random seed (42) was used throughout dataset partitioning and episodic generation to guarantee complete experimental reproducibility. All baseline machine learning algorithms, deep learning models, few-shot learning methods, and the proposed Transformer–FSL–GAN framework were evaluated using exactly the same Evaluation Pool, identical support sets, and identical query sets. This protocol eliminates evaluation bias and ensures that observed performance differences originate from the learning algorithms rather than from differences in sampled evaluation tasks (Infant et al., 2025). The resulting support and query sizes for each episode are summarized in Table 6.

Table 6. Fixed episodic evaluation protocol.

3.4 Transformer-based feature encoder

The proposed model uses a transformer-based encoder, unlike traditional multilayer perceptron (MLP) encoders, which treat input features separately, to learn more intricate inter-feature interactions found in water quality data. The transformer architecture illustrates in Figure 7 allows for the modeling of global contexts by a self-attention mechanism, which dynamically learns the relationship between physicochemical parameters (Yang et al., 2023; Lee et al., 2024).

3.4.1 Self-attention mechanism

In Equation 9, the core operation of the transformer encoder is the scaled dot-product self-attention, defined as:

Attention ( Q , K , V ) = softmax ( QK T d k ) V (9)

Where query matrix, Q = XW Q, key matrix. K = XW K, value matrix, V = XW V, scaling factor = d k, and learnable weights are W Q , W K , W V .

This mechanism allows the model to assign adaptive importance weights to different features, enabling the capture of nonlinear dependencies among water quality parameters, such as pH, DO, and BOD.

3.4.2 Multi-head attention

In Equations 10, 11, to enhance representation capability, multi-head attention is employed, allowing the model to learn multiple types of feature relationships simultaneously:

MHA ( X ) = Concat ( hea d 1 , … , hea d h ) W O (10)

Where each attention head is computed as:

hea d i = Attention ( Q i , K i , V i ) (11)

W O is a learnable output projection matrix.

This design enables the encoder to capture diverse interaction patterns across physicochemical parameters from multiple representation subspaces.

3.4.3 Feature token embedding

In Equation 12, since the input consists of numerical water quality parameters, each feature is transformed into a learnable token embedding:

e i = W e x i + b e (12)

Where x i denotes the physicochemical parameter, e i represents the corresponding embedding vector, and W e and b e are learnable projection parameters. A special classification token, denoted as [CLS], is appended to the input sequence to aggregate global contextual information across all feature tokens.

3.4.4 CLS global representation

In Equations 13, 14, after passing through stacked Transformer encoder layers, the final representation is obtained from the [CLS] token:

z = TransformerEncode ( [ x CLS , e 1 , e , … , e n ] ) (13)

Final embedding:

h = z CLS (14)

Where h denotes the global latent embedding used for downstream few-shot classification. This embedding is further normalized and projected into a discriminative latent space by fully connected layers, enabling robust class separation. Overall, this design significantly enhances representation learning capability by enabling global context modeling, adaptive feature weighting, and robust interaction capture, thereby outperforming conventional MLP-based encoders in complex environmental datasets.

3.5 GAN-based data augmentation (WGAN-GP)

3.5.1 GAN framework

The limited availability of labeled water quality observations, particularly for the Fair, Marginal, and Poor CCME-WQI classes, poses a significant challenge for developing robust few-shot learning models. To mitigate the effects of severe data scarcity while preserving the statistical properties of the original data, a class-specific Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP) was employed to generate realistic synthetic support samples. WGAN-GP learns the underlying data distribution and generates synthetic observations that preserve the complex nonlinear relationships among physicochemical parameters. The overall augmentation framework is illustrated in Figure 8 (Radhakrishnan et al., 2026; Chen et al., 2026).

This study employed independent class-specific WGAN-GP models, where a separate generator and critic were trained for each water quality class within every few-shot episode. Consequently, each generator learned only the distribution of its corresponding class, enabling the synthesis of class-consistent water quality samples while avoiding interference from other categories. For each episode, the generator G transformed a latent random vector z sampled from a Gaussian distribution into a synthetic water quality sample Equations 15–17:

x ⌢ = G ( z ) (15)

Where, z ~ N ( 0 , I ) denotes a 16-dimensional latent noise vector, G ( ⋅ ) represents the generator network, and x ⌢ is the generated synthetic water quality sample containing the eight physicochemical parameters. The critic D ( ⋅ ) estimated the Wasserstein distance between the distributions of real and generated samples:

D ( x ) (16)

Where x denotes either a real or synthetic water quality sample.

3.5.2 Wasserstein objective function

Instead of the binary cross-entropy objective used in class-specific WGAN-GP optimizes the Wasserstein distance between the real and generated data distributions, resulting in significantly more stable training and reducing the likelihood of mode collapse. The critic objective is defined as:

L W = E x ~ P data [ D ( x ) ] − E x ⌢ ~ P g [ D ( x ⌢ ) ] (17)

Where P data denotes the distribution of real water quality samples, and P g represents the generator distribution.

3.5.3 Gradient penalty for Lipschitz constraint

In Equations 18–20, to enforce the 1-Lipschitz continuity condition required by the Wasserstein distance, a gradient penalty term is incorporated:

L GP = λ E x ˜ ( ∥ ∇ x ˜ D ( x ˜ ) ∥ 2 − 1 ) 2 (18)

Where x ˜ is sampled along the interpolation between real and generated samples, and λ is the gradient penalty coefficient.

The final critic loss becomes:

ℒ D = ℒ W + ℒ GP (19)

The generator is trained to minimize:

L G = − E x ˜ ~ P g [ D ( x ˜ ) ] (20)

The gradient penalty stabilizes adversarial training by preventing exploding gradients and improving convergence when only a few real support samples are available.

3.5.4 Generator and critic architecture

The class-specific WGAN-GP consisted of fully connected multilayer perceptron networks. The generator received a 16-dimensional Gaussian latent vector as input and produced an 8-dimensional synthetic feature vector corresponding to the selected physicochemical parameters. Its architecture consisted of four fully connected layers with dimensions.

16 → 64 → 128 → 64 → 8

Using ReLU activation functions in the hidden layers and a linear output layer. The critic accepted an eight-dimensional water quality feature vector and estimated its Wasserstein score using the architecture.

8 → 64 → 128 → 64 → 1

With LeakyReLU activation (α = 0.2) applied after each hidden layer.

3.5.5 Training strategy

The class-specific WGAN-GP was trained only on the support set of each few-shot episode. Base subset and query samples were never used during the adversarial training, thereby preventing information leakage. For every episode and every water quality class, an independent class-specific WGAN-GP model was optimized using the Adam optimizer with a learning rate of 0.0005, β 1=0.5, and β 2=0.9. The critic was updated five times for every generator update using a gradient penalty coefficient of 10, and each model was trained for 1,000 epochs. In the extreme 1-shot setting, only one real support sample was available for each class. To stabilize critic optimization, the single support instance was temporarily replicated during GAN training. These replicated samples were used only to facilitate optimization and were not treated as additional independent observations. After convergence, each class-specific generator produced 30 synthetic samples, which were combined with the original support samples to construct an augmented support set for prototype computation and classifier adaptation. Consequently, the synthetic-to-real ratios were 30:1, 6:1, and 3:1 for the 1-shot, 5-shot, and 10-shot configurations, respectively. The 1-shot setting was however interpreted with some caution, as only one real observation is enough to estimate the underlying class distribution, class-specific variability and data manifold at a rough level. Hence these samples are called model-based augmentation samples and not necessarily representative, realistic, or diverse samples of the underlying population. As a result, the 1-shot results are evaluated by the effect of the augmentation on the few-shot classification performance downstream, and the class-specific WGAN-GP is able to reconstruct the underlying class distribution.

3.6 Few-shot learning via metric learning

The proposed framework is based on a metric-based FSL approach with a Prototypical Network to overcome data scarcity often faced in environmental monitoring applications. In this method, class-discriminative representations are learned from a small number of labeled samples by performing classification in a learned embedding space. Figure 9 shows the FSL framework (Karim et al., 2025; Rahman et al., 2026b).

3.6.1 Episodic training strategy

In Equation 21, training is conducted using an episodic learning paradigm designed to simulate real-world low-data scenarios. Each training episode ε consists of two sets:

ε = { S , Q } (21)

Where S denotes the support set and Q represents the query set. In Equation 22, an N- way K- shot setting, the support set contains K labeled samples from each of N classes:

S = { ( x i , y i ) } i = 1 N × K (22)

The query set contains unlabeled samples drawn from the same N classes and is used to evaluate classification performance within the episode. The episodic training strategy encourages the model to learn transferable feature representations rather than memorizing specific training samples.

3.6.2 Prototype computation

In Equation 23, for each class c, a prototype vector is computed as the mean embedding of its support samples in the learned latent space:

P k = 1 ∣ S k ∣ ∑ ( x i , y i ) ∈ S k f θ ( x i ) (23)

Where P k represents the prototype of class k, S k denotes the set of support samples belonging to class k, and f θ ( ⋅ ) is the embedding function implemented by the Transformer encoder. The prototype serves as a representative centroid of each class in the feature space.

3.6.3 Distance-based classification

In Equation 24, Given a query sample x, its embedding is first computed as:

z q = f θ ( x q ) (24)

In Equation 25, the probability of assigning the query sample to class k is then calculated using a distance-based softmax function:

p ( y = k ∣ x ) = exp ( − d ( z , P k ) ) ∑ j exp ( − d ( z , P j ) ) (25)

Where d ( ⋅ )signifies distance metric.

3.6.4 Euclidean distance

In Equation 26, the Euclidean distance enables classification based on proximity to class prototypes rather than learning explicit decision boundaries.

d ( z q , P k ) = ‖ z q − P k ‖ 2 2 (26)

3.6.5 Training objective

In Equation 27, the model parameters are optimized by minimizing the negative log-likelihood loss over the query set:

L = − ∑ ( x , y ) ∈ Q log p ( y ∣ x ) (27)

This objective leads to embeddings of samples from the same class to cluster around their corresponding prototypes while maximizing separation between different classes.

3.7 Experiment setup

All experiments were implemented in Python 3.11 using the PyTorch 2.2 deep learning framework. Data preprocessing, statistical analysis, and classical machine learning models were developed using NumPy 1.26, Pandas 2.2, and Scikit-learn 1.4. The experiments were conducted using Google Colaboratory Pro, which provided an NVIDIA Tesla T4 GPU (16 GB VRAM), approximately 12.7 GB system RAM, and an Intel Xeon CPU running a Linux environment. Additional validation and code development were performed on a local workstation equipped with an Intel Core i7 processor, 16 GB RAM, and the Windows 11 operating system. To ensure experimental reproducibility, all stochastic processes, including dataset partitioning, episodic sampling, neural network initialization, and optimization, were initialized using a fixed random seed of 42. The MLP and Transformer feature encoders were first pretrained using the 70% Base Subset, which contained observations from all five CCME-WQI classes, while the remaining 30% Evaluation Pool was reserved exclusively for few-shot evaluation. Because the same five CCME-WQI classes were present in both subsets, the experimental protocol represents within-class few-shot adaptation under limited labeled data rather than conventional novel-class few-shot learning. During evaluation, 1, 5, or 10 labeled support observations per class were used for adaptation and prototype construction, while the corresponding query observations were reserved for performance evaluation. During episodic evaluation, all competing models were assessed using the same evaluation pool, identical support-query splits, and identical randomly generated episodes, ensuring a fair comparison among traditional machine learning, deep learning, few-shot learning, GAN-based augmentation, and the proposed Transformer-FSL-GAN framework. Performance was evaluated using Accuracy, Precision, Recall, F1-score, and macro-average ROC-AUC over the predefined few-shot tasks.

3.8 Hyperparameter settings

Table 7 summarizes the experimental settings and hyperparameters used throughout this study. Unless otherwise stated, the same configuration was applied to all baseline models and the proposed Transformer-FSL-GAN framework to ensure a fair comparison under identical few-shot evaluation protocols.

Table 7. Hyperparameter settings used in the proposed transformer-FSL-GAN framework.

3.9 Algorithm

In Table 8 illustrates the Pseudocode of proposed Transformer-FSL-GAN framework that consists of two sequential stages: WGAN-GP-based data augmentation and episodic few-shot metric learning. First, an independent WGAN-GP model is trained for each water quality class using only the support samples within an episode. The generator transforms a Gaussian latent vector into synthetic water quality observations, while the critic distinguishes real from generated samples by optimizing the Wasserstein distance with gradient penalty. The gradient penalty enforces the Lipschitz continuity constraint, stabilizes adversarial training, and mitigates mode collapse, which is particularly important when only a few real observations are available. After convergence, the trained generator produces additional synthetic support samples for each class, thereby alleviating class imbalance and increasing the diversity of the limited training data. In the second stage, episodic few-shot meta-learning is performed to mimic realistic low-data environmental monitoring scenarios. For each episode, a 5-way K-shot support set (K = 1, 5, or 10) and a fixed query set are sampled from the evaluation pool. The synthetic samples generated by the corresponding WGAN-GP models are combined with the real support samples to form an augmented support set. Each water quality sample, represented by eight physicochemical parameters, is projected into feature tokens and processed by the Transformer encoder, whose multi-head self-attention mechanism captures global dependencies among water quality variables. The encoder outputs a latent embedding for every sample, from which class prototypes are computed by averaging the embeddings of the support samples belonging to the same class. For each query sample, its latent embedding is compared with all class prototypes using the Euclidean distance. The negative distances are converted into class probabilities through a softmax function, and the class with the highest probability is selected as the predicted water quality category. During meta-training, the Transformer encoder parameters are optimized by minimizing the cross-entropy loss through backpropagation and the Adam optimizer. This episodic optimization encourages samples from the same class to cluster around their corresponding prototypes while maximizing the separation between different classes, enabling robust multiclass water quality classification under severely limited labeled data conditions.

Table 8. Pseudocode for transformer-FSL-GAN water quality classification.

Bold value indicate the annotation.

3.10 Traditional machine learning models

To provide a comprehensive benchmark, the proposed Transformer-FSL-GAN framework was compared with several widely used supervised machine learning algorithms for multiclass water quality classification. The evaluated models included Random Forest (RF), XGBoost, LightGBM, CatBoost, Support Vector Machine (SVM), K-Nearest Neighbor (KNN), Naïve Bayes (NB), and Decision Tree (DT). All models were implemented using the Scikit-learn and XGBoost, LightGBM, and CatBoost libraries in Python and were trained using the same training data to ensure a fair comparison. To maintain experimental consistency, identical preprocessing procedures were applied to all models, including missing-value handling, feature normalization, and label encoding. The dataset was divided into 70% for model training (2,762 samples) and 30% for evaluation (1,187 samples). During few-shot evaluation, every classifier was trained using the same fixed support sets and evaluated on the identical query sets generated from the evaluation pool. Three few-shot settings were considered: 1-shot, 5-shot, and 10-shot, corresponding to one, five, and ten labeled support samples per class, respectively. Each evaluation episode followed a 5-way classification task with 15 query samples per class (75 query samples per episode). To minimize sampling bias and improve statistical reliability, the reported performance represents the average results over 100 randomly generated evaluation episodes using a fixed random seed of 42. The hyperparameters of each baseline model were selected based on commonly recommended settings in the literature and preliminary tuning experiments. Random Forest employed 100 decision trees with the Gini impurity criterion and unlimited tree depth. XGBoost used 100 boosting estimators, a maximum tree depth of 6, learning rate of 0.1, and subsample ratio of 0.8. LightGBM was configured with 100 estimators, 31 leaves, a learning rate of 0.1, and maximum depth of 6. CatBoost utilized 100 boosting iterations, tree depth of 6, and a learning rate of 0.1. The Support Vector Machine employed the Radial Basis Function (RBF) kernel with C = 1.0 and gamma = “scale”. K-Nearest Neighbor used k = 5 with Euclidean distance. Naïve Bayes employed the standard Gaussian Naïve Bayes formulation, while the Decision Tree classifier used the Gini impurity criterion with unrestricted tree depth. These baseline configurations provide representative comparisons against the proposed Transformer-FSL-GAN framework while maintaining consistent experimental conditions across all evaluated models.

3.11 Model explainability

To improve the transparency and interpretability of the proposed Transformer-FSL-GAN framework, two complementary explainable artificial intelligence (XAI) techniques, namely SHapley Additive exPlanations (SHAP) and Local Interpretable Model-Agnostic Explanations (LIME), were employed. While the proposed model achieved high predictive performance, these methods provide insight into the contribution of individual physicochemical parameters toward the model’s predictions, thereby enhancing its reliability for environmental decision-making. SHAP was employed to quantify the contribution of each input feature to the prediction outcome based on cooperative game theory. SHAP assigns an importance value (Shapley value) to every feature by measuring its marginal contribution to the prediction relative to all possible feature combinations. The Shapley value for feature i is computed as (Choudhary et al., 2025) (Equations 28,29)

ϕ i = ∑ S ⊆ F { i } ∣ S ∣ ! ( ∣ F ∣ − ∣ S ∣ − 1 ) ! ∣ F ∣ ! [ f ( S ∪ { i } ) − f ( S ) ] (28)

Where F denotes the complete feature set, Srepresents a subset of features excluding featurei, and f ( · )is the prediction function of the trained model. The SHAP values satisfy local accuracy, consistency, and feature attribution properties, making them suitable for interpreting complex deep learning models. Global feature importance was obtained by computing the mean absolute SHAP value across all evaluation samples (Talukdar et al., 2024):

I i = 1 N ∑ j = 1 N ∣ ϕ i ( j ) ∣ (29)

Where N denotes the number of evaluated samples. Larger values indicate greater influence of the corresponding physicochemical parameter on the overall prediction. In addition to global importance ranking, class-specific SHAP analysis was performed to investigate how each parameter contributes to the prediction of the five CCME-WQI water quality classes.

While SHAP provides global feature importance, Local Interpretable Model-Agnostic Explanations (LIME) was employed to explain individual predictions made by the proposed Transformer-FSL-GAN framework. LIME approximates the behavior of the complex model around a selected query sample by fitting an interpretable local surrogate model. For a given prediction, perturbed samples are generated in the neighborhood of the original instance, and the surrogate model is optimized by minimizing (Makumbura et al., 2024)

ξ ( x ) = arg min g ∈ G L ( f , g , π x ) + Ω ( g ) (30)

Where f represents the original Transformer-FSL-GAN model, g denotes the interpretable surrogate model, L ( · ) measures the local approximation error weighted by the neighborhood kernel π x, and Ω ( g ) represents the complexity penalty imposed on the surrogate model. LIME provides feature contribution scores indicating whether each physicochemical parameter positively or negatively influences the predicted water quality class for an individual query sample. Unlike SHAP, which summarizes feature importance over the entire dataset, LIME explains a single prediction, thereby providing local interpretability and supporting model transparency for practical environmental monitoring applications. By combining SHAP and LIME, the proposed framework offers both global and local interpretability. SHAP identifies the overall importance of physicochemical parameters across the dataset, whereas LIME explains the reasoning behind individual predictions. This complementary analysis enhances the trustworthiness and practical applicability of the proposed Transformer-FSL-GAN framework for intelligent water quality assessment and environmental decision support.

3.11.1 Evaluation matrix

To evaluate the performance of the proposed models, several standard classification metrics were employed, including accuracy, precision, recall, F1-score, mean, and area under the receiver operating characteristic curve (AUC-ROC). In Equations 30–33 are derived from the four basic outcomes of a classification task: true positives (TP), true negatives (TN), false positives (FP), and false negatives (FN). The formulas are defined as follows:

Accuracy = TP + TN TP + TN + FP + FN (30)

Precision = TP TP + FP (31)

Recall = TP TP + FN (32)

F 1 − Score = 2 Precision · Recall Precision + Recall (33)

This comprehensive evaluation strategy ensures both classification reliability and statistical robustness of the proposed FSL framework in comparison to traditional machine learning approaches.

4 Results and discussion

4.1 Performance of traditional machine learning baselines

To establish benchmark performance under limited-data conditions, eight traditional machine learning classifiers were evaluated using the identical episodic evaluation protocol adopted for the proposed framework. The evaluated models included Random Forest (RF), XGBoost, LightGBM, CatBoost, SVM, KNN, NB, and DT. To ensure a fair comparison, every classifier was trained using the same support sets and evaluated on the same fixed query sets generated from the evaluation pool. Table 9 summarizes the performance of all baseline classifiers in terms of accuracy by different shot conditions.

Table 9. Performance comparison of traditional machine learning baselines.

Under the extremely limited 1-shot setting, all baseline models experienced substantial performance degradation. Random Forest, Support Vector Machine (SVM), K-Nearest Neighbor (KNN), and Naïve Bayes achieved the highest accuracy of 41.33%, whereas Decision Tree, CatBoost, and XGBoost exhibited lower predictive performance. These results indicate that conventional supervised learning algorithms struggle to generalize when only one labeled sample per class is available. When the support size increased to 5-shot, classification performance improved considerably. CatBoost produced the highest accuracy (76.00%) and macro F1-score (75.81%), followed by Random Forest and Decision Tree (both 69.33% accuracy). Gradient boosting methods such as XGBoost also demonstrated competitive performance, while LightGBM remained unstable under the limited-data setting. The best overall baseline performance was obtained under the 10-shot configuration, where CatBoost achieved an accuracy of 85.33%, followed by DT (81.33%) and XGBoost (80.00%). Random Forest and LightGBM also produced strong classification performance with accuracies above 75%. These findings indicate that traditional machine learning algorithms require a substantially larger number of labeled support samples before reliable decision boundaries can be learned. Overall, the results demonstrate that although conventional machine learning methods benefit from increasing support size, their performance deteriorates rapidly under severe data scarcity. This limitation motivates the use of few-shot learning approaches capable of learning transferable feature representations from only a handful of labeled examples, which is the primary objective of this study.

4.2 Performance of deep learning and few-shot learning models

Table 10 presents the performance comparison of the MLP baseline, Prototypical Network (ProtoNet), Transformer baseline, and the proposed Transformer + FSL model under the 1-shot, 5-shot, and 10-shot evaluation settings. All models were evaluated using the identical episodic protocol, where the same support sets and fixed query sets were sampled from the evaluation pool to ensure a fair comparison. As the number of labeled support samples increased, all models demonstrated improved classification performance. Under the most challenging 1-shot setting, the proposed Transformer + FSL achieved the highest accuracy (85.33%), macro F1-score (85.12%), and macro ROC-AUC (97.18%), substantially outperforming the MLP baseline (65.33%) and both the Transformer baseline and ProtoNet (74.67%). Similar improvements were observed in the 5-shot and 10-shot settings, where the proposed model consistently achieved the best overall performance with accuracies of 90.67 and 96.00%, respectively. The superior performance of the proposed framework can be attributed to the complementary strengths of Transformer-based feature extraction and metric-based few-shot learning. The self-attention mechanism effectively captures complex relationships among physicochemical parameters, while prototype-based classification enables robust generalization from limited labeled examples. Overall, these results demonstrate that integrating Transformer representations with Prototypical Networks provides a more discriminative embedding space than conventional MLP-based models, making the proposed approach particularly suitable for multiclass water quality classification under limited-data conditions.

Table 10. Performance comparison of deep learning and few-shot learning models under different few-shot settings.

Bold values indicate the best performance.

4.3 Ablation study of the proposed transformer-FSL-GAN

Figure 10 illustrates the classification accuracy of the four ablation models under the 1-shot, 5-shot, and 10-shot evaluation settings. As the number of labeled support samples increased, all models exhibited improved classification performance, demonstrating the positive effect of additional labeled examples. The proposed Transformer-FSL-GAN consistently achieved the highest accuracy across all shot configurations, improving from 82.67% in the 1-shot setting to 89.33% in the 5-shot setting and 97.33% in the 10-shot setting. Notably, under the strictest data constraint of the 1-shot setting, the FSL - GAN baseline outperformed both GAN Alone and Transformer - GAN, achieving approximately 73.33% accuracy and demonstrating the immediate value of metric learning when data is extremely scarce. However, as the shot configuration expanded to 5-shot and 10-shot regimes, the models incorporating Transformer-based feature learning exhibited much steeper performance scaling. The Transformer - GAN model showed substantial improvement, increasing from 66.67 to 96.00%, indicating that Transformer-based feature learning effectively captures complex relationships among physicochemical parameters when more support samples are available. In contrast, the architectures lacking Transformer components (GAN Alone and FSL - GAN) experienced rapid saturation, failing to pass the 85% accuracy threshold at 10-shot. Overall, the results demonstrate that integrating Transformer-based feature extraction, few-shot metric learning, and WGAN-GP augmentation provides complementary benefits, resulting in the most robust and accurate multiclass water quality classification under limited-data conditions.

4.4 Multi-class classification performance

Table 11 summarizes the multiclass classification performance of the proposed Transformer-FSL-GAN under the 10-shot evaluation setting one representative evaluation episode. The episode contained 75 query samples, with 15 query samples for each of the five CCME-WQI classes. The model correctly classified 73 of the 75 query samples, resulting in an episode-level accuracy and balanced accuracy of 97.33%. The macro precision, macro recall, and macro F1-score were 97.50, 97.33, and 97.33%, respectively, indicating balanced classification performance across the five classes. The macro ROC-AUC of 99.53% indicates strong discrimination among the CCME-WQI classes for this episode. The class-wise evaluation reveals that the proposed model achieved perfect classification for the Excellent water quality class, obtaining 100% precision, recall, F1-score, and ROC-AUC. Similarly, the Poor class achieved a perfect recall of 100% and a ROC-AUC of 1.0000, indicating that all highly polluted water samples were correctly identified. The Good and Marginal classes achieved perfect precision (100%) with slightly lower recall (93.33%), suggesting that only a small number of samples were misclassified into adjacent categories. Meanwhile, the Fair class achieved 100% recall and an F1-score of 96.77%, demonstrating robust recognition despite representing an intermediate water quality category. Overall, these results indicate that the proposed Transformer-FSL-GAN maintains high and well-balanced classification performance across both majority and minority CCME-WQI classes. The high macro-average metrics, together with consistently high class-specific ROC-AUC values, demonstrate that the integration of Transformer-based feature extraction, WGAN-GP data augmentation, and metric-based few-shot learning enables the model to learn highly discriminative representations even under limited-data conditions. The minimal performance variation among classes further confirms the robustness of the proposed framework for practical multiclass surface water quality assessment.

Table 11. Multi-class classification performance of the proposed Transformer-FSL-GAN for one representative 10-shot evaluation episode (75 query samples).

Bold values indicate the best performance.

4.5 Confusion matrix analysis

Figure 11 presents the confusion matrices of the proposed Transformer-FSL-GAN framework under the 1-shot, 5-shot, and 10-shot evaluation settings. The confusion matrices demonstrate a progressive improvement in classification performance as the number of labeled support samples increases. Across all shot configurations, most query samples are correctly assigned to their corresponding CCME-WQI classes, as indicated by the strong concentration of values along the main diagonal. Under the 10-shot configuration, the proposed model achieved an overall classification accuracy of 97.33% (73/75 correctly classified samples). The Excellent class was classified perfectly, achieving 100% precision and recall, with all 15 query samples correctly identified. The Fair and Poor classes also achieved 100% recall, although each received one false-positive prediction from an adjacent class. Minor confusion was observed only between the Good and Fair classes and between the Marginal and Poor classes, where one query sample in each case was assigned to the neighboring water quality category. These neighboring classes exhibit similar physicochemical characteristics and overlapping CCME-WQI score ranges, making such boundary misclassifications expected. Comparing the three shot settings reveals that the number of off-diagonal misclassifications decreases substantially as additional labeled support samples become available. This trend indicates that increasing support examples enables more accurate prototype estimation in the embedding space while improving the representation learned by the Transformer encoder. Furthermore, the integration of WGAN-GP-generated synthetic support samples enhances class diversity during prototype construction, thereby improving discrimination for minority classes. Overall, the confusion matrix analysis confirms that the proposed Transformer-FSL-GAN framework provides accurate and robust multiclass water quality classification under limited-data conditions while maintaining high performance across all five CCME-WQI categories.

4.6 Multi-class ROC evaluation and AUC analysis

Figure 12 presents Receiver Operating Characteristic (ROC) curves of the proposed Transformer-FSL-GAN under the 1-shot, 5-shot, and 10-shot evaluation settings. As the number of labeled support samples increased, the model exhibited progressively stronger discriminative capability, with the macro-average ROC-AUC improving from 97.33% in the 1-shot configuration to 98.80% in the 5-shot configuration and reaching 99.53% in the 10-shot configuration. The ROC curves consistently shifted toward the upper-left corner, indicating higher true positive rates and lower false positive rates across all five CCME-WQI water quality classes. Under the 10-shot setting, the Excellent and Poor classes achieved perfect discrimination with AUC values of 1.0000, while the Good and Marginal classes achieved 0.9933, and the Fair class achieved 0.9900. The slightly lower AUC values for the Good, Fair, and Marginal classes are consistent with the minor confusion observed between adjacent water quality categories in the confusion matrix, reflecting their similar physicochemical characteristics. Overall, the continuous improvement in ROC performance demonstrates that the proposed framework effectively integrates Transformer-based feature learning, WGAN-GP data augmentation, and prototype-based few-shot learning to produce highly discriminative latent representations, enabling accurate multiclass water quality classification even under limited-data conditions.

4.7 Latent space analysis

Figure 12 presents the UMAP visualization of the latent feature space learned by the proposed Transformer-FSL-GAN framework under the 1-shot, 5-shot, and 10-shot few-shot learning settings. Each subplot illustrates the distribution of support samples, query samples, synthetic samples generated by WGAN-GP, and the corresponding class prototypes for the five CCME-WQI water quality categories. UMAP was employed to project the high-dimensional feature embeddings into a two-dimensional space while preserving the local neighborhood structure of the learned representations. Under the 1-shot setting at Figure 13a, the model forms distinguishable clusters despite having only a single labeled support sample per class. Although the clusters are relatively sparse, the synthetic samples generated by WGAN-GP populate the surrounding feature space, enabling more reliable prototype estimation and reducing the effect of extreme data scarcity. Minor overlap is observed only between neighboring water quality categories, reflecting the inherent similarity of adjacent CCME-WQI classes. As the number of labeled support samples increases to the 5-shot configuration at Figure 13b, the latent representations become noticeably more compact and better separated. The additional support samples improve prototype estimation, while the synthetic samples further densify each class distribution without introducing significant overlap. Consequently, both support and query samples are positioned closer to their respective class prototypes, indicating improved feature consistency and stronger discriminative capability. The 10-shot configuration in Figure 13c exhibits the highest-quality latent representation, where all five water quality classes form highly compact and well-separated clusters around their corresponding prototypes. The synthetic samples closely follow the distribution of the real support samples, demonstrating that the WGAN-GP successfully preserves the underlying data manifold while increasing sample diversity. The minimal overlap among clusters is consistent with the confusion matrix and ROC analyses, which reported the highest classification accuracy (97.33%) and macro ROC-AUC (99.53%) under the 10-shot setting. Overall, the progressive improvement in cluster compactness and inter-class separation across the three shot configurations provides qualitative evidence that the integration of Transformer-based feature extraction, prototype-based few-shot learning, and WGAN-GP augmentation learns highly discriminative latent representations for multiclass water quality classification under limited-data conditions.

4.8 Robustness and statistical validation

Table 12 shows the statistical validation results for the proposed Transformer-FSL-GAN framework and all baseline models by paired t-tests, over 100 repeated few-shot test episodes. The proposed model was evaluated and outperformed all the considered traditional machine learning, deep learning and few shot learning baselines with the mean classification accuracy of 96.77%. As all models were evaluated under the same conditions, that is, using the same episode-specific support-query episodes, the episode-level accuracy values for the models could be directly compared.

Table 12. Statistical validation of the proposed model against baselines.

The proposed framework achieved a mean accuracy improvement of 1.17 percentage points over Transformer-FSL, increasing from 95.60 to 96.77% (t = 6.0145, p = 3.04 × 10⁻⁸, Cohen’s d = 0.6014). The result shows that using WGAN-GP augmentation in the Transformer-based few-shot framework is statistically significant. The other baselines were compared for further improvement. The proposed framework showed a mean accuracy increase of 2.28 percentage points, t = 11.8393, p = 1.13 × 10⁻²⁰, and Cohen’s d = 1.1839 compared to the Transformer Baseline. Compared with Transformer-GAN, the improvement was 9.17 percentage points (t = 48.0397, p = 2.06 × 10⁻⁷⁰, Cohen’s d = 4.8040), while the improvement over ProtoNet was 13.26 percentage points (t = 67.1947, p = 2.09 × 10⁻⁸⁴, Cohen’s d = 6.7195). These comparison results suggest that the feature extraction module, prototype-based metric learning module, and WGAN-GP augmentation module all contributed to the superior performance of the full-component network compared to the reduced-component network setting with the same test procedure. The proposed framework also recorded mean accuracy gains of 15.72, 21.10, 16.74, 30.37 and 34.18 percentage points when compared to the accuracy obtained from MLP, Random Forest, XGBoost, SVM, and KNN, respectively. For all paired comparisons, p < 0.001. These conventional and deep-learning baselines yielded Cohen’s d scores ranging from 7.6579 to 17.8598, suggesting that there was a large performance difference between the paired baselines when compared to the observed performance variability across the evaluation episodes. These effect sizes, however, should be viewed in the context of the repeated-episode evaluation design and should not be interpreted as the same as those from independent data or independent experimental replications. Importantly, the 100 evaluation episodes were repeatedly drawn from the same finite Evaluation Pool and thus were not seen as 100 independent datasets or independent experiments. Observations in the Evaluation Pool may be multiple episodes. The paired t-tests are therefore not intended to provide a statistical inference from 100 independent replications but are intended to quantify differences between the models in the repeated, matched evaluation tasks within the set few-shot protocol. Likewise, the paired correlations obtained (0.73–0.79) are between performances of the compared models during the same evaluation episodes. They are not viewed as evidence of independent replication and repeatability. Based on the statistical comparison, the Transformer-FSL-GAN could be seen as having higher mean episode-level accuracy than all considered baseline models for the given few-shot setting of the evaluation protocol. The stable performance level observed throughout the repeated matched trials is indicative of the strong performance of the proposed framework in the chosen episodic evaluation environment. The episodes were based on the same finite Evaluation Pool and so cannot be considered to represent generalization across independent datasets. Rather, they offer strong evidence for the performance benefits of the combination of Transformer-based feature learning, WGAN-GP-based support augmentation, and prototype-based few-shot classification under the limited-label classification of water quality conditions provided.

4.9 Explainability analysis

To enhance the transparency and interpretability of the proposed Transformer-FSL-GAN framework, explainability analysis was performed using SHAP (SHapley Additive exPlanations) and LIME (Local Interpretable Model-agnostic Explanations). SHAP was employed to provide global explanations by quantifying the contribution of each physicochemical parameter across the entire dataset, whereas LIME was used to explain the prediction of an individual query instance. Together, these methods offer comprehensive insights into both the overall model behavior and individual classification decisions. Figure 14a presents the SHAP summary (beeswarm) plot, which illustrates the distribution of SHAP values for all environmental variables. The results indicate that Orthophosphate is the most influential feature, followed by Ammonia, Nitrate, Dissolved Oxygen (DO), Biochemical Oxygen Demand (BOD), pH, Temperature, and Nitrogen. The wide spread of SHAP values demonstrates that each variable contributes differently depending on its measured value, highlighting the nonlinear relationships learned by the Transformer encoder. Positive SHAP values increase the confidence of a predicted class, whereas negative values decrease its contribution.

Figure 14b shows the global SHAP feature importance based on the mean absolute SHAP values. The ranking confirms that Orthophosphate contributes the most to the multiclass water quality classification, followed by Ammonia and Nitrate, while Nitrogen exhibits the smallest overall influence. The consistency between the SHAP summary and feature importance plots indicates that the proposed framework relies primarily on chemically meaningful water quality parameters when making classification decisions. To complement the global interpretation, Figure 15 presents a representative LIME explanation for an individual query sample. The local surrogate model highlights the contribution of each environmental variable toward the predicted water quality class by assigning positive and negative feature weights. The results show that Orthophosphate provides the strongest contribution to the prediction, while the remaining physicochemical parameters contribute with varying magnitudes depending on the local characteristics of the evaluated sample. The agreement between the SHAP and LIME analyses demonstrates that the proposed Transformer-FSL-GAN framework produces consistent and interpretable predictions rather than relying on spurious correlations. Overall, the explainability analysis confirms that the proposed framework is not only highly accurate but also transparent and trustworthy. The combination of SHAP for global feature interpretation and LIME for local decision explanation provides comprehensive evidence that the model learns meaningful relationships among physicochemical parameters, supporting its applicability for reliable water quality assessment and environmental decision-making.

4.10 Discussion

The experimental results prove that the proposed Transformer-FSL-GAN framework is consistently outperforming all baseline machine learning, deep learning and few-shot learning models in both of the evaluation settings. The integration of Transformer-based feature extraction, prototype-based few-shot learning, and WGAN-GP data augmentation enabled the framework to effectively address two major challenges in water quality classification: limited labeled data and class imbalance. The proposed framework outperformed all conventional classifiers in terms of classification accuracy, precision, recall, F1 score and ROC-AUC in all the 1-shot, 5-shot and 10-shot support settings, which demonstrated its capability of generalization with only a few labeled samples. The ablation study also showed that all the components performed positively in overall performance. While GAN-based data augmentation boosted the diversity of training samples and enhanced minority class representation, synthetic data alone was not enough to obtain the best classification performance. Likewise, the feature representation was greatly improved by the Transformer encoder for modeling long-range dependencies between physicochemical parameters with multi-head self-attention, and the few-shot classification was greatly facilitated by the Prototypical Network that learned discriminative metric spaces. The combination of these three components resulted in the most successful performance, indicating that they complement each other and can all be used to achieve good water quality classification. The results from the confusion matrix, the ROC analysis and latent space visualization further support the efficiency of the proposed framework. The confusion matrices indicated that the majority of the predictions were along the main diagonal, and there was a small amount of confusion between neighbouring CCME-WQI classes where the physicochemical properties overlap between the classes. The multi-class ROC curves showed good separation between all the water quality categories and the macro ROC-AUC value under the 10-shot configuration was excellent (0.9953). Additionally, the UMAP latent space visualization showed that larger-sized support groups led to better class separation, whereas the proposed Transformer-FSL-GAN produced small clusters with little overlap between classes, and it successfully separated the classes when compared to the baseline models. The learned embedding space shows good ability to represent those discriminative features that are useful for accurate prototype-based classification, based on these observations. The explainability analysis also gives valuable practical implications on the proposed model. Water quality prediction influencing variables that SHAP and LIME consistently identified were orthophosphate, ammonia, nitrate, dissolved oxygen and biochemical oxygen demand. The results are in line with known environmental studies that identify nutrient enrichment and oxygen related variables as key parameters for determining the health of an aquatic ecosystem. The corresponding relationship between global SHAP and local LIME interpretations improves transparency of the proposed framework and boosts the trust in the framework’s decision-making process, which becomes essential for environmental monitoring applications in the real world. The statistical validation also enhances the trustworthiness of the proposed approach. Paired t-tests revealed significant differences from all baseline models (p < 0.001); Cohen’s d values were all moderate to extremely large. The high paired correlation for each identical testing episode indicates that the observed gains are not likely to be due to random variation, but are consistent and repeatable. The results validate that the proposed Transformer-FSL-GAN is statistically and practically significant and improves the existing machine learning, deep learning, and few-shot learning approaches. The proposed framework performed very well but there are some drawbacks. A single dataset was used for the present study with eight physicochemical parameters and five classes of CCME-WQI. Therefore, the model needs to be significantly tested with more water quality data from other geographical areas and environmental conditions. Furthermore, the current implementation was tested on a workstation using GPU acceleration and was not implemented on a resource-limited edge device most frequently used in environmental monitoring systems for AIoT. Future work will therefore concentrate on the validation of the framework on the basis of multi-source datasets, study of domain adaptation for the cross-region generalization, optimization of the model using compression and quantization methods for deployment at the edge and the integration of real-time sensor networks to enable water quality monitoring for intelligent applications in the real environment.

5 Conclusion

This study proposed a novel Transformer-FSL-GAN framework to distinguish water quality into multiple classes in limited data conditions. The proposed approach combines three complementary components: A Transformer-based feature encoder to learn the complex nonlinear relationships between physicochemical parameters, a Prototypical Network (PN) for metric-based few-shot learning, and a Wasserstein Generative Adversarial Network with Gradient Penalty (WGAN-GP) for synthetic data augmentation. This model is designed to overcome some of the typical limitations of traditional machine learning algorithms in environmental monitoring, such as data scarcity and class imbalance. The proposed model addresses these challenges by leveraging these techniques. The Canadian Surface Water Quality dataset (3,949 samples), which is classified into five CCME-WQI health categories, was used for experimental evaluation. The proposed approach was tested on difficult learning settings: 1-shot, 5-shot and 10-shot, and compared with classical machine learning models, deep learning models and few-shot learning baselines. The results consistently showed that Transformer-FSL-GAN achieved the best classification result in terms of all the evaluation metrics. With the 10-shot setting, the suggested model achieved a precision of 97.50%, a recall of 97.33%, and an F1-score of 97.33%, as well as a macro ROC-AUC of 99.53%, which is higher than all the other methods. These improvements were statistically validated by paired t-tests and showed to be highly significant (p < 0.001), and large Cohen’s d values showed large practical effect sizes. The ablation study revealed that all three components in the proposed architecture, namely Transformer feature learning, WGAN-GP augmentation, and prototype-based metric learning, the full architecture always recorded the highest overall performance. Confusion matrix analysis showed only slight misclassifications between adjacent water quality categories of the CCME-WQI, and ROC analysis showed good class separability for all five water quality classes. The UMAP latent space visualization showed a progressive improvement in clustering with the growing size of the support, and the explainability analyses applied to SHAP and LIME revealed that orthophosphate, ammonia, nitrate, dissolved oxygen, and biochemical oxygen demand are the most influential variables for classification. These results show that the proposed model has a high predictive power and also makes the model more interpretable and transparent. Overall, the proposed Transformer-FSL-GAN framework provides an accurate, robust, and interpretable solution for water quality classification in limited-data environments. The framework is well suited for applications in environmental monitoring where the collection of data is costly, time-consuming or limited in geographic scope, as it can learn from a small number of labeled samples. In addition to water quality assessment, the proposed methodology can be further expanded to other environmental monitoring, ecological assessment, and smart sensing applications with low data density that require reliable classification in the field of AIoT. Future work will include validating the framework with more extensive and varied water quality data from various geographical regions to assess its ability to generalise across regions, and evaluating the domain adaptation and transfer learning approaches for improved cross-region generalisation, as well as optimising the model to be deployed on resource-constrained edge AIoT devices by model compression and quantization. In addition, the combination of real-time sensor networks and the proposed framework would enable the development of real-time sensor networks that would allow for continuous and intelligent monitoring of water quality parameters, improving the ability to manage environmental resources effectively and sustainably.

Generative AI statement

The author(s) declared that Generative AI was not used in the creation of this manuscript.

Any alternative text (alt text) provided alongside figures in this article has been generated by Frontiers with the support of artificial intelligence and reasonable efforts have been made to ensure accuracy, including review by the authors wherever possible. If you identify any issues, please contact us.

Publisher’s note

All claims expressed in this article are solely those of the authors and do not necessarily represent those of their affiliated organizations, or those of the publisher, the editors and the reviewers. Any product that may be evaluated in this article, or claim that may be made by its manufacturer, is not guaranteed or endorsed by the publisher.

Footnotes

Edited by: Alyssa Dausman, The Water Institute of the Gulf, United States

Reviewed by: Huichao Yin, New Mexico State University, United States

Jinqiang Wang, Beijing Normal University, China

Ahmed M. Osman, Suez University, Egypt