Unified Framework for Histopathology Image Augmentation and Classification via Generative Models

Data scarcity remains a significant obstacle in applying deep learning technologies in the medical field. To mitigate this challenge, data augmentation has emerged as a powerful technique for expanding training datasets. However, conventional augmentation methods used in general image processing often struggle or cannot be applied due to the complexities and unique characteristics of medical images. This limitation underscores the appeal of generative models, which have demonstrated strong performance in augmenting medical data.

Currently, three primary generative models are utilized for medical image synthesis: variational autoencoders (VAEs), generative adversarial networks (GANs), and diffusion models [8]. GANs, in particular, are widely favored for their capability to generate realistic images even with limited data availability [9], despite challenges such as training instability, convergence issues, and mode collapse [10]. In contrast, VAEs offer greater output diversity and avoid mode collapse compared to GANs, but they often produce blurry images, which limits their adoption for augmentation purposes. Recently, diffusion models [11, 12] have shown promise in generating realistic and diverse outputs, though they require extensive data and computational resources for training [8]. Given the inherent data scarcity in medical datasets, our study focuses on augmenting data with limited availability using generative models. Considering these constraints, GANs emerge as the preferred choice due to their proven effectiveness in generating high-quality medical images under restricted data conditions.

A GAN model consists of two key components: a discriminator and a generator. During training, these components are alternatively trained. The discriminator learns to differentiate between the synthetic images generated by the generator and real training data. Conversely, the generator learns to create synthetic images that are as realistic as possible to deceive the discriminator. During testing, only the well-trained generator is used to synthesize realistic images.

GAN was originally proposed using the convolutional neural network (CNN) architecture to formulate a min-max optimization problem aimed at narrowing the gap between real and synthetic data distributions [13]. Recently, transformer networks [14] have attracted attention for their ability to effectively model non-local contextual dependencies. In the realm of medical image synthesis, there is a growing trend towards replacing CNNs with transformer networks, given the significant role of global understanding of training data in generating realistic medical images. For instance, Korkmaz et al. [15] introduced the GVTrans model, which employs cross-attentive visual transformers to map low-dimensional noise and latent variables onto MRI data. Zhang et al. [16] proposed PTNet3D, leveraging pyramid transformer networks to generate high-resolution 3D longitudinal infant brain MRI data. However, these methods often adopt a hybrid structure that integrates a transformer-based generator alongside a CNN-based discriminator.

In contrast to hybrid models, a pure transformer-based GAN offers architectural simplicity that can potentially reduce method complexity, thereby facilitating stable training and reducing computational overhead. MedViTGAN stands out as a pure vision transformer-based conditional GAN designed specifically for generating histopathology images for data augmentation [17]. Inspired by MedViTGAN, our study adopts a pure transformer GAN framework to capture non-local or long-range information in histopathology images. Furthermore, different from previous methods, our framework integrates conditional class projection and selective data augmentation techniques to enhance both class separation and generation quality.

To enhance performance in various image synthesis applications, several GAN-based studies have explored combining the Wasserstein loss with different loss functions. For instance, Liu et al. [18] integrated total variation loss, pixel loss, and feature-level losses linearly with the Wasserstein loss function to train an auto-painter, resulting in improved painting quality. Similarly, Ebenezer et al. [19] employed a combination of Wasserstein loss, L1 pixel loss, and VGG loss to train their GAN. However, these approaches typically involve manual tuning of the weights assigned to each loss component. This manual hyperparameter tuning relies heavily on expert knowledge and can lead to time-consuming and computationally expensive training processes. Moreover, it may limit the generalizability of the method and result in sub-optimal performance.

The success of multi-task learning suggests a more sophisticated approach to combining loss functions across different tasks. Kendall et al. [6] introduced a multi-task loss formulation that leverages the homoscedastic uncertainty of each task to dynamically weigh their respective losses. Inspired by their work, we propose a multi-loss weighting function that automates the adjustment of these weights during the training process.

The GAN architecture utilized in our framework is a pure transformer-based model inspired by TransGAN [5]. The core component is the Vision Transformer (ViT) encoder [20], comprising multi-head self-attention modules and a feed-forward multi-layer perceptron (MLP) with GELU non-linearity. The self-attention modules and MLPs in the ViT encoder are equipped with residual connections and layer normalization for improved stability and performance.

