Class Distribution-Aware Adversarial Training for Semantic Segmentation of Imbalanced Remote Sensing Imagery

Abstract

Class imbalance is a fundamental challenge in semantic segmentation of remote sensing imagery, causing deep convolutional networks to systematically neglect minority categories such as vehicles and small water bodies. To address this, we propose a class distribution-aware adversarial training framework that explicitly penalizes the model’s bias toward dominant classes at the representation level. The framework comprises two key components. First, a Hierarchical Attention U-Net (HAU-Net) integrates a Parallel Hybrid Attention Module (PHAM) at multiple scales, combining spatial and channel attention to amplify feature responses for small and sparsely distributed objects. Second, a distribution discriminator is trained to distinguish the predicted per-class probability distribution from a balanced uniform prior; through a minimax game, the segmentation network is forced to produce class-balanced predictions without manual loss re-weighting. Extensive experiments on the ISPRS Potsdam and Vaihingen datasets demonstrate that our method achieves an overall mIoU of 73.82% and delivers substantial improvements on severely under-represented minority classes, including a +12.11% IoU gain on vehicles and a +19.26% IoU gain on clutter/background, compared to our non-adversarial baseline. Ablation studies confirm that both the hierarchical attention design and the adversarial training contribute independently to these improvements. Our framework offers an effective solution to class imbalance in remote sensing segmentation, with the potential to be integrated into diverse segmentation architectures.

1. Introduction

Semantic segmentation of remote sensing imagery using deep convolutional neural networks (CNNs) has become indispensable for applications such as urban planning, agricultural monitoring, and environmental assessment. However, severe class imbalance remains a critical bottleneck that substantially degrades model performance in real-world scenarios. In typical land cover classification tasks, foreground objects of interest—such as vehicles, roads, or water bodies—often occupy only a tiny fraction of the total pixels, whereas background or dominant classes cover the vast majority. Trained with standard losses, CNNs are biased toward the well-represented classes, leading to severely degraded recognition of minority classes and, consequently, unreliable segmentation maps that fail to capture these sparsely distributed categories [1].

This detrimental effect is starkly illustrated by the ISPRS Potsdam dataset. As shown in Figure 1, the Clutter/Background class (red) overwhelmingly dominates the pixel distribution, while the Vehicles class (yellow) constitutes an extremely small proportion—typically below 2% of total pixels. In such a setting, standard CNN training drives the model to ignore the minority class almost entirely, producing high overall accuracy but a near-zero intersection over union (IoU) for Vehicles. This disconnect between overall metrics and minority-class performance poses a serious risk for downstream decision-making that relies on precise identification of small or rare objects.

Figure 1. Class imbalance in remote sensing imagery: (a,b) display an image sample and its corresponding mask from the ISPRS Potsdam dataset, where severe imbalance exists between the Clutter/Background class (red) and the Vehicles class (yellow); (c) plots the statistical distribution of category ratios for the Potsdam dataset.

To mitigate class imbalance, existing methods generally fall into three categories: data-level resampling (e.g., random over- or under-sampling), algorithm-level loss modification (e.g., weighted cross-entropy, focal loss), and hybrid approaches. While effective to some degree, each approach exhibits notable limitations when applied to high-resolution remote sensing data. Over-sampling minority instances risks overfitting to repeated patterns, whereas under-sampling discards vast amounts of valuable contextual information from majority classes. Weighted loss functions, though widely adopted, often fail under extreme imbalance ratios as they become numerically unstable or require extensive hand-tuning of class weights. More fundamentally, these methods operate primarily on the loss landscape or the input sample distribution, but do not explicitly enforce the learning of balanced, discriminative feature representations for minority classes—a gap that is particularly damaging in remote sensing images where within-class spectral variability and complex spatial textures already challenge feature extraction [2].

Adversarial training provides a distinct and promising avenue to address this representation-level deficiency. In the standard adversarial framework, a generator and a discriminator are jointly trained in a minimax game, forcing the generator to produce outputs indistinguishable from a target distribution. For semantic segmentation, the segmentation network can serve as the generator, and a discriminator can be trained to assess whether the predicted label maps adhere to a class-balanced prior. Through this competition, the segmentation model is compelled to learn features that are not only discriminative for majority classes but also sufficiently expressive for minority classes, thereby counteracting the bias inherent in the data. Several works have exploited adversarial learning to improve segmentation quality—for instance, by aligning predicted label distributions with ground truth statistics or by using adversarial domain adaptation. Moreover, recent studies have tailored adversarial training specifically to address class imbalance in medical image analysis, where lesion pixels are similarly scarce [3]. Nevertheless, the application of adversarial training to class-imbalanced remote sensing segmentation remains largely unexplored. The high spatial resolution and multi-spectral richness of remote sensing data offer both opportunities and challenges: the added spectral dimensions can provide complementary cues for minority classes [4], yet they also increase the feature space complexity, making naive adversarial objectives prone to instability. Hence, a carefully designed adversarial framework, customized to the peculiarities of remote sensing imagery, is urgently needed.

In this work, we introduce an adversarial training framework specifically devised to tackle class imbalance in remote sensing image semantic segmentation. The framework leverages an improved Attention U-Net as the generator (segmentation model), which integrates a novel attention mechanism to selectively amplify feature responses for small and sparse objects that would otherwise be suppressed by dominant classes. Crucially, a companion discriminator is designed not merely to distinguish real vs. predicted masks, but to minimize category distribution bias by evaluating the per-class probability histograms of output predictions, thereby explicitly penalizing the generator whenever minority classes are under-represented. This design forces the segmentation model to maintain a balanced class distribution during inference, significantly boosting the segmentation accuracy of minority categories without sacrificing overall performance.

The main contributions of this paper are threefold:

(i)

We propose a dedicated adversarial training framework for class imbalance in remote sensing segmentation, which jointly optimizes an improved Attention U-Net generator and a class-distribution-aware discriminator, directly addressing the representation bias toward majority classes.

(ii)

We design an effective attention mechanism integrated into the U-Net architecture to enhance feature extraction for sparse and geometrically fine objects typical of remote sensing minority classes.

(iii)

We provide comprehensive experimental analysis on the ISPRS Potsdam and Vaihingen benchmarks, demonstrating significant and consistent improvements for minority class segmentation over multiple strong baselines, and offering practical insights for adversarial training under extreme imbalance.

