A general approach to improve adversarial robustness of DNNs for medical image segmentation and detection

Abstract

It is known that deep neural networks (DNNs) are vulnerable to adversarial noises. Improving adversarial robustness of DNNs is essential. This is not only because unperceivable adversarial noise is a threat to the performance of DNNs models, but also adversarially robust DNNs have a strong resistance to the white noises that may present everywhere in the actual world. To improve adversarial robustness of DNNs, a variety of adversarial training methods have been proposed. Most of the previous methods are designed under one single application scenario: image classification. However, image segmentation, landmark detection, and object detection are more commonly observed than classifying the entire images in the medical imaging field. Although classification tasks and other tasks (e.g., regression) share some similarities, they also differ in certain ways, e.g., some adversarial training methods use misclassification criteria, which is well-defined in classification but not in regression. These restrictions/limitations hinder application of adversarial training for many medical imaging analysis tasks. In our work, the contributions are as follows: (1) We investigated the existing adversarial training methods and discovered the challenges that make those methods unsuitable for adaptation in segmentation and detection tasks. (2) We modified and adapted some existing adversarial training methods for medical image segmentation and detection tasks. (3) We proposed a general adversarial training method for medical image segmentation and detection. (4) We implemented our method in diverse medical imaging tasks using publicly available datasets, including MRI segmentation, Cephalometric landmark detection, and blood cell detection. The experiments substantiated the effectiveness of our method.

1. INTRODUCTION AND RELATED WORK

It is known that deep neural networks (DNNs) are vulnerable to adversarial perturbations (a.k.a. adversarial noises) [1]. Improving robustness of DNNs against adversarial perturbations, especially the unperceivable adversarial perturbations, is essential for various imaging tasks. The improvement of adversarial robustness also leads to the improvement of robustness against white noises that can be encountered in various fields in the real world [2][3][24].

To improve adversarial robustness of DNNs, a variety of adversarial training methods (a.k.a. defense methods) have been proposed. Most of the previous methods were developed under the simple situation of whole-image classification. However, image segmentation and object detection are more frequently encountered than the classification of whole images in the medical image domain. Some adversarial training methods, e.g., IAAT [4], FAT [5], GAIRAT [6], AME [7], and MART [8], use misclassification criterion to generate adversarial training perturbation. Misclassification criterion is well-defined in classification tasks but not in segmentation and detection tasks. Some adversarial training methods require loss functions specific to classification tasks: MMA [9] is largely based on the soft margin loss that is rarely used in segmentation and detection tasks; LBGAT [10] uses categorical cross-entropy loss for classification of the whole images. These restrictions hinder the application of those previous adversarial training methods to segmentation, landmark detection, and object detection tasks, which are common tasks in the field of medical image analysis.

In this work, our contributions are as follows. (1) We investigated some of these existing adversarial training methods and discovered the challenges that make them hardly be adapted for use in segmentation and detection tasks. (2) We proposed a general adversarial training method for medical image segmentation and detection tasks. Our method not only is effective across various types of medical image tasks, but also eliminates the need for users to adjust the common hyperparameter: training noise level, which makes our method even more general. As a comparison, most of the previous adversarial training methods are very sensitive to training noise level [7]. (3) We applied our method to multiple medical image analysis tasks, including MRI segmentation by nnUnet [11], Cephalometric landmark detection by multi-task Unet [12] and blood cell detection by YOLO v5 [15]. (4) For comparison, we modified and adapted two adversarial training methods, VAT and TEAT, for medical image segmentation and detection tasks. The results of the experiments proved the effectiveness of our method. Sample images are shown in Figure 1. Code is available at https://github.com/SarielMa/SPIE2024-code.

Figure 1.

Examples of adversarial attacks against cephalometric landmark detection (a), segmentation of prostate (b) and heart (c), and blood cell detection, where (d) is the input image and (e) is the detection result. By adding adversarial perturbations to the first-row clean images (clean samples), the second-row noisy images (adversarial samples) are obtained. DNN outputs are significantly changed when the input images are infused with unperceivable adversarial noises.

2. METHODS

2.1. Glossary

