Mixing-AdaSIN: Constructing a De-biased Dataset using Adaptive Structural Instance Normalization and Texture Mixing

Recently, several methods have been proposed to improve COVID-19 patient diagnosis or treatment planning using CT scans [12, 16, 24, 25]. The most methods are often trained and evaluated on single source data; producing models that exploit underlying biases in the data for better predictions. Yet, they fail to generalize when the bias shifts in external data reporting lower performance. This has critical implications, especially in medical imaging where biases are hard to define or accurately quantify. To address this, extensive data augmentation [5] or re-sampling is often employed; though is still limited by collection (multi-institute data) and how to express the bias when it is unknown. Thus, there is a need for methods that mitigate bias in part or fully towards improved model performance and better generalization.

In general, models trained on biased data achieve high accuracy and despite their capacity lack the motivation to learn the complexity of the intended task. For instance, a model trained on our in-house biased dataset with COVID-19 and bacterial pneumonia reported 97.18% (f1-score) on the validation set, yet degrades to 33.55% when evaluated on an unbiased test-set. Here, we believe bias may originate from varied CT protocols based on exam purpose, scanners, and contrast delivery requirements [3]. Though contrast CT is a standard protocol, it is challenging for practitioners to meet requirements during the pandemic due to extra processes such as contrast agent injection and disinfection [21, 11]. Further, protocols may also vary for other pneumonia that exhibit similar imaging characteristics with COVID-19 in CT. Consequently, we believe biased datasets are often constructed unexpectedly and sometimes unavoidably due to the aforementioned factors.

Among the existing techniques proposed to remove model dependence on bias, augmentation is a de-facto technique for medical images; with other methods pre-defining the bias the trained model should be independent of. This assumes bias is easily defined, but one has to take extra care in the medical setup where such assumptions do not hold. To address this, we propose to construct a de-biased dataset where spurious features based on texture information become uninformative for accurate prediction. A key motivation is that accurate prediction of COVID-19 from other pneumonia’s is dependent on the CT protocols related to texture features and contrast. Thus, we propose to generate COVID-19 CTs with bacterial pneumonia protocol characteristics and vice versa for bacterial pneumonia with COVID-19, respectively.

Specifically, we propose Mixing-AdaSIN; a generative model based bias removal framework that leverages Adaptive Structural Instance Normalization (AdaSIN) and texture mixing to generate de-biased images used to train classifiers robust to the original bias. For image generation, we employ two main components: (a) texture mixing, which enables realistic image generation, and (b) AdaSIN, which guarantees structural consistency and prevents bias retainment in the input image via modifying the distribution of the structure feature. To prevent incorrect image generation, we first pre-train a contrastive encoder [8] to learn key CT features and later use it to search similar image pairs for the texture mixing step in the proposed generative framework. For evaluation, we construct biased train/validation sets based on the CT protocol and an unbiased test set from the non-overlapping CT protocols of the train/validation sets, respectively. The proposed method reports high bias mitigation performance (66.97% to 80.97%) and shows improved generalization performance when verified on an external dataset. The main contributions of this work are summarized as follows:

• •

  We propose a generative model that can sufficiently mitigate the bias present in the training dataset using AdaSIN and texture mixing.

• •

  We show that the use of contrastive learning for texture transfer pair selection prevents incorrect image generation.

• •

  We constructed a biased COVID-19 vs. bacterial pneumonia dataset to verify bias mitigation performance. Our approach not only enabled improvements for standard classification models such as ResNet18 but also current state-of-the-art COVID-19 classification models.

• •

  We also demonstrate the generalization performance of our classifier trained with the de-biased data on an external public dataset.

Another line of work is the augmentation based models that utilize techniques such as arbitrary style transfer. Geirhos et al. [5] proposed shape-ResNet, an approach that finetunes a model pre-trained on AdaIN [10] generated images. Li et al. [18] proposed a multitask learning approach that enables accurate prediction by either using shape, texture, or both types of features. Though successful, a key drawback is the heavy reliance on artistic image generation techniques that may be detrimental for medical images and subsequent diagnoses. To address this, our approach is able to capture subtle differences in the image to generate consistent texture updated images. Thus, classifiers trained on the generated images can avoid subtle bias information.

Given a chest CT dataset $\mathcal{D}=\{X^{1},...,X^{n}\}$ where $X^{i}$ is a set of 2D CT slices $X^{i}=\{x_{1}^{i},...,x_{m}^{i}\}$ each with its label $y^{i}\in\{0,1\}$ denoting bacterial pneumonia and COVID-19 samples, respectively. Our goal is to generate a dataset $\mathcal{D}^{\prime}=\{X^{1\prime},...,X^{n\prime}\}$ where each $X^{i\prime}$ is a set of texture updated CT slices that may contain the bias information of the other label CT protocol. To achieve this, we first pre-train a contrastive encoder using $\mathcal{D}$ to learn representative slice features and then use it for slice similarity search i.e. $x_{1}^{i},x_{2}^{j}\sim\mathcal{D},y^{i}\neq y^{j}$ with semantically similar image structures. Second, to generate $x_{1}^{i\prime}$ we feed searched pairs ${x_{1}^{i},x_{2}^{j}}$ to an encoder network $E$ that outputs structural and texture features used as input for a generator network $G$. Here, AdaSIN and texture mixing is employed on the respective features for improved generation. Lastly, following standard practice for adversarial-based methods, we also employ a discriminator network $D$ to improve the quality of generation. To enforce style and content similarity, a pre-trained VGG19 [23] is used to optimize content loss between $x_{1}^{\prime}$ and $x_{1}$, and a style loss between $x_{1}^{\prime}$ and $x_{2}$, respectively. Through this process, we can generate images to construct $\mathcal{D}^{\prime}$ and then combine $\mathcal{D}$ and $\mathcal{D}^{\prime}$ to train a classifier that will learn to ignore bias information. The overall framework is presented in Figure 1 with specific details of each step categorized below.

In this work, we have proposed a novel methodology to train a COVID-19 vs. bacterial pneumonia classifier that is robust to bias information present on training data. We constructed an in-house biased training dataset in conjunction with an unbiased testing dataset and proved that our method allowed the classifier to learn the appropriate features to correctly predict the labels without considering bias information and achieved better generalization. All of this was possible thanks to an adequate image generation design that relies on two major components: (a) texture mixing, which enables realistic image generation, and (b) AdaSIN, which prevents bias flow from the input to the output image in the generation stage, while maintaining structural consistency. We proved the benefits of our pipeline by achieving the best bias mitigation performance when compared to other related methods in both our in-house dataset as well as in an external dataset. Considering that biases can be easily included when constructing datasets, we hope that our findings help to improve performance in various medical tasks.