2. Related Work

2.1. Mitigating Class Imbalance in Semantic Segmentation

Existing strategies for handling class imbalance in segmentation can be broadly organized into data-level resampling, loss function redesign, and hybrid approaches. While extensively studied on natural images, their direct application to remote sensing imagery is severely challenged by extreme imbalance ratios and intricate spatial–spectral structures.

Data-level methods in segmentation typically operate at the patch level rather than on abstract feature vectors. Techniques such as class-uniform sampling and over-sampling of minority-class patches are widely used [5]. More sophisticated augmentation, such as copy-pasting foreground objects onto varying backgrounds, can improve minority-class recall but often introduces unrealistic spatial context in remote sensing scenes, leading to boundary artifacts and spurious correlations [6]. Furthermore, because certain rare classes (e.g., vehicles, small water bodies) appear in very limited variations, naive over-sampling easily causes overfitting. These limitations indicate that data resampling alone cannot fundamentally correct the model’s bias toward dominant classes.

Loss function redesign is the most prevalent alternative. Weighted cross-entropy assigns class-specific penalties, and Focal Loss [7] automatically down-weights well-classified pixels, effectively focusing training on hard, often minority, examples. However, under extreme imbalance ratios, Focal Loss can suffer from numerical instability, and the optimal focusing parameter is highly dataset-dependent. Class-Balanced Loss [8] re-weights based on effective sample counts, yet it assumes that the importance of a class scales inversely with its frequency—a premise that does not universally hold in remote sensing tasks where certain rare classes are critical. Dice loss and its variants show robustness to foreground–background imbalance but struggle with multiple severely under-represented categories simultaneously [9]. Crucially, all these loss-centric methods operate on the scalar output space, without explicitly encouraging the model to learn balanced, discriminative feature representations for minority classes. Consequently, the feature extractor can remain heavily biased toward majority-class patterns, a flaw that is exacerbated in remote sensing images with high intra-class variability.

A few very recent works explore decoupled training or classifier re-balancing for long-tail recognition [10,11], but they are primarily designed for image classification and do not trivially extend to dense prediction tasks. In summary, existing imbalance mitigation methods largely overlook representation-level intervention, leaving a clear gap that our adversarial training framework aims to fill.

2.2. Adversarial Training for Segmentation and Class Imbalance

Adversarial training has evolved beyond its original generative goal and is now a powerful tool for enforcing structural and statistical constraints on model outputs. We focus on two specific threads that directly inform our approach: adversarial semantic segmentation and adversarial handling of class imbalance.

In adversarial semantic segmentation, the segmentation network serves as a generator, and a convolutional discriminator is employed to critique the predicted label maps. Pioneering research on adversarial semantic segmentation demonstrated that adversarial loss can substantially enhance the spatial consistency and boundary fidelity of segmentation outputs [12,13]. This is achieved because the discriminator captures higher-order inconsistencies invisible to per-pixel losses. Our framework inherits this principle but repurposes the discriminator to specifically capture category distribution bias, rather than solely spatial plausibility.

More relevantly, adversarial training has recently been explored as a mechanism to combat class imbalance. In medical image analysis, where lesion pixels are extremely scarce, several studies have designed discriminators that operate on the predicted class probability distributions, penalizing the segmentation network when the output distribution deviates from a balanced prior [14,15]. Through this min-max game, the generator is forced to produce more balanced predictions, thereby improving minority-class sensitivity without manual re-weighting. Additionally, feature-level adversarial learning has been used to disentangle class-relevant representations from class-irrelevant ones, preventing majority-class features from dominating the shared representation space [16]. Despite these promising developments, the application of adversarial training specifically to class-imbalanced remote sensing segmentation remains largely uncharted. Remote sensing data introduce unique challenges, such as multi-spectral inputs, cluttered backgrounds, and an extreme range of object scales, which can destabilize naive adversarial objectives. Hence, a specialized adversarial design that explicitly models the class distribution peculiarities of remote sensing imagery is urgently needed—precisely the contribution of our work.

2.3. Semantic Segmentation in Remote Sensing

A diverse array of architectures has been developed for remote sensing image segmentation. We briefly review the main families and discuss their inherent strengths and weaknesses when handling severe class imbalance, thereby motivating our choice of the generator backbone.

CNN-based architectures, particularly the U-Net family [17,18,19] and DeepLab series [20,21,22], have been the de facto standard. U-Net’s skip connections effectively recover fine spatial details, which is crucial for delineating small minority-class objects. Attention mechanisms integrated into U-Net (e.g., Attention U-Net [18]) further enhance the network’s focus on relevant foreground regions, partially alleviating the suppression of rare classes. However, standard attention gates are still learned predominantly from majority-class statistics and may not sufficiently amplify the extremely weak activations of minority classes under severe imbalance.

Transformer-based models, such as Swin Transformer [23] and various visual transformers, excel at capturing long-range dependencies and global context. This capability is, in principle, beneficial for recognizing sparse objects that require contextual cues. However, transformers are data-hungry and can overfit to majority classes when minority-class training samples are scarce, leading to degraded performance on rare categories. Hybrid CNN-Transformer architectures attempt to combine local detail with global context, yet they do not explicitly address the representation bias induced by class imbalance.

Recent foundation models like the Segment Anything Model (SAM) [24] demonstrate remarkable zero-shot generalization but lack built-in mechanisms to handle semantic class imbalance during fine-tuning. Their outputs are class-agnostic masks that must be matched to land-cover categories, a process that can further amplify annotation biases. Therefore, even with powerful backbones, an explicit imbalance-handling module remains necessary.

Given the trade-off between detail preservation and computational efficiency, and considering the stability requirements of adversarial training, we select an improved Attention U-Net as our generator. We specifically augment the attention mechanism to be more responsive to sparse and geometrically fine minority objects, a design that synergizes with the adversarial imbalance penalty introduced in our framework.

In summary, existing segmentation architectures, regardless of their sophistication, do not inherently resolve extreme class imbalance. Our work integrates a carefully refined attention-based generator with a class-distribution-aware adversarial discriminator, thereby addressing the imbalance issue at both the architectural and the representation-learning levels.