“Clean sample” denotes a sample without adversarial perturbations. “Adversarial sample” is a sample with adversarial perturbations. “Clean accuracy” is a model’s accuracy on a test set that contains only clean samples. “Adversarial accuracy” is a model’s accuracy on a test set that contains only adversarial samples. “Clean model” is a naturally trained model without adversarial training. “Robust model” is an adversarially trained model. “Noise level” is the upper bound of adversarial perturbations, used for testing or training. “Adversarial robustness” is the robustness of a DNN model against adversarial perturbations.

2.2. Algorithm

The basic idea of adversarial training is to train a model with adversarial samples. The most well-known adversarial training method is vanilla adversarial training (VAT) [14]. VAT uses a fixed noise level for all the training samples. The performance of VAT is largely affected by the choice of the noise level for training.

For a classifier, the above issue, which is caused by a single training noise level, could be mitigated through the following idea: during training, the generated adversarial sample may be on different noise levels but should not cross the true decision boundary [5][9][7]. Since different clean samples have different distances to the true decision boundary, the adversarial samples should be fused with different levels of noise, i.e., using sample-wise noise level for training. To make it brief, we designate the distance between a clean sample and the decision boundary as “margin”. Motivated by this, the basic idea of our method is to let each training sample find its own “margin”, instead of using a fixed noise level done by VAT. However, decision boundary is not well-defined for medical image segmentation and detection tasks. So, we generalize the concept of sample margin: the margin of a clean sample is the maximum noise level that can be added to the clean sample during training, without harming the DNN’s clean accuracy.

Our proposed adversarial training method is shown in Algorithm 1. x is the input image (clean), $\widehat{y}$ is the DNN model’s output (e.g., segmentation map), and i is the unique ID for the clean sample x with true label y. Given the sample i and its estimated margin E(i), an adversarial sample $\tilde{x}$ is generated (Line 3). The DNN model f(.) applies on the clean and noisy samples to generate the outputs. Then, the loss L₀ of clean data and the loss L₁ of noisy data are calculated and combined. Finally, back-propagation updates the DNN model. It is important not to add too much perturbation, because large perturbation may harm the DNN performance on clean data. To address this issue, we add the condition L₁ < ξ to control the perturbation magnitude. Generally, the loss L₁ on the noisy sample will increase if the noise level increases. To preserve DNN performance on clean data, the loss L₁ should not be significantly larger than the average loss value on clean data. We note that our novel strategy enables our proposed method to handle different tasks if the loss functions are well-defined. Also, our proposed method obviously has the same time complexity as that of VAT.

2.3. Value of ξ

The threshold ξ is used to ascertain whether the noise is too large to be added to a sample for training. Assuming that on the clean training set S = {(x_i, y_i)|i = 1,2, …}, the DNN model f(.) has been well trained; and Loss(.) is the loss function for training f(.) on S. Then, we can obtain the loss values of all training samples, denoted by V = {Loss(f(x), y)| (x, y) ∈ S}. In our experiments, we observed that the distribution of the loss values approximately follows Normal distribution. The mean and standard deviation of these loss values (mean(V), std(V)) can be obtained. Then, the threshold ξ can be mean(V) or mean(V) + 2 ∙ std(V), which depends on the dataset.

2.4. Loss function with multiple items

The training loss function may contain multiple loss terms for some DNN models. For instance, YOLO v5 [15] has three loss terms. If the loss function contains K loss terms, a small modification is made to Algorithm 1: (1) The threshold ξ becomes a vector of K scalars, and each scalar is calculated with each of the K distributions of the loss terms on clean data; (2) Line 9 is modified to:

$$\mathbf{\text{IF}}\left\{\left(L_{1\left[1\right]}<{\xi }_{\left[1\right]}\right)\&\left(L_{1\left[2\right]}<{\xi }_{\left[2\right]}\right)\&\dots \&\left(L_{1\left[K\right]}<{\xi }_{\left[K\right]}\right)\right\}$$ (1)

Algorithm 1.

Robust Training in one epoch

