BSDA: Bayesian Random Semantic Data Augmentation for Medical Image Classification

In image recognition tasks, data augmentation techniques like random flipping, panning, and rotation make the network learn certain invariants [18, 22]. These methods rely heavily on empirical knowledge, and a specific augmentation strategy can only be used for a particular dataset. AutoAugment [4] was the first technique to perform automatic augmentation by searching for a better strategy across numerous solution spaces using reinforcement learning. However, due to the large of its search space, other works (e.g., RA [5], DADA [23]) have also shown strong performance by improving the search algorithm. While automated data augmentation algorithms have shown strong capabilities, the training of deep neural networks remains computationally intensive. Image Erasure [24] and Image Mixing [25, 19, 20] enhance the performance of the network by performing random erasure or mixing of images in the image space. Generative models have also shown strong performance in image generation, with some works [26, 27] developing data augmentation techniques based on image generation. However, these methods require additional training of a generative model for each dataset. A Bayesian data augmentation method proposed in [28] learns the distribution of features from the training set by generating an adversarial network implementing generalized Monte Carlo expectation maximization and samples from it to obtain augmented data. It is worth noting that our method is similar in name but differs in approach. Several studies have explored the efficiency of data augmentation. [17] notes that data augmentation modeled as additive perturbation improves network performance by amplifying and perturbing the singular values of the network’s Jacobian determinant. Similarly, [29] showed a lower bound on the data required for data augmentation, demonstrating that the general use of DA requires an exponential amount of data before it leads to larger bounds.

Semantic data augmentation is achieved by label-preserving translation in the semantic space, generating new sample points in the feature space. Three semantic augmentation schemes involve adding random Gaussian noise with mean 0 and variance $\sigma^{2}$ to the features, interpolating neighbors, and extrapolating [15]. Similarly, OnlineAugment [30] achieves excellent performance by using Gaussian noise to perturb features, but it tends to generate non-class-preserving sample points. Furthermore, [16] used a moving average to estimate the feature covariance matrix online during training and modeled the cross-feature joint noise distribution. ISDA [14] is a novel semantic data augmentation algorithm that enhances dataset diversity by translating training samples in the deep feature space along semantically meaningful directions, improving the generalization performance of deep models with minimal computational cost. A shape space-based feature augmentation method projects multiple image features into a pre-shape space. [31] uses Wasserstein geodesic interpolation to augment data and regularize performance. Moment Exchange is an implicit data augmentation method that enhances recognition models by replacing and interpolating learned features’ moments (mean and standard deviation) between training images. This forces models to extract training signals from these moments, thus improving generalization across benchmark datasets [32].

The BSDA method generates augmented feature $\mathbf{\tilde{a}}$ by adding the semantic magnitude $\mathbf{m}$ to the original feature $\mathbf{a}$ after element-wise multiplication of them with randomly selected semantic directions $\mathbf{d}_{\lambda}$. The formula is as follows:

where $\mathbf{d}\in\left\{0,1\right\}^{k}$ is a binary vector, initially set to all ones and then set to zero with probability $\lambda$, and $\odot$ denotes element-wise multiplication. $\mathbf{m}\in\mathcal{R}^{k}$ represents the semantic magnitude sampled from the augmentable magnitude distribution $p(\mathbf{m}|\mathbf{a})$ when given the original feature vector $\mathbf{a}$. Furthermore, to preserve specific properties of the features, such as low rank [33, 34, 35], we mask semantic directions corresponding to zero feature values. We have:

where $\mathbb{I}_{\mathbf{a}\not=0}$ is an indicator function when $\mathbf{a}\not=0$, the value of this function is 1.

Estimate the magnitude distribution. To obtain $p(\mathbf{m}|\mathbf{a})$, we introduce a model $q_{\phi_{\mathbf{m}}}(\mathbf{m}|\mathbf{a})$ to approximate the true distribution $p(\mathbf{m}|\mathbf{a})$. The Kullback-Leibler (KL) divergence measures the similarity between these two distributions, aiming to make $q_{\phi_{\mathbf{m}}}(\mathbf{m}|\mathbf{a})$ closely match $p(\mathbf{m}|\mathbf{a})$ by maximizing the KL divergence. Thus, our optimization goal is as follows:

Removing the terms that are not related to the parameter $\tilde{\phi_{\mathbf{m}}}$ in $D_{KL}$, we have:

Loss function of BSDA. The first term of Eq. 4 can be calculated easily. The second term estimates the features $\mathbf{a}$ given $\mathbf{m}$, which is more challenging to compute. Drawing inspiration from the design of Variational Autoencoders (VAE) [36], we introduce a reconstruction network for learning this term, rewriting the second term as $p_{\phi_{\mathbf{a}}}(\mathbf{a}|\mathbf{m})$. Then, we obtain the loss function of BSDA as follows:

The second part of Eq. 5 depends on the model, and we use MSE loss in BSDA. Assuming the marginal distribution $p(\mathbf{m})$ follows a normal distribution $\mathcal{N}(0,\mathbf{I})$ and $q_{\phi_{\mathbf{m}}}(\mathbf{m}|\mathbf{a})$ also follows a normal distribution $\mathcal{N}(0,\boldsymbol{\sigma}^{2})$, setting the mean to zero because we aim to learn the offset relative to the original rather than the augmented feature. Thus, the loss function of BSDA is given by:

where $\boldsymbol{\sigma}^{2}$ is estimated variance of BSDA, and $\mathbf{\hat{a}}$ is reconstructed feature using $\mathbf{m}$.

The reparameterization trick. We employ the reparameterization trick to facilitate the computation of the loss function while ensuring the gradient flow for effective backpropagation. The random variable $\mathbf{m}$ can be represented as a deterministic variable $\mathbf{m}=g_{\phi_{\mathbf{m}}}(\boldsymbol{\epsilon},\mathbf{a})$, with $\boldsymbol{\epsilon}\sim\mathcal{N}(0,\mathbf{I})$ being an auxiliary variable with an independent marginal distribution $p(\boldsymbol{\epsilon})$, and $g_{\phi_{\mathbf{m}}}$ being a vector-valued function parameterized by $\phi_{\mathbf{m}}$.

Loss function. Our augmentation method is training with $G_{\Theta}$. For convenience, we denote the loss of $G_{\Theta}$ as $\mathcal{L}_{task}$, which typically uses cross-entropy loss in classification tasks. Therefore, the total loss function is:

Different superscripts on $\mathcal{L}_{task}$ distinguish between augmented and original features. The hyperparameter $\alpha$ is a dynamic value introduced to mitigate the impact of BSDA on the network during the initial stages of training when the network has yet to learn valuable features.

To summarize, the BSDA technique can be integrated into deep networks. We provide the pseudocode of BSDA in Algorithm 1.

VAE. Our approach resembles VAE [36], where a latent variable is estimated and used for reconstruction. However, the relationship between our latent variable $\mathbf{m}$ and the inputs differs from VAE [36], the latent variable representation of the inputs in VAE [36]. Our latent variable $\mathbf{m}$ represents the augmentable range in a given feature $\mathbf{a}$ without altering the label.

ISDA. Our approach differs from ISDA [14] in several key aspects: (a) it does not depend on the task’s loss function; (b) it serves as an implicit data augmentation method, complementing the explicit ISDA. This makes BSDA more convenient for special treatments in particular network branches. For instance, the BSDA module can be applied to both branches simultaneously in feature-decoupled networks, achieving post-decoupled feature enhancement.

Dataset. The MedMNIST+ [37] dataset comprises twelve pre-processed 2D datasets and six pre-processed 3D datasets from selected sources covering primary data modalities (e.g., X-ray, OCT, Ultrasound, CT, Electron Microscope), diverse classification tasks (binary/multi-class, ordinal regression, and multi-label) and data scales (from 100 to 100,000) [37]. We selected nine 2D medical image datasets and six 3D medical image datasets in MedMNIST+ [37] covering twelve modalities, as shown in Figure 4. For more details on the dataset, please refer to Table 1.

Evaluation Protocols We used the MedMNIST+ split training set and validation set to select hyperparameters and reported the results of the test set. The Randomness arising from model selection is often ignored. For instance, does method A outperform method B only because the random search for A got lucky [38]? Therefore, we repeated the process three times with different random seeds. Every number we report is the mean of these repetitions and their estimated standard error. The area under the ROC curve (AUC) and accuracy (ACC) are used as the evaluation metrics.

Implementation Details. We implemented BSDA using PyTorch 2.2.2 and Torchvision 0.17.2 and experimented on an NVIDIA RTX 4090 GPU $\times 2$ with an Intel 14900k CPU. During training, we utilized the AdamW [39] optimizer with a learning rate of 0.001 and employed a learning rate warm-up strategy for the first five epochs. The 2D image size is $224\times 224$, and the 3D image size is $64\times 64\times 64$. To ensure fairness, we maintained consistent training configurations across all experiments. The distribution estimator and reconstruction modules of BSDA consist of two fully connected layers, followed by BatchNorm and GeLU activation.

First, we evaluated state-of-the-art methods on nine 2D medical image datasets in MedMNIST2D [37](BloodMNIST, BreastMNIST, DermaMNIST, OCTMNIST, OrgansMNIST, PathMNIST, PneumoniaMNIST, RetinaMNIST, and TissueMNIST), which include a total of 513,429 samples. We compared BSDA with state-of-the-art methods, including Cutout [40], Mixup [25], CutMix [20], and ISDA [14]., on nine 2D medical image datasets.