3. Proposed Methodology

To explicitly counteract the model’s bias toward majority classes, we introduce an adversarial training framework that jointly optimizes a primary segmentation network and a class-distribution-aware discriminator. The overall architecture is illustrated in Figure 2.

Figure 2. A block diagram of the proposed approach.

Given an input remote sensing image $x$, the segmentation model $S$ produces a per-pixel class probability map $S\left(x\right)\in {[0,1]}^{H\times W\times C}$, where $C$ is the number of classes. In parallel, the discriminator $D$ operates on the global class distribution of $S\left(x\right)$, rather than on individual pixels, and outputs a scalar score indicating whether this distribution is balanced. Through a minimax game, the segmentation model is forced to generate class probability maps whose marginal class distribution approaches a balanced prior, thereby significantly enhancing the recognition of minority classes.

3.1. Hierarchical Attention U-Net (HAU-Net)

HAU-Net serves as the primary segmentation model and the generator in the adversarial framework. It extends the standard U-Net with a novel attention mechanism applied hierarchically across multiple scales, specifically designed to amplify the responses of small and sparsely distributed objects that are typical of minority classes in remote sensing images.

3.1.1. Overall Architecture

As illustrated in Figure 3, HAU-Net follows a symmetric encoder–decoder structure with four down-sampling and four up-sampling stages. The encoder comprises repeated blocks of two 3 × 3 convolutions followed by batch normalization and LeakyReLU activation (slope 0.2), and a 2 × 2 max pooling with stride 2. The number of feature channels starts at 64 and doubles after each down-sampling, reaching 1024 at the bottleneck. The decoder mirrors the encoder, using transposed convolutions for up-sampling, concatenation with skip-connected encoder features, and two 3 × 3 convolutions. A final 1 × 1 convolution projects the feature maps to a C-class SoftMax output, yielding the per-pixel class probability map $S\left(x\right)\in {[0,1]}^{H\times W\times C}$.

Figure 3. Hierarchical Attention U-Net architecture.

Critically, a Parallel Hybrid Attention Module (PHAM) is inserted into each skip connection before the concatenation step, and an additional PHAM is placed after the bottleneck to refine the highest-level semantic features. Table 1 summarizes the detailed layer-wise configuration of HAU-Net.

Table 1. Architectural details of HAU-Net. “Skip” indicates the features taken from the corresponding encoder level; PHAM is applied to those skip features before concatenation, and also to the bottleneck features.

Stage	Layer Type	Output Size (H × W)	Channels	PHAM
Input	Input Image	256 × 256	3 (RGB) or 4 (IRRG)	No
Encoder 1	Conv Block + MaxPool	128 × 128	64	No
Encoder 2	Conv Block + MaxPool	64 × 64	128	No
Encoder 3	Conv Block + MaxPool	32 × 32	256	No
Encoder 4	Conv Block + MaxPool	16 × 16	512	No
Bottleneck	Conv Block	16 × 16	1024	Yes
Decoder 4	UpConv + Skip (from Enc3)	32 × 32	512	Yes (on skip)
Decoder 3	UpConv + Skip (from Enc2)	64 × 64	256	Yes (on skip)
Decoder 2	UpConv + Skip (from Enc1)	128 × 128	128	Yes (on skip)
Decoder 1	UpConv + Skip (from input)	256 × 256	64	No
Output	1 × 1 Conv + SoftMax	256 × 256	C	No

The hierarchical placement of PHAM is designed to recalibrate feature representations at multiple semantic levels: shallow skip connections carry precise localization information that is critical for small object boundaries, while the bottleneck carries global context. Applying attention at these points helps the network suppress irrelevant background and reinforce the activations of minority-class objects throughout the decoding pathway.

3.1.2. Parallel Hybrid Attention Module (PHAM)

The core innovation of HAU-Net lies in the PHAM, which processes an input feature map $\mathit{F}\in {\mathbb{R}}^{H\times W\times C}$ through two parallel attention branches—spatial and channel—and fuses the refined maps via element-wise multiplication. Figure 4 illustrates the module structure.

Figure 4. Structure of PHAM.

Spatial Attention Branch: To capture long-range spatial dependencies and highlight informative spatial locations, we design a query-key–value (QKV) self-attention branch. Given the input feature map $\mathit{F}\in {\mathbb{R}}^{H\times W\times C}$, We first generate a query–key feature by sequentially applying two $3\times 3$ convolutional layers and a $1\times 1$ convolutional layer, followed by a SoftMax function, to obtain the spatial similarity matrix $\mathit{S}\in {\mathbb{R}}^{HW\times HW}$. The value feature $\mathit{V}\in {\mathbb{R}}^{HW\times C}$ is obtained by a single $1\times 1$ convolution. The spatial attention-enhanced feature is then computed as:

$$\begin{gathered} {\mathit{A}}_S=\hspace{0.17em}{Conv}_{1\times 1}\hspace{0.17em}\left(\mathit{S}\mathit{V}\right), \\ \mathit{S}=SoftMax\left({Conv}_{1\times 1}\left({Conv}_{3\times 3}\left({Conv}_{3\times 3}\left({\mathit{F}}_{\mathit{I}}\right)\right)\right)\right), \mathit{V}={Conv}_{1\times 1}(\mathit{F}) \end{gathered}$$

(1)

Channel Attention Branch: To determine which feature channels are most relevant to the current segmentation task, we compute a channel descriptor through global average pooling (GAP) and global max pooling (GMP). Let $z\in {\mathbb{R}}^C$ be the channel-wise average of ${\mathit{F}}_C$. A bottleneck multi-layer perceptron (MLP) with a reduction ratio $r$ generates the channel attention vector:

$${\mathit{A}}_C=\sigma \hspace{0.17em}(MLP\left({\mathit{z}}_{avg}\right)+MLP({\mathit{z}}_{max})$$

(2)

where ${MLP\left(\mathit{z}\right)=\mathit{W}}_2\hspace{0.17em}ReLU\left({\mathit{W}}_1\mathit{z}+{\mathit{b}}_1\right)+{\mathit{b}}_2$, with ${\mathit{W}}_1\in {\mathbb{R}}^{C/r\times C}$, ${\mathit{W}}_2$∈${\mathbb{R}}^{C\times C/r}$, and $r=16$ in our implementation. The output ${\mathit{A}}_C\in {[0,1]}^{1\times 1\times C}$ selectively emphasizes channels that are most discriminative for minority classes, which often correspond to subtle spectral signatures.

Fusion: The refined feature map is obtained by broadcasting both attention maps to the same spatial and channel dimensions and multiplying element-wise with the original features:

$${\mathit{F}}^{\prime }=\mathbf{F}\odot {\mathit{A}}_S\odot {\mathit{A}}_C$$

(3)

This dual-path attention allows the network to jointly optimize “where to look” and “which features to trust,” which is effective for the small and sparsely distributed objects from class imbalance.

3.2. Class Distribution Discriminator

The discriminator $D$ is purposefully designed to detect class distribution bias rather than pixel-level authenticity. This makes it directly suited for class imbalance mitigation. Its architecture is a lightweight MLP that maps a $C$-dimensional distribution vector to a single scalar.

For any segmentation output $S\left(x\right)$ of spatial size $H\times W,$ we compute the per-class average probability via global average pooling:

$${\mathit{v}}_{S\left(x\right)}={\displaystyle \frac{1}{HW}}{\displaystyle \sum _{p=1}^{HW}}{S_1\left(x\right)}_p,\dots ,{\displaystyle \frac{1}{HW}}{\displaystyle \sum _{p=1}^{HW}}{S_C\left(x\right)}_p$$

(4)

The vector ${\mathit{v}}_{S\left(x\right)}\in {[0,1]}^C$ sums to 1 and represents the empirical label distribution predicted by the segmentation model for image $x$. The discriminator $D$ consists of fully connected layers:

Input layer: $C$ neurons.

Hidden layer 1: 128 neurons, followed by LeakyReLU (negative slope 0.2).

Hidden layer 2: 64 neurons, followed by LeakyReLU (negative slope 0.2).

Output layer: 1 neuron with sigmoid activation, producing a scalar $D\left(\mathit{v}\right)\in [0,1]$.

We define the balanced prior as the uniform distribution $\mathit{u}=[{\displaystyle \frac{1}{C}},{\displaystyle \frac{1}{C}},\dots ,{\displaystyle \frac{1}{C}}]$. During the discriminator’s update step, we assign label 1 to $u$ and label 0 to ${\mathit{v}}_{S\left(x\right)}$ from the current segmentation model. This forces the discriminator to learn the manifold of balanced distributions and distinguish it from biased ones. The generator (HAU-Net), in turn, tries to produce segmentation maps whose class distribution ${\mathit{v}}_{S\left(x\right)}$ is classified as balanced by D, i.e., $D\left({\mathit{v}}_{S\left(x\right)}\right)\to 1$. The formulation offers a global and interpretable way to combat class imbalance, with the added cost of less than 0.5% to the total parameters.

3.3. Adversarial Training Procedure

The training of HAU-Net and the distribution discriminator follow a minimax formulation with alternating optimization. We present the loss functions, then the full training algorithm.

3.3.1. Weighted Cross-Entropy Loss

To ensure pixel-level accuracy while mitigating the dominance of majority classes, we use a weighted cross-entropy (WCE) loss. The weight for each class $i$ is set to $w_i={\left(n_i^{0.5}\right)}^{-1}$, where $n_i$ is the total number of pixels of class $i$ in the training set. This square-root compression prevents the weights for extremely rare classes from becoming excessively large, which could cause gradient instability. For a given pixel with one-hot ground truth $y_i$ and prediction $S_i\left(x\right)$, the WCE loss is:

$${\mathcal{L}}_{wce}=-{\displaystyle \frac{1}{N}}{\sum }_{p=1}^N{\sum }_{i=1}^C{w}_i·y_i\left(p\right)·\log ({S_i\left(x\right)}_p)$$

(5)

where $N=H\times W$ is the total number of pixels per image.

3.3.2. Adversarial Loss

For the discriminator, we employ the standard binary cross-entropy loss to distinguish balanced prior samples from biased predictions:

$${\mathcal{L}}_D=-{\mathbb{E}}_x\left[\log D\left(\mathit{u}\right)+\log (1-D({\mathit{v}}_{S\left(x\right)})\right]$$

(6)

The generator (segmentation model) aims to fool the discriminator, i.e., to maximize the probability that its predicted distribution is classified as balanced:

$${\mathcal{L}}_{adv}=-{\mathbb{E}}_x\left[\log D\left({\mathit{v}}_{S\left(x\right)}\right)\right]$$

(7)

Note that we optimize $D$ using the uniform prior $\mathit{u}$ as the only positive example, rather than using the ground-truth label distribution, because the ground-truth distribution in remote sensing images is itself heavily imbalanced. Using the ground truth as the target would encourage the model to replicate the imbalance rather than correct it.

3.3.3. Full Training Process

The overall objective for HAU-Net is a weighted sum of the segmentation loss and the adversarial loss:

$$\underset{S}{\min }{\mathcal{L}}_{total}={\mathcal{L}}_{wce}\left(S\right)+\lambda {\mathcal{L}}_{adv}(S,D)$$

(8)

where $\lambda$ balances the two terms. We empirically set $\lambda =0.1$ to ensure that pixel accuracy remains the primary objective while the adversarial term provides a meaningful regularization effect. The discriminator is optimized by: $\underset{D}{\max }\mathcal{L}\left(D\right)$.

The training is performed via alternating gradient descent. In each iteration, we first update the discriminator $D$ for $k_D$ steps, and then update the generator $S$ for one step. We use the Adam optimizer with ${\beta }_1=0.9$, ${\beta }_2=0.999$ and a learning rate of $1\times {10}^{-4}$ for both models. Batch size is set to 8, and training proceeds for 200 epochs with a learning rate decay factor of 0.1 at epochs 100 and 150. The complete training procedure is formalized in Algorithm 1.

Algorithm 1. Adversarial Training Procedure

1 $\mathrm{Initialize} \mathrm{parameters} \mathrm{of} S$ (HAU-Net) and $D$ (MLP) with random weights.

2 $\mathrm{For} \mathrm{each} \mathrm{training} \mathrm{iteration} \mathrm{do}:$

3 $\mathrm{Sample} \mathrm{a} \mathrm{mini}\text{-}\mathrm{batch} \mathrm{of} \mathrm{remote} \mathrm{sensing} \mathrm{images} \{x_1,\dots ,x_m\}$ with ground truth $\{y_1,\dots ,y_m\}$

4 $\mathrm{Forward} \mathrm{pass}: \mathrm{compute} S\left(x_i\right)$ and the distribution vectors ${\mathit{v}}_{S\left(x\right)}$.

5 $\mathrm{Generate} \mathrm{balanced} \mathrm{prior} \mathrm{samples} {\mathit{u}}_1,\dots ,{\mathit{u}}_m$.

6 $\mathrm{Update} \mathrm{discriminator} D$ by maximizing ${\displaystyle \frac{1}{m}}{\sum }_{j=1}^m\left[logD\right({\mathit{u}}_j)+log(1-D\left(v_{S\left(x_j\right)}\right)]$.

7 $\mathrm{Compute} {\mathcal{L}}_{wce}$ and ${\mathcal{L}}_{adv}$ for the current batch.

8 $\mathrm{Update} \mathrm{generator} S$ by minimizing ${\displaystyle \frac{1}{m}}{\sum }_{j=1}^m{[\mathcal{L}}_{wce}\left(S\left(x_j\right),y_j\right)+\lambda (-logD(v_{S\left(x_j\right)}\left)\right]$

9 $\mathrm{End} \mathrm{for}$

4. Experimental Results and Analysis

4.1. Dataset

We evaluate our framework on two widely used sub-datasets of the ISPRS 2D Semantic Labeling benchmark (International Society for Photogrammetry and Remote Sensing (ISPRS), Hannover, Germany) [25]: Potsdam and Vaihingen. Both provide high-resolution aerial imagery with pixel-level land cover annotations across six common categories: impervious surfaces, buildings, low vegetation, trees, vehicles, and clutter/background. The clutter/background class aggregates diverse objects such as water bodies, containers, tennis courts, and swimming pools. The distribution of these six classes is highly imbalanced, as illustrated in Figure 1, with vehicles and clutter/background occupying extremely small pixel proportions.

The Potsdam dataset consists of 38 ortho-rectified RGB images, each of size 6000 × 6000 pixels with a ground sampling distance of 5 cm. The Vaihingen dataset includes 33 images of mixed rural–urban scenes, provided as near-infrared, red, and green bands (IRRG). Because the official ISPRS test labels are not publicly available, we randomly split the official training set into training and validation subsets at a fixed 7:3 ratio. To match the input size of our model, we performed batch cropping on both original images and ground-truth masks to generate non-overlapping 256 × 256 patches. This process yields roughly 20,102 patches from the Potsdam dataset and 1617 patches from the Vaihingen dataset, amounting to a total of 21,719 patches.

4.2. Experimental Setup

All models were implemented using PyTorch 1.11 and Python 3.10, and trained on a single NVIDIA RTX 4090 GPU (24 GB). The system ran on Ubuntu 20.04 LTS with an Intel Xeon W-2255 CPU. Training was performed for 200 epochs with a batch size of 8. We employed the Adam optimizer with ${\beta }_1=0.9$, ${\beta }_2=0.999$, and an initial learning rate of $1\times {10}^{-4}$, which is decayed by a factor of 0.1 at epochs 100 and 150. A weight decay of $1\times {10}^{-4}$ was applied as L2 regularization to mitigate overfitting. The learning rate schedule was further smoothed using a cosine annealing warm-up during the first 10 epochs.

To improve model generalization, we applied online data augmentation including random horizontal and vertical flipping, random rotation (±30°), random scaling (0.8–1.2×), and photometric jitter (brightness adjustment within $\pm 15\%$, contrast adjustment factor 0.8–1.5). All augmentation operations were applied consistently to both the images and the corresponding labels.

The training progress was monitored via the total loss ${\mathcal{L}}_{total}={\mathcal{L}}_{ce}+\lambda {\mathcal{L}}_{adv}$ on the training and validation sets. As shown in Figure 5, the loss decreased rapidly during the initial 40 epochs and reached a stable dynamic equilibrium after approximately 80 epochs, with no significant divergence between training and validation curves, indicating that the model did not overfit. The convergence of the loss correlated well with the improvement in test-set mIoU, confirming the stability and effectiveness of the adversarial training process. The key experimental configurations are summarized in Table 2.

Figure 5. Loss function curve.

Table 2. Summary of experimental configurations.

Category	Details
CPU	Intel Xeon W-2255 (Intel Corporation, Santa Clara, CA, USA)
GPU	NVIDIA RTX 4090 (24 GB) (NVIDIA Corporation, Santa Clara, CA, USA)
Operating System	Ubuntu 20.04 LTS (Canonical Ltd., London, UK)
Framework	PyTorch 1.11 (Meta AI, Menlo Park, CA, USA), Python 3.10 (Python Software Foundation, Beaverton, OR, USA)
Optimizer	Adam (${\beta }_1=0.9$, ${\beta }_2=0.999$)
Initial Learning Rate	$1\times {10}^{-4}$
Learning Rate Scheduling	Step decay (×0.1 at epoch 100, 150) + Cosine annealing warm-up
Weight Decay	$1\times {10}^{-4}$
Batch Size	8
Training Epochs	200
Data Augmentation	Random geometric distortion and photometric distortion perturbation
Overfitting Prevention	L2 adaptive regularization constraint

Evaluation Metrics: We evaluate segmentation quality using four complementary metrics. The standard mean Intersection over Union (mIoU) is defined as the arithmetic mean of the per-class IoU values:

$$\mathrm{mIoU}={\displaystyle \frac{1}{C}}{\displaystyle \sum _{i=1}^C}{\mathrm{IoU}}_i$$

(9)

where $C$ is the number of classes. To additionally reflect overall pixel-level accuracy, we report the Frequency Weighted IoU (FWIoU):

$$\mathrm{FWIoU}={\displaystyle \frac{1}{{\displaystyle {\sum }_{i=1}^C}{n}_i}}{\displaystyle \sum _{i=1}^C}{n}_i\cdot {\mathrm{IoU}}_i$$

(10)

where $n_i$ denotes the total number of pixels of class $i$ in the evaluation set. FWIoU weights each class by its pixel frequency and therefore emphasizes majority classes, whereas mIoU treats all classes equally and is more sensitive to minority-class performance.

We also report the mean Pixel Accuracy (mPA), defined as the arithmetic mean of per-class pixel accuracies:

$$\mathrm{mPA}={\displaystyle \frac{1}{C}}{\displaystyle \sum _{i=1}^C}{\displaystyle \frac{{\mathrm{TP}}_i}{n_i}}$$

(11)

where ${\mathrm{TP}}_i$ is the number of correctly classified pixels of class $i$. In addition, the Kappa coefficient is used to measure the agreement between the predicted segmentation and the ground truth, corrected for chance agreement:

$$\kappa ={\displaystyle \frac{p_o-p_e}{1-p_e}}$$

(12)

where $p_o$ is the overall pixel accuracy and $p_e$ is the expected agreement by chance, computed from the confusion matrix.

All metrics are computed on the stitched full-resolution test predictions, and classes absent from a given image or patch are excluded from the IoU and pixel accuracy calculations for that image to avoid undefined values.

The convergence curves of the loss function on the training and validation sets are shown in Figure 5.

4.3. Comparative Experiments

We conduct a comprehensive comparison of our method against a representative set of baseline models. The baselines are organized into four categories: (i) CNN-based segmentation architectures (U-Net [17], PSPNet [26], DeepLabV3+ [21], ACNet [27], HRNet + OCR [28]); (ii) Transformer-based and hybrid architectures (SegFormer-B4 [29], Segmenter [30], SwinUperNet [31], UNetFormer [32]); (iii) lightweight networks (BiSeNetV2-L [33], DABNet [34]). The backbone network configurations of these models are tabulated in Table 3. All models were trained with identical data augmentation and for the same number of epochs to ensure a fair comparison.

Table 3. Models’ backbone configurations in comparison experiments.

Models	Backbone
U-Net	--
PSPNet	ResNet101
DeepLabV3+	ResNet101
ACNet	DenseNet121
HRNet + OCR	HRNetV2
BiSeNetV2-L	--
DABNet	MobileNetV2
SegFormer-B4	MIT
Segmenter	ViT-Tiny
SwinUperNet	Swin-Tiny
UNetFormer	Swin-Tiny
HAU-Net (Ours)	Hierarchical Attention U-Net

4.3.1. Quantitative Results

Table 4 reports the global performance metrics. Bold values indicate the best performance for each metric, and the same applies to subsequent tables.

Table 4. Comparative experimental results on the ISPRS dataset.

Models	mIoU (%)	FWIoU (%)	mPA (%)	Kappa (%)	IoU (%)
					Imp. Surf.	Buildings	Low Veg.	Trees	Clutter /Bkgd.	Vehicles
U-Net	51.62	66.05	67.83	47.35	76.04	65.74	58.94	64.34	18.54	26.14
PSPNet	58.41	71.74	77.52	52.36	81.31	71.90	69.26	75.37	20.10	32.54
DeepLabV3+	63.82	75.23	84.52	66.10	79.13	81.33	71.03	77.93	22.97	50.53
BiSeNetV2-L	58.60	70.59	75.22	58.28	77.03	78.34	62.53	73.13	21.11	39.43
DABNet	60.99	71.39	79.86	55.93	78.88	76.08	64.28	71.98	24.06	50.68
ACNet	64.90	75.67	80.77	61.76	80.73	81.80	70.60	76.59	28.27	51.41
HRNet + OCR	70.54	79.90	87.62	64.92	87.23	84.23	72.93	78.23	33.07	67.53
SegFormer-B4	68.01	76.66	83.27	58.11	79.17	79.37	76.95	76.91	34.54	61.14
Segmenter	67.39	77.30	82.40	62.16	81.03	83.83	73.03	77.53	31.79	57.13
SwinUperNet	71.22	78.60	86.67	61.51	84.51	87.11	69.44	74.68	45.07	66.51
UNetFormer	71.93	78.93	88.57	66.64	86.24	84.74	70.04	75.24	44.48	70.84
HAU-Net (Ours)	73.82	79.57	89.04	65.88	83.78	85.51	74.18	78.27	48.16	73.04

From the perspective of class-fair evaluation, our method achieves the highest standard mIoU of 73.82%, outperforming all baselines including the strong HRNet + OCR (70.54%) and UNetFormer (71.93%). In terms of Frequency Weighted IoU (FWIoU), which weights classes by pixel frequency and thus favors majority classes, HRNet + OCR attains the best value (79.90%), with our method ranking a very close second (79.57%). This small gap in FWIoU is expected because our adversarial objective intentionally shifts model capacity from over-represented to under-represented classes. Moreover, our method achieves the highest mean Pixel Accuracy (89.04%), indicating superior per-class average recall and confirming that the model does not sacrifice minority-class sensitivity for overall pixel-level accuracy.

The advantage of our method becomes particularly evident when examining the severely imbalanced minority classes. For vehicles, our method achieves 73.04% IoU, surpassing HRNet + OCR by +5.51% (67.53%); for clutter/background, our method reaches 48.16% IoU, outperforming HRNet + OCR by +15.09% (33.07%). These substantial gains directly validate that the proposed adversarial training framework effectively suppresses majority-class bias and greatly enhances the recognition of rare classes. In contrast, on certain majority classes such as impervious surfaces, our method shows a slight decrease compared to the best individual models (e.g., 83.78% vs. 87.23% for HRNet + OCR). This is an expected trade-off: the adversarial loss explicitly penalizes over-confidence on dominant classes, reallocating predictive capacity to minority categories without causing catastrophic degradation on any single class. The competitive mPA (89.04%) and Kappa (65.88%) further support that our method provides a better overall class-wise balance than competing approaches.

To examine whether adversarial training over-equalizes the class distribution or inflates minority-class IoU through excessive false positives, we report the statistics of the predicted class distributions on the test set. The full model yields a mean entropy of 1.42 and a mean KL divergence of 0.63 to the uniform prior, compared with 1.21 and 0.81 for the WCE-only baseline. These values indicate a more balanced yet non-degenerate distribution. For the vehicle class, precision and recall are 82.0% and 87.0% for the full model, versus 80.0% and 71.9% for the WCE-only baseline, confirming that the IoU gains arise primarily from improved recall rather than excessive false positives.

4.3.2. Qualitative Visualization

Figure 6 and Figure 7 present qualitative segmentation results of our method on the Vaihingen and Potsdam datasets, respectively. In each figure, we show the original image, the ground-truth label mask, and our predicted segmentation map.

Figure 6. Visualization of semantic segmentation on the ISPRS Vaihingen dataset.

Figure 7. Visualization of semantic segmentation on the ISPRS Potsdam dataset.

On the Vaihingen dataset (Figure 6), our model accurately delineates the boundaries of buildings and low vegetation. More crucially, it successfully identifies small vehicles scattered in the images, which are often missed by standard models. The tree category also benefits from fine boundary preservation, with minimal confusion with low vegetation.

On the Potsdam dataset (Figure 7), our model produces a clean segmentation of impervious surfaces and buildings. While occasional misclassifications appear in the clutter/background class—particularly around complex structures such as containers and water edges—the overall segmentation remains robust. These residual errors highlight the inherent difficulty of recognizing extremely rare and texture-variant categories, and they point to directions for future improvement, such as incorporating edge-aware constraints.

4.3.3. Comparison of Visual Results

Figure 8 compares the segmentation maps produced by our method and several strong baselines on representative challenging patches. To facilitate detailed comparison, we highlight critical regions with black dashed boxes.

Figure 8. Visualization of the comparison experiments.

In the dashed-box region, HRNet + OCR exhibits noticeable omissions of buildings (blue); Our method, in contrast, recovers these buildings with significantly better completeness and boundary fidelity. And also, most baseline algorithms suffer from substantial misclassification of the clutter/background class (red). Typical failure modes include false positives, where background clutter is erroneously labeled as other land-cover categories, and false negatives, where true clutter areas are incorrectly assigned to surrounding classes. Our method markedly reduces both types of errors, producing a cleaner and more coherent segmentation of the clutter/background regions.

These qualitative gains directly reflect the capacity of the hierarchical attention and adversarial training to strengthen the recognition of minority classes under severe imbalance.

4.4. Ablation Study

We conduct a series of ablation experiments to isolate and quantify the contribution of each key component in our framework. All ablation models in Section 4.4.1 and Section 4.4.2 are trained without the adversarial discriminator, using only the weighted cross-entropy loss, so that the improvements attributable to the architecture design and the adversarial training can be assessed independently.

4.4.1. Backbone Network Ablation

To demonstrate the effectiveness of our proposed Hierarchical Attention U-Net (HAU-Net) as the segmentation backbone, we compare it against several widely used architectures, including the standard U-Net [17], HRNet-w48 [35], Vision Transformer (ViT) [36], and Swin-Tiny [37]. All backbones are trained under identical settings without adversarial training. The results are shown in Table 5.

Table 5. Ablation study on backbone networks.

Backbone	mIoU (%) ↑	FWIoU (%) ↑	mPA (%) ↑	Kappa (%) ↑	#Params (M) ↓	FLOPs (G) ↓
U-Net	51.62	66.05	67.83	47.35	31.6	78.5
HRNet-w48	65.13	78.30	83.40	62.16	65.4	215.3
ViT	69.88	78.83	87.43	63.64	91.2	138.7
Swin-Tiny	70.17	79.24	88.96	67.41	60.8	108.2
HAU-Net (Ours)	67.23	77.44	82.55	62.27	38.2	189.6

Under the WCE-only setting, Swin-Tiny achieves the highest mIoU, mPA, and Kappa among the compared backbones, benefiting from its strong global context modeling. Our HAU-Net obtains competitive results—77.44% FWIoU and 82.55% mPA—while requiring only 38.2 M parameters, which is significantly fewer than Swin-Tiny’s 60.8 M and ViT’s 91.2 M. Although HAU-Net’s FLOPs (189.6 G) are higher than Swin-Tiny’s (108.2 G), this is a consequence of the parallel spatial–channel attention design, which is intended to amplify minority-class features rather than maximize computational efficiency.

More importantly, when HAU-Net is integrated into the full adversarial framework, it achieves 79.57% FWIoU and 89.04% mPA, surpassing Swin-Tiny by +0.33% and +0.08%, respectively. This demonstrates that the performance gain of the proposed framework comes primarily from the adversarial training, rather than from the backbone architecture alone.

4.4.2. PHAM Ablation

We next evaluate the efficacy of the Parallel Hybrid Attention Module (PHAM) by replacing it with several established attention mechanisms while keeping the U-Net backbone and all other settings fixed. The compared variants are:

Model A: Standard U-Net without any attention.

Model B: Channel attention SE [38] inserted into skip connections.

Model C: Spatial attention STN [39] inserted into skip connections.

Model D: Joint attention CBAM [40] inserted into skip connections.

Model E: Dual-stream attention DANet [41] inserted into skip connections.

Model F: PHAM inserted into both low-level and high-level feature maps.

Table 6 reports the global metrics and per-class IoU for all variants.

Table 6. Ablation study on attention mechanisms.

Models	mIoU (%) ↑	FWIoU (%) ↑	mPA (%) ↑	Kappa (%) ↑	#Params (M) ↓	GFLOPs ↓
Model A	51.62	66.05	67.83	47.35	31.6	78.5
Model B	61.26	74.55	79.21	58.96	31.8	79.1
Model C	60.09	71.83	78.16	58.20	38.5	85.3
Model D	64.87	76.70	85.48	62.71	32.9	84.5
Model E	65.93	77.35	86.11	62.83	45.2	112.4
Model F	67.23	77.44	82.55	62.27	38.2	189.6

Among the attention variants under the WCE-only setting, our PHAM (Model F) achieves the highest standard mIoU (67.23%), outperforming DANet by +1.30% and CBAM by +2.36%. Its FWIoU (77.44%) is also the best among the compared attention modules, surpassing DANet by +0.09% and CBAM by +0.74%. Although its mPA (82.55%) is lower than DANet’s 86.11% and CBAM’s 85.48%, and its Kappa (62.27%) is slightly below the best values, we emphasize that the primary objective of PHAM is not to maximize aggregate pixel accuracy, but to amplify the feature responses of sparse and small objects that are most severely affected by class imbalance. To verify this, we refer to the per-class results reported later in Table 7, where Model I corresponds exactly to Model F in Table 6. Compared with the standard U-Net baseline (Model A in Table 7), PHAM improves the IoU of vehicles from 26.14% to 60.93% and the IoU of clutter/background from 18.54% to 28.90%, confirming its effectiveness in enhancing minority-class features.

Table 7. Ablation study on adversarial training.

Models	mIoU (%) ↑	FWIoU (%) ↑	mPA (%) ↑	Kappa (%) ↑	IoU (%) ↑
					Imp. Surf.	Buildings	Low Veg.	Trees	Clutter/Bkgd.	Vehicles
Model I	67.23	77.44	82.55	62.27	85.17	83.52	69.57	75.28	28.90	60.93
Model II	73.82	79.57	89.04	65.88	83.78	85.51	74.18	78.27	48.16	73.04

Regarding computational cost, PHAM incurs a higher FLOP count (189.6 G) than other attention modules due to its parallel spatial and channel branches. Nevertheless, its parameter count (38.2 M) remains lower than that of DANet (45.2 M), indicating acceptable complexity given the substantial gains on challenging minority categories.

4.4.3. Effectiveness of Adversarial Training

To isolate the contribution of adversarial training, we fix HAU-Net with PHAM as the backbone and compare two training configurations:

Model I: HAU-Net trained with weighted cross-entropy loss ${\mathcal{L}}_{wce}$ only (no adversarial training).

Model II: HAU-Net trained with the full objective ${\mathcal{L}}_{wce}+\lambda {\mathcal{L}}_{adv} (\lambda =0.1,$ adversarial training enabled).

The experimental results are shown in Table 7.

With adversarial training enabled, Model II achieves substantial improvements over Model I in both standard mIoU and FWIoU. Specifically, the standard mIoU increases from 67.23% to 73.82% (+6.59 percentage points), while FWIoU improves from 77.44% to 79.57% (+2.13 percentage points). The gains are disproportionately concentrated on the minority classes: vehicles improve by +12.11% (from 60.93% to 73.04%) and clutter/background by +19.26% (from 28.90% to 48.16%). Concurrently, majority classes such as impervious surfaces exhibit only a marginal decrease (−1.39%, from 85.17% to 83.78%), while buildings improve by +1.99% (from 83.52% to 85.51%), confirming that the adversarial loss effectively reallocates model capacity from over-represented to under-represented categories without catastrophic forgetting of dominant classes. These results provide direct evidence that the proposed adversarial training framework succeeds in its core objective—mitigating class imbalance at the representation level.

To further verify the advantage of the learnable critic over a fixed analytical penalty, we trained a non-adversarial variant in which the adversarial loss was replaced by a direct KL-divergence penalty between the predicted class distribution and the uniform prior, using the same $\lambda =0.1$. This variant achieved a vehicle IoU of 71.62% and a clutter/background IoU of 45.83%, compared with 73.04% and 48.16% for the adversarial model, confirming the benefit of the learnable critic.

4.4.4. Sensitivity Analysis of the Balance Coefficient $\lambda$

The hyper-parameter $\lambda$ controls the strength of the adversarial regularization relative to the primary segmentation loss. To assess its influence, we train Model II with varying $\lambda$ values sampled from the set $\left\{0.0, 0.01, 0.05, 0.1, 0.3, 0.5, 0.8, 1.0\right\}$, where $\lambda =0.0$ corresponds to Model I. Figure 9 plots the resulting mIoU as a function of $\lambda$.

Figure 9. Impact of the balance coefficient λ on mIoU. Evaluated λ values are {0.0, 0.01, 0.05, 0.1, 0.3, 0.5, 0.8, 1.0}; the best performance is achieved at λ = 0.1.

As $\lambda$ increases from 0 to 0.1, the mIoU improves from 77.44% to a peak of 79.57%, demonstrating that a moderate adversarial penalty effectively regularizes the model toward balanced predictions. When $\lambda$ exceeds 0.1, the mIoU begins to decline, falling to 74.31% at $\lambda =0.5$ and dropping sharply to 64.52% at $\lambda =1.0$. This inverted-U pattern is expected: an excessively large $\lambda$ forces the model to overly homogenize class distributions at the expense of pixel-level accuracy, leading to degraded overall performance. Based on this analysis, we set $\lambda =0.1$ for all experiments reported in this paper.

5. Conclusions

In this paper, we addressed the critical yet under-explored problem of class imbalance in semantic segmentation of remote sensing imagery. Motivated by the observation that standard CNN training systematically neglects minority classes under extreme pixel distribution skew, we proposed a novel adversarial training framework that couples a Hierarchical Attention U-Net (HAU-Net) generator with a class-distribution-aware discriminator. The discriminator is specifically designed to detect global deviations from a balanced prior, thereby forcing the segmentation model to learn features that are equally expressive for rare and dominant categories.

Extensive experiments on the ISPRS Potsdam and Vaihingen benchmarks demonstrate the effectiveness of our approach. Our method achieves the highest standard mIoU of 73.82% and an FWIoU of 79.57%, along with an mPA of 89.04%, ranking among the top-performing models on this benchmark. On the severely under-represented minority classes, our method improves vehicle IoU by +5.51% (73.04% vs. 67.53%) and clutter/background IoU by +15.09% (48.16% vs. 33.07%) relative to the strongest external baseline (HRNet + OCR).

We note a limitation that is inherent to our design: the adversarial objective explicitly penalizes over-confidence on majority classes, which may result in a marginal accuracy decline on certain dominant categories. We view this not as a failure, but as an expected and desirable trade-off: the model sacrifices a small amount of accuracy on already well-segmented classes to achieve large gains on critically under-segmented ones. The high mPA and Kappa scores confirm that the overall class-wise balance is substantially improved.

Several promising directions emerge from this work. First, the current discriminator operates on a per-image global distribution; extending it to model spatial priors could further refine the balance between local accuracy and global fairness. Second, the fixed uniform prior could be replaced with a learned or task-specific target distribution that reflects realistic planning constraints. Third, the proposed adversarial framework is architecture-agnostic and could be integrated with emerging foundation models to enable class-balanced fine-tuning on large-scale remote sensing benchmarks.