Our GAN’s generator is structured with four-stage blocks designed to progressively learn and upsample the given latent vector. Each stage block consists of a class projection layer, an upsampling module, and four ViT encoders. As the generator advances through these stages, it incrementally increases the feature map resolution until it reaches the target dimensions. Initially, the latent noise vector $z$ is concatenated with the one-hot class label $c$ and passed through a linear projection layer to generate embedding tokens $X_{0}\in\mathbb{R}^{H_{0}\times W_{0}\times C}$. These tokens are then sent to the first stage of the generator. In the first stage, the embedded feature map undergoes upsampling from $X_{0}\in\mathbb{R}^{H_{0}\times W_{0}\times C}$ to $X_{1}\in\mathbb{R}^{2H_{0}\times 2W_{0}\times C}$ using cubic interpolation, ensuring early feature learning without reducing dimensions. In the following three stages, a pixel-shuffle module [21] is used to upsample the resolution by a factor of 2$\times$, while decreasing the channels to one-quarter. This approach reduces memory requirements and enhances network efficiency. To tackle the issue where the self-attention module sacrifices local information to learn global correspondences during the high-resolution generation stage (the third and fourth stages), we replace the self-attention module with the grid self-attention module [5] in the ViT encoder. This modification enables our model to effectively capture both global and local information. Specifically, we set the predefined window size to $32\times 32$ in the third stage and $16\times 16$ in the fourth stage.

The discriminator is similarly structured into four stages. The first three stages consist of a ViT encoder and an average pooling layer. The final stage includes two ViT encoders that handle the class and source heads, respectively. Between the first three stages, we apply a multi-scale technique that combines the outputs from the last block with patches of varying sizes extracted from the same input image. After passing through an MLP layer, the patch information is encoded into a sequence of embeddings for concatenation, allowing the model to learn both semantic structure and texture details efficiently.

For conditional learning, a class token and a source token are appended to the beginning of the 1D sequence at the end of the third block. This modified sequence then passes through the fourth stage, enabling the model to make both class predictions and real/synthesis judgments.

To further improve the conditional image generation, we propose the conditional class projection module, including the conditional direct skip connection and the class projection layer. Inspired by the skip-z approach [22], we propose conditional skip-z, which integrates skip-z with class-specific information. Skip-z [22] feeds the noise vector $z$ into multiple layers of the generator, instead of just the initial layer. This design enables the generator to utilize the diversity of input noise to enhance the quality and diversity of the generated images. In our approach, we divide $z$ into chunks corresponding to each stage and concatenate each chunk with the class vector $c$. The combined vector $[c,z]$ is then processed through the class projection layer at each stage, enhancing both the quality and diversity of the synthetic images within the class.

Next, we project $[c,z]$ onto token embeddings by the class projection layer which is a conditional layer normalization. Unlike class-conditional methods that use batch normalization [23, 22, 24], we find that layer normalization [25] yields superior performance for ViT-based conditional GANs. This is possibly explained by the fact that layer normalization can preserve positional information learned by the attention mechanism, whereas batch normalization may compromise this information [26]. The approach transforms a layer’s activations $a$ into a class-specific normalized activation $\bar{a}$, which is described as:

where $\hat{\mu}$ and $\hat{\sigma}$ represent the mean and variance of the input, and $\epsilon$ ensures numerical stability. We introduce class-conditional information through parameters $\gamma$ and $\beta$ as linear transformations of the class embedding $c$, where $\gamma:=W_{\gamma}^{\top}[c,z]$ and $\beta:=W_{\beta}^{\top}[c,z]$, with $[c,z]$ denoting the concatenation of $c$ and $z$ in the conditional direct skip-z.

Our framework’s learning strategy is inspired by AC-GAN [27], which incorporates an auxiliary class head in the discriminator. However, the auxiliary classifier in [27] primarily functions to guide diversified image generation, often resulting in suboptimal classification performance. To address this issue, we introduce a multi-loss weighting function that integrates the WGAN-GP loss [28] with the classification cross-entropy loss. This approach ensures that both the generation and classification aspects of the model perform satisfactorily. We first apply the WGAN-GP loss to our GAN. The objective function is defined as follows:

where $f^{1}$ is the source head of the discriminator, $D$ denotes the discriminator feature extractor and $G$ is the generator. $\mathbb{P}_{r}$, $\mathbb{P}_{z}$, and $\mathbb{P}_{\hat{I}}$ donate the distribution of real data $I_{real}$, the normal distribution of a random noise vector $z$, and the distribution of pairs of points from $\mathbb{P}_{r}$ and $\mathbb{P}_{z}$. Each generated sample is associated with a class label $c\sim\mathbb{P}_{c}$ besides $z$. $G$ uses conditioned noise to generate images $I_{\text{syn}}=G(c,z)$. Note that the source head does not involve conditioned information.

To enable conditional learning, we employ a class head $P(C\mid I)$ on the discriminator, the log-likelihood objective function is given by:

both $D$ and $G$ are trained to maximize $\mathcal{L}_{C}$. We find that simply combining $\mathcal{L}_{C}$ and $\mathcal{L}_{S}$ is prone to result in training failure as the WGAN-GP loss value is relatively large at most times, which results in cross-entropy loss overwhelmed during training. In addition, the model tends to perform well on the discriminator for a particular classification task, leading to an imbalance problem.

To this end, we extend the concept from the multi-task learning realm [6] to weigh classification losses and propose a multi-loss weighing loss function, which balances losses of multi-classification tasks by considering the uncertainty of each task. Specifically, let $f^{2}(I)$ be the output of the class head and $\exp\left(-\sigma\right)$ be the trainable scaling parameter, our adapted classification likelihood of the model output through the softmax function can be written as:

with $\sigma\in[-\infty,+\infty]$. The scaling process can be regarded as a Maxwell–Boltzmann distribution, where $\exp(-\sigma)$ is commonly referred to as the temperature of the input. The learnable parameter’s magnitude determines how “flat” the discrete distribution is. The output is then related to uncertainty by using the log-likelihood, which can then be written as:

we write $\mathcal{L}=-\operatorname{logSoftmax}\left(y,f^{2}(I)\right)$ for the cross entropy loss (not scaled). Then, an explicit simplifying assumption is introduced as $\exp(-\sigma)\sum_{c^{\prime}}\exp\left[\exp(-\sigma)f_{c^{\prime}}^{2}(I)\right]\approx\left(\sum_{c^{\prime}}\exp\left(f_{c^{\prime}}^{2}(I)\right)\right)^{\exp(-\sigma)}$. When $\sigma\to 0$, this equation becomes equality. This helps simplify $\operatorname{log}\frac{\sum_{c^{\prime}}\exp[\exp(-\sigma)f^{2}_{c^{\prime}}(I)]}{\left(\sum_{c^{\prime}}\exp({f^{2}_{c^{\prime}}(I))}\right)^{\exp(-\sigma)}}$ to $\sigma$ and empirically demonstrates promising results.

Given multiple classification outputs from the class head, we often define the likelihood of factorizing over the outputs. Our multi-task likelihood is shown as follows:

where $y_{1,2}$ represent classification predictions of real and synthetic data from the class head, respectively. Next, the joint loss of weighted log-likelihood of multi-classification tasks is given as:

$L_{mlw}$ can be regarded as learning the weights of each output loss. A larger temperature value, $\sigma$, results in a reduced contribution to the loss function, whereas a smaller $\sigma$ increases its weight. When $\sigma$ is very small, the function is predominantly regulated by the last three terms. Our experiments indicate that the range of $\sigma$ values impacts model performance during training. Consequently, differently from the approach of [6], we utilize $\exp\left(-\sigma\right)$ as weights for different losses. This allows $\sigma$ to act as an infinitely regularized value, enhancing the model’s performance by providing a more balanced loss function across various tasks. The total loss $L_{total}$ is shown as the following:

We further enhance the performance of synthetic augmentation by employing a selective data augmentation mechanism. To ensure high-fidelity image generation, we first use a truncation method that resamples $z$ from a truncated normal distribution with a truncation value $\tau$. After generating the images conditionally, we leverage the class head to select high-quality augmentation images. Only generated images with prediction values exceeding a threshold $\lambda$ will be used to compute the classification loss function. This approach helps to mitigate the negative impact of unrealistic synthetic data on model training. Our experiments indicate that setting $\tau=0.7$ and $\lambda=0.6$ provides the optimal performance. The pseudo-code of selective data augmentation is shown in the Algorithm 2.