We selected five 3D MedMNIST datasets to validate the performance of BSDA. As shown in Table 4. BSDA enhances the performance of ResNet-18 across various 3D datasets, with improvements in both AUC and ACC for most datasets. For example, AdernalMNIST3D shows a substantial accuracy improvement from $71.5\%$ to $89.8\%$ and a slight AUC increase from $89.1\%$ to $89.2\%$, suggesting that BSDA significantly enhances the model’s performance, particularly in accuracy. OrganMNIST3D shows an accuracy increase from $87.2\%$ to $88.7\%$ and an AUC improvement from $99.2\%$ to $99.4\%$. This indicates that BSDA improves the correct classification rate and enhances the model’s ability to distinguish between classes. Similarly, NoduleMNIST3D experiences a rise in accuracy from $84.5\%$ to $86.1\%$ and in AUC from $87.9\%$ to $89.2\%$, showing consistent improvements in both metrics. For FractureMNIST3D, accuracy rises from $53.5\%$ to $55.0\%$, while AUC improves from $72.7\%$ to $74.0\%$. These moderate improvements indicate that BSDA helps in better identifying fractures, though the dataset remains challenging. VesselMNIST3D also benefits, with accuracy increasing from $92.1\%$ to $93.2\%$ and AUC from $90.2\%$ to $91.7\%$, reflecting improved precision and robustness in vessel identification. These results underscore the value of data augmentation techniques in improving both the discriminative power and the accuracy of deep learning models in medical imaging tasks, demonstrating BSDA’s effectiveness in enhancing model performance across diverse medical datasets.

We selected several mainstream architectures in computer vision to conduct experiments to verify the performance of BSDA combined with different models, including ResNet [18], EfficientNet [22], DenseNet [41], TinyViT [42], ViT [43], Swin Transformer [44]. Table 5 shows that the BSDA technique significantly improves accuracy (ACC) and area under the curve (AUC) across different convolutional neural networks, particularly on the DenseNet-121, ViT, and Swin Transformer series. For example, the accuracy of DenseNet-121 improved from 84.9% to 89.4%, and the AUC increased from 96.6% to 96.9%. In most cases, BSDA introduces a small additional time cost, up to 7.3% for EfficientNet-B0. Although BSDA leads to slight performance degradation on some networks (e.g., ResNet-50), overall, it improves model performance with manageable additional computational costs. Therefore, BSDA is a technique worth considering for improving model performance.

We analyzed the effectiveness of BSDA components and compared BSDA by adding random standard Gaussian noise to the latent space features. Table 6 indicates that random noise can somewhat improve model performance, but the improvement is limited. The BSDA model, with the introduction of the indicator function (Eq. 2) and reconstructor components, shows significant improvement over the baseline model, especially in terms of accuracy. Both the indicator function and reconstructor components are essential for BSDA. Notably, model performance degrades considerably when the indicator function and reconstructor components are removed. This is equivalent to eliminating the reconstructor as well as the second term in Eq. 5, which is comparable to simply making $q_{\phi_{\mathbf{m}}}(\mathbf{m}|\mathbf{a})$ similar to the standard Gaussian distribution. The model performance should be comparable to that of sampling noise from a standard Gaussian distribution added to the original features, because it is simultaneously constrained by the loss of the task, thus further making the samples sampled by the model $q_{\phi_{\mathbf{m}}}(\mathbf{m}|\mathbf{a})$ more favorable to the task.

To better understand the BSDA method, we conducted a series of experiments analyzing it from different perspectives. The following sections provide a detailed analysis of the results, including hyperparameter sensitivity analysis and loss weight analysis.

Hyper-parameters Analysis As shown in Figure 5, subplots 5(a) and 5(b) show the AUC and ACC results of the joint features for different hyperparameters $\lambda$ and $U$ compared to the baseline, with red or blue indicating improvements over the baseline. Comparing subplots 5(a) and 5(c) and 5(b) and 5(d), we see that the joint features improve model performance in most cases.

Weight of Loss We conducted a sensitivity analysis of the hyperparameter $\alpha$ of BSDA using ResNet-18 on the BreastMNIST dataset. The results are shown in Figure 6. We chose the hyperparameter $\alpha$ values from 0.1 to 1. It can be observed that the AUC and ACC metrics are stable in almost all cases, indicating that the BSDA method is not sensitive to the hyperparameter $\alpha$. The AUC value is slightly lower than the baseline at $\alpha=0.3$ or $\alpha=0.9$, and the ACC value is somewhat lower than the baseline at $\alpha=0.1$ or $\alpha=0.9$. Empirically, we recommend $\alpha=0.5$ as the default parameter.

We visualize the deep features using t-SNE [45] in Figure 7. Circular markers represent the original features, while the cross markers indicate the augmented features generated by BSDA. The results show that the augmented features are distributed around the original features, indicating that BSDA generates augmented features that are close to the original ones. Figure 7 also shows that the original features wrap tightly around the newly generated features. Intuitively, classifiers learned from the augmented features will be farther from the original features.

We also visualize the features learned by our method compared to other methods (without feature-level or image-level augmentation). Figure 8 shows that BSDA methods (Fig. 8(f)) make the intra-class features more cohesive and more easily separable.