Input:	S is the training set; f(.) is the model; L(.) is the loss function for training f(.).
Parameters:	E is the array of the estimated sample margins; E(i) is the margin of sample i; A is the expansion step size; $\xi$ is a threshold. Adv(.) is an adversarial perturbation function (e.g., PGD). E(i) is initialized as A, and A is selected from grid search.
1.	For each training sample (x, y) with ID i in S
2.	Run the model f(.) on clean sample: $\widehat{y}=f\left(x\right)$
3.	Generate an adversarial sample: $\tilde{x}=\mathit{\text{Adv}}\left(x,y,E\left(i\right)\right)$
4.	Run the model f(.) on adversarial sample: $\tilde{y}=f\left(\tilde{x}\right)$
5.	$L_0=\mathit{\text{Loss}}\left(\widehat{y},y\right)$
6.	$L_1=\mathit{\text{Loss}}\left(\tilde{y},y\right)$
7.	$L=\left(L_0+L_1\right)/2$
8.	Back-propagate from L and update the model f(.)
9.	IF $L_1<\xi$
10.	E(i) = E(i) + Δ
11.	ELSE
12.	$E\left(i\right)=\left({‖x-\tilde{x}‖}_p+E\left(i\right)\right)/2$
Note:	For easy understanding, we use one sample to describe the algorithm. The real algorithm runs in mini batches. ${‖\cdot ‖}_p$ denotes vector Lp norm. Adv(.) ensures ${‖x-\tilde{x}‖}_p\le E\left(i\right)$ where x is sample i.

3. EXPERIMENTS

3.1. Basic settings

All experiments are conducted on a server with Nvidia Tesla V100 GPU and Intel(R) Xeon(R) E5–2698 v4 CPU processor (2.20GHz). Three popular DNN models and five publicly available medical image datasets are employed for method evaluation including 3 segmentation tasks, 1 landmark detection task, 1 object detection task.

3.2. Evaluated defense methods

Two defense methods, TEAT [13] and VAT [14], with different training noise levels are compared with our method. In the experiments, the model trained only on clean data by the standard training method is denoted as “STD”; The model trained by the VAT method with the noise level of ϵ is denoted as “VATϵ ”. The model trained by the TEAT method with the noise level of ϵ is indicated as “TEϵ ”. As for our method, there is no need to tune the training noise level. In the segmentation experiments (Table 3–5), TEAT is excluded because it cannot converge in these tasks.

Table 3.

ADI of nnUnet on Heart MRI test set: “Dice Score” denotes the ADI on clean testing samples; “PGDϵ” denotes PGD attack with noise level ϵ; “IFGSMϵ” denotes IFGSM attack with noise level ϵ; “Avg.” denotes average performance on adversarial testing samples. ADI is the higher the better.

Methods	Dice Score	PGD5	PGD10	PGD15	Avg.	IFGSM5	IFGSM10	IFGSM15	Avg.
VAT25	0.6652	0.6175	0.5711	0.5355	0.5747	0.6410	0.6173	0.6017	0.6200
VAT15	0.7188	0.6579	0.5869	0.5194	0.5881	0.6933	0.6640	0.6329	0.6634
VAT5	0.7525	0.6565	0.5542	0.4617	0.5575	0.7087	0.6667	0.6370	0.6708
Ours	0.9148	0.8476	0.7594	0.6588	0.7553	0.8834	0.8532	0.8255	0.8540
STD	0.7954	0.5197	0.1997	0.0172	0.2455	0.6930	0.6238	0.5779	0.6316

Table 5.

ADI of nnUnet on Prostate MRI test set: “Dice Score” denotes the ADI on clean testing samples; “PGDϵ” denotes PGD attack with noise level ϵ; “IFGSMϵ” denotes IFGSM attack with noise level ϵ; “Avg.” denotes average performance on adversarial testing samples. ADI is the higher the better.