We use a public PatchCamelyon (PCam) benchmark dataset [29] for our experiments. The PCam dataset consists of 327,680 color images obtained from lymph node sections, digitized with a 40x objective, resulting in a pixel resolution of 0.243 microns. Each image patch has a resolution of 96 $\times$ 96 pixels and is classified into two categories based on the presence of metastatic tissue in the central region. The dataset is split into 75% for training, 12.5% for validation, and 12.5% for testing, using a hard-negative mining regime. To simulate scenarios with limited training data, we randomly select only 10% of the training images (32,768 images). The entire test set is used to evaluate model performance.

All experiments are conducted using two Tesla V100 GPUs, each with 16GB of RAM, and implemented with PyTorch. An Adam optimizer with parameters $\beta_{1}=0$, $\beta_{2}=0.99$, and a learning rate of $1e-4$ is used to tune both the generator and discriminator. The batch size is set to 12 for both the generator and discriminator. Training is performed over 500 epochs for all experiments. We select the best epoch based on two criteria: the Frechet Inception Distance (FID) score between the generated images and the validation data, and the classification performance on the validation set.

To comprehensively evaluate the performance of the trained classifier, we employed four evaluation metrics in our experiments: accuracy, area under the ROC curve (AUC), sensitivity, and specificity. Accuracy is useful for datasets with a balanced distribution of categories but can be misleading when categories are unbalanced. AUC, on the other hand, provides an assessment of the model’s overall performance across different thresholds and is suitable for datasets with unbalanced categories. Sensitivity is critical in tasks where identifying all positive samples is essential, as it measures recall. Specificity is important in tasks requiring precise identification of negative samples, as it measures the true negative rate. This combined assessment offers a thorough evaluation of the model’s performance in terms of its accuracy and its ability to correctly identify both positive and negative samples.

We conducted a comparative analysis between our proposed method, MedViTGAN [17], and several representative baseline models under identical settings. The baseline models include ResNet34 [30], ResNet50 CBAM [31], and DenseNet169 [32]. For the baseline models, we present results both with and without utilizing synthetic data generated by our proposed method during training. Each model was executed five times with random initialization to ensure a fair comparison. Mean and standard deviation values of the results are reported for comprehensive evaluation.

Table. I presents the experimental results. Comparing the results of baseline models with and without synthetic data demonstrates a significant improvement in classification performance when using data generated by our method. For instance, compared to the baseline DenseNet169, the DenseNet169 trained with synthetic data shows an improvement of 3.4%, 0.5%, 1.7% and 5.9% in accuracy, AUC, sensitivity and specificity, respectively. Furthermore, our proposed method surpasses MedViTGAN in terms of classification performance. As illustrated in Fig. 2, the synthetic images exhibit comparable quality in fidelity and diversity to real training images. Thus, we believe that our proposed model is effective for handling challenges related to limited training data in histopathology image tasks.

To further validate the effectiveness of our proposed method, we conducted an ablation study to assess the contributions of its two main components: the conditional class projection module and selective data augmentation. The findings are presented in Table II. First, we removed both the conditional class projection and selective data augmentation from the framework to establish a baseline. Then, we gradually added these two modules back into the framework to analyze their individual and combined effects. Compared to the baseline, adding the conditional class projection improved accuracy by 1.5% and specificity by 3.4%, although it slightly underperformed the baseline in terms of AUC and sensitivity. Similarly, incorporating the selective data augmentation mechanism resulted in a 1.3% improvement in accuracy, 0.6% increase in AUC, and 3.8% enhancement in specificity over the baseline, while maintaining sensitivity at nearly the same level as the baseline. The decrease in AUC and sensitivity when using only conditional class projection may be due to the generation of some noisy synthetic data, slightly diminishing the model’s performance in these areas. Using only selective data augmentation avoids this issue, but its improvements in accuracy and specificity are not as pronounced as those achieved with conditional class projection alone. Overall, the results of our proposed method showed a 3.2% increase in accuracy over the baseline, demonstrating the effectiveness of our framework in enhancing classification performance.

Additionally, we analyzed the impact of different loss functions on the generated images. In these experiments, all settings were kept constant except for the loss function. We compared the multi-task loss function proposed by Kendall et al. [6] (denoted as log_) with our multi-loss weighing function (denoted as exp_). The results in Fig.3 demonstrate that with our loss function applied, the fake loss converges rapidly, indicating that the learning of feature separation achieves robust results.