Methods	Dice Score	PGD10	PGD20	PGD40	Avg.	IFGSM10	IFGSM20	IFGSM40	Avg.
VAT40	0.6152	0.5423	0.4614	0.3085	0.4374	0.5817	0.5453	0.4768	0.5346
VAT20	0.6640	0.5606	0.4510	0.2779	0.4298	0.6181	0.5698	0.4772	0.5550
VAT10	0.6569	0.5405	0.4034	0.2063	0.3834	0.5981	0.5522	0.4676	0.5393
Ours	0.8336	0.7991	0.7572	0.6528	0.7364	0.8180	0.8014	0.7649	0.7948
STD	0.7418	0.3012	0.1428	0.0455	0.1632	0.5252	0.4138	0.2796	0.4062

To the best of our knowledge, there are no other adaptive adversarial training methods for image segmentation, landmark detection, and object detection in the medical image application field.

3.3. Method evaluation

PGD-based adversarial attack is widely used for adversarial defense method evaluation [16][17]. The noises for evaluation are generated by PGD whitebox attack [14] with 100 iterations (100-PGD). We also applied IFGSM whitebox attack [18] with 10 iterations (10-IFGSM). The maximum attack noise level in the evaluation is selected such that the prediction accuracy drops to almost 0 (except for Landmark Detection task).

3.4. L2 norm for noise level measurement

It has been observed that adversarial noises have inner structures and patterns (see Figure 1 in [20], Figure 1 in [12], Figure 2 in [19]), which is completely ignored by L-inf norm. This indicates L-inf is not an effective measure of adversarial noise level. Therefore, the adversarial noises are quantified using the L2 Norm in this paper.

3.5. Evaluation on image segmentation

For image segmentation experiment, we apply VAT and our proposed method to a self-configuring DNN, nnUnet [11]. The nnUnet can automatically configure itself, including preprocessing, network architecture, training, and post-processing for the dataset. The inputs of nnUnet are 2D slices of 3D images. Some of the image slices have poor quality, and the nnUnet does not filter them out, which leads to a large std(V) (see Section 2.2) in the training set. As a result, for our method, ξ is set to mean(V) in the experiment (see Section 2).

The “Average Dice Index (ADI)” is used as the evaluation metric.

$$\mathit{\text{ADI}}=\frac{1}{n}\times {\sum }_{i=1}^n\frac{2\times {\mathit{\text{TP}}}_i}{2\times {\mathit{\text{TP}}}_i+{\mathit{\text{FP}}}_i+{\mathit{\text{FN}}}_i}$$ (2)

In formula (2), n is the number of samples in the test set. For the sample i, TP_i is the number of pixels in true-positive area, FP_i is the number of pixels in false-positive area, and FN_i is the number of pixels in false-negative area.

Three public datasets are used in this experiment: Heart MRI dataset, Hippocampus MRI dataset, and Prostate MRI dataset.

The Heart MRI dataset [21] has 20 labeled 3D images: 16 for training, 1 for validation and 3 for testing. The median shape of each 3D image is 115 × 320 × 320, of which 115 is the number of 2D slices. In this experiment, only 2D segmentation is considered, so the input of the model is one 2D slice. The batch size (40), input image size (320 × 256) are self-configured by nnUnet for this dataset. The model is trained for 50 epochs, where each epoch contains 50 iterations. Other training settings are the same as those in [11]. For VAT, we tried three different noise levels (5, 15, 25).

The Hippocampus MRI dataset [21] has 260 labeled 3D images: 208 for training, 17 for validation and 35 for testing. The median shape of each 3D image is 36 × 50 × 35, where 36 is the number of slices. The batch size (366), the input image size (56 × 40) and network structure are self-configured by nnUnet for this dataset. The model is trained for 100 epochs, where each epoch has 50 iterations. Other training settings are the same as those in [11]. For VAT, we tried three different noise levels (5, 10, 15).

The Prostate MRI dataset [21] has 32 labeled 3D images: 25 for training, 2 for validation and 5 for testing. The median shape of each 3D image is 20 × 320 × 319, where 20 is the number of slices. The batch size (32), patch size (320 × 320) and network structure are self-configured by nnUnet for this dataset. The model is trained for 50 epochs, where each epoch has 50 iterations. Other training settings are the same as those in [11]. For VAT, we tried three different noise levels (10, 20, 40).

3.6. Evaluation on landmark detection

For landmark detection experiment, we apply our method, TEAT and VAT to the Multi-Task U-Net [12]. This U-Net model detects the landmarks by regressing both Gaussian heatmap and offset maps of the landmarks simultaneously. The loss function for this model has three loss terms, and therefore ξ is a vector of size 3, in which each value is set to mean(V) + 2 ∙ std(V) (see Section 2). The original loss function contains a Binary Cross-Entropy (BCE) loss term. We find that, by replacing the BCE loss with Dice loss, the robustness of the Multi-Task U-Net is better (see Table 1). The Multi-Task U-Net trained with the original BCE loss is denoted as “STD (BCE)”, and the one trained with the Dice loss is denoted as “STD (Dice)”. Three methods, including TEAT, VAT and our proposed method, are applied to the “STD (Dice)” models.

Table 1.

MRE (mean radial error, mm) on cephalometric landmark detection test set: “MRE” denotes the MRE on clean testing samples; “PGDϵ” denotes PGD attack with noise level ϵ; “IFGSMϵ” denotes IFGSM attack with noise level ϵ; “Avg.” denotes average performance on adversarial testing samples. MRE is the lower the better.

Methods	MRE (mm)	PGD0.5	PGD1	PGD2	Avg.	IFGSM0.5	IFGSM 1	IFGSM 2	Avg.
TE3	2.21	2.31	2.40	2.63	2.45	2.30	2.39	2.63	2.44
TE5	2.41	2.53	2.61	2.83	2.66	2.53	2.61	2.83	2.66
VAT3	2.11	2.19	2.27	2.48	2.31	2.19	2.27	2.47	2.31
VAT5	2.18	2.26	2.34	2.54	2.38	2.26	2.34	2.55	2.38
Ours	1.77	1.89	2.04	2.42	2.12	1.89	2.04	2.39	2.11
STD	1.55	4.35	14.30	38.87	19.17	3.93	10.91	35.02	16.62

The evaluation metric for landmark detection is “Mean Radial Error (MRE)” defined in [12], which measures the Euclidean distance between ground truth and predicted landmarks. MRE is scaled to the unit of mm. The dataset is created for cephalometric landmark detection in IEEE ISBI 2015 Challenge [23]. It contains 400 cephalometric radiographs, which are officially split into 3 sets (Train, Test1, Test2): 150 for Train, 150 for Test1 and 100 for Test2. Because the performance of the models reported in [12] is poor on Test2, we only use Test1. Test1 is further split into 2 sets: 50 in the validation set and 100 in the test set. Each radiograph has 19 manually labeled landmarks of clinical anatomical significance by two expert doctors. The average annotations by two doctors are considered as the ground truth. The images are resized to 200 × 160 in the experiment. The model is trained for 500 epochs. The batch size is 8. Other training settings are the same as those in [12]. For TAET and VAT, we tried different noise levels for training including 3 and 5.

3.7. Evaluation on object detection

For object detection experiment, we apply our method, TEAT and VAT, to the YOLO v5s object detector [15]. The loss function has three loss terms. Therefore ξ is a vector of size 3, of which each value is set as mean(V) + 2 ∙ std(V) (see Section 2).

The metric used for evaluation is mAP 0.5:0.95. The metric used for tuning the hyperparameter is Average Intersection over Union (IOU) of all detected objects. The dataset is named BCCD, which is a small-scale dataset for blood cell detection [22]. The dataset contains 364 2D images, which is split into 3 sets: 294 for training, 33 for validation and 37 for testing. The blood cells on each image are from three classes: Red Blood Cell, White Blood Cell, and Platelets. Each image is resized to 320 × 320 in our experiment. The model is trained for 300 epochs. The batch size is 32. Other training settings are the same as [15]. For TEAT and VAT, we employ different noise levels for training including 10 and 15.

4. RESULTS AND DISCUSSION

We have the following observations from the experiment results (Table 1–5): (1) Compared with other adversarial training methods, our method has the best overall performance: the best average performance on noisy data and the best performance on clean data. This suggests that properly crafted adversarial training noises may avoid harming the model’s performance. (2) The VAT-trained models with larger training noise level have worse performance on smaller-perturbation data, and the VAT-trained models with smaller training noise level have poor performance on larger-perturbation data. This indicates that a fixed noise level is inadequate for robust DNN training. (3) TEAT is also sensitive to the hyperparameter, the training noise level, in our experiment results. For life-critical medical image analysis applications, the DNN models should have a good performance on both clean samples and adversarial samples. Thus, our method with adaptive training perturbation is more suitable and effective than VAT and TEAT which require a manual-fixed training noise level.

5. CONCLUSIONS

Improving the adversarial robustness of DNNs is essential since deep neural networks are vulnerable to adversarial noises. To improve the adversarial robustness of deep neural networks, many adversarial training methods have been proposed recently. However, most of these previous methods can hardly be extended to image segmentation, landmark detection, and object detection tasks, which are common tasks in the medical image analysis field. In this work, we investigated many well-known adversarial training methods and identified the reasons preventing their extension to segmentation and detection tasks. To solve this problem, we proposed a general adversarial training method to make DNNs robust against adversarial noises in medical image segmentation, landmark detection, and object detection tasks. The key idea is to adaptively adjust the training noises for individual training samples so that the robust DNN models demonstrate strong resistance to adversarial noises while preserving high accuracy on clean data. Our novel strategy enables our proposed method to handle different medical analysis tasks given that the loss functions are well-defined, which means our proposed method can be extended to a broader range of applications in the medical application domain. Furthermore, our proposed method eliminates the necessity to tune training noise level, making it more user-friendly. We applied our proposed method to three well-known DNNs on publicly available medical image datasets. For comparison, we also extended and adapted VAT and TEAT to medical segmentation and detection tasks. The experiment results show that our method outperforms the other adversarial training methods in both adversarial accuracy and clean accuracy while having almost the same training time complexity as VAT. Our approach facilitates the development of robust DNNs for more medical applications.

Table 2.

mAP0.5:0.95 on blood cell detection test set: “mAP” denotes the mAP0.5:0.95 on clean testing samples; “PGDϵ” denotes PGD attack with noise level ϵ; “IFGSMϵ” denotes IFGSM attack with noise level ϵ; “Avg.” denotes average performance on adversarial testing samples. mAP is the higher the better.

Methods	mAP	PGD5	PGD10	PGD15	Avg.	IFGSM5	IFGSM10	IFGSM15	Avg.
TE10	0.4859	0.3117	0.1536	0.0626	0.1760	0.3109	0.1499	0.0709	0.1772
TE15	0.2026	0.1595	0.1135	0.0664	0.1131	0.1588	0.1135	0.0717	0.1146
VAT10	0.4722	0.2988	0.1869	0.0765	0.1874	0.2988	0.1868	0.0808	0.1888
VAT15	0.2565	0.2013	0.1375	0.0572	0.1320	0.2015	0.1354	0.0603	0.1324
Ours	0.5989	0.3341	0.2017	0.1428	0.2262	0.3350	0.2106	0.1511	0.2322
STD	0.5993	0.0072	0.0007	0.0005	0.0028	0.0456	0.0102	0.0082	0.0213

Table 4.

ADI of nnUnet on Hippocampus MRI test set: “Dice Score” denotes the ADI on clean testing samples; “PGDϵ” denotes PGD attack with noise level ϵ; “IFGSMϵ” denotes IFGSM attack with noise level ϵ; “Avg.” denotes average performance on adversarial testing samples. ADI is the higher the better.

Methods	Dice Score	PGD2	PGD6	PGD10	Avg.	IFGSM2	IFGSM6	IFGSM10	Avg.
VAT15	0.7112	0.6762	0.5938	0.4833	0.5844	0.6891	0.6600	0.6276	0.6589
VAT10	0.7175	0.6768	0.5953	0.4707	0.5809	0.6949	0.6599	0.6232	0.6593
VAT5	0.7305	0.6866	0.5636	0.4213	0.5572	0.7013	0.6580	0.6195	0.6596
Ours	0.7845	0.7525	0.6752	0.5580	0.6619	0.7587	0.7280	0.6968	0.7278
STD	0.7589	0.5217	0.0379	0.0015	0.1870	0.6572	0.5105	0.3870	0.5182