Learning to Generalize towards Unseen Domains via a Content-Aware Style Invariant Model for Disease Detection from Chest X-rays

TL;DR

This paper introduces a style-invariant model with content-aware augmentation techniques to improve the generalization of chest X-ray disease detection across unseen domains, outperforming existing methods.

Contribution

The proposed approach employs novel style randomization modules at image and feature levels, enabling diversified augmentation and enhanced cross-domain robustness in disease classification.

Findings

Achieves higher AUCs on unseen datasets compared to state-of-the-art models.

Utilizes style sampling from the full CXR value range for better augmentation.

Demonstrates statistically significant improvements in thoracic disease detection.

Abstract

Performance degradation due to distribution discrepancy is a longstanding challenge in intelligent imaging, particularly for chest X-rays (CXRs). Recent studies have demonstrated that CNNs are biased toward styles (e.g., uninformative textures) rather than content (e.g., shape), in stark contrast to the human vision system. Radiologists tend to learn visual cues from CXRs and thus perform well across multiple domains. Motivated by this, we employ the novel on-the-fly style randomization modules at both image (SRM-IL) and feature (SRM-FL) levels to create rich style perturbed features while keeping the content intact for robust cross-domain performance. Previous methods simulate unseen domains by constructing new styles via interpolation or swapping styles from existing data, limiting them to available source domains during training. However, SRM-IL samples the style statistics from the possible value range of a CXR image instead of the training data to achieve more diversified augmentations. Moreover, we utilize pixel-wise learnable parameters in the SRM-FL compared to pre-defined channel-wise mean and standard deviations as style embeddings for capturing more representative style features. Additionally, we leverage consistency regularizations on global semantic features and predictive distributions from with and without style-perturbed versions of the same CXR to tweak the model's sensitivity toward content markers for accurate predictions. Our proposed method, trained on CheXpert and MIMIC-CXR datasets, achieves 77.32$\pm$0.35, 88.38$\pm$0.19, 82.63$\pm$0.13 AUCs(%) on the unseen domain test datasets, i.e., BRAX, VinDr-CXR, and NIH chest X-ray14, respectively, compared to 75.56$\pm$0.80, 87.57$\pm$0.46, 82.07$\pm$0.19 from state-of-the-art models on five-fold cross-validation with statistically significant results in thoracic disease classification.

Tables

Table 1. TABLE II: Training Hyperparameter Settings

Training configs	Values
Weight initialization	pre-trained on ImageNet
Optimizer	Adam
Optimizer momentum	${\beta }_1$=0.9, ${\beta }_2$=0.999
Learning rate	0.0001
Batch size	50
Gradient accumulation step	4
Training epochs	50
Input image dimension	224×224
Augmentations	Horizontal flip and Random cropping [1]
Classification Loss	Focal loss with ${\alpha }^t$=0.25 and ${\gamma }^{\prime }$=2 [51]
Exp. Mov. Avg. (EMA) [52]	0.997

Table 2. TABLE VII: Analysis of the Feature-Level Style Randomization Module. The Best Result for Each Disease Class is Shown in Red .

Part-1: Investigation of the effect of randomization stage
Method	BRAX	NIH Chest X-Ray14	VinDr-CXR
DB-1	77.00±0.39	82.34±0.08	87.63±0.63
DB-2	77.32±0.35	82.63±0.13	88.38±0.19
DB-3	76.75±0.44	82.34±0.13	88.15±0.35
Part-2: Effect of learnable style embeddings
Method	BRAX	NIH Chest X-Ray14	VinDr-CXR
Pre-defined	77.04±0.48	82.44±0.13	88.09±0.48
Learnable	77.32±0.35	82.63±0.13	88.38±0.19
Part-3: Analysis of the style net architecture
Method	BRAX	NIH Chest X-Ray14	VinDr-CXR
Conv	77.32±0.35	82.63±0.13	88.38±0.19
Conv+ReLU	77.17±0.30	82.30±0.10	88.23±0.30
Conv+BN	77.26±0.40	82.38±0.14	88.19±0.23

Equations

$$\mu (X_1{)}_i=\frac{1}{WH}\sum _{j=1}^W\sum _{k=1}^H(X_1{)}_{i,j,k}$$

$$\displaystyle\sigma(X_{1})_{i}=\sqrt{\frac{1}{WH}\sum_{j=1}^{W}\sum_{k=1}^{H}\Bigl{[}(X_{1})_{i,j,k}-\mu(X_{1})_{i}\Bigl{]}^{2}+\epsilon}$$

$$\mathrm{IN}({\mathbf{X}}_1)=\mathbf{\gamma }\left(\frac{{\mathbf{X}}_1-\mu ({\mathbf{X}}_1)}{\sigma ({\mathbf{X}}_1)}\right)+\mathbf{\beta }$$

$$\mathrm{AdaIN}({\mathbf{X}}_1,{\mathbf{X}}_2)=\sigma ({\mathbf{X}}_2)\left(\frac{{\mathbf{X}}_1-\mu ({\mathbf{X}}_1)}{\sigma ({\mathbf{X}}_1)}\right)+\mu ({\mathbf{X}}_2)$$

$$\mathrm{KLD}({\mathbf{y}}_{\mathbf{pred}},{\mathbf{y}}_{\mathbf{true}})$$

$${\mathrm{KLD}}^{\prime }({\mathbf{y}}_1,{\mathbf{y}}_2)$$

$$SRM-IL(\mathbf{I},\mu (\mathcal{S}),\sigma (\mathcal{S}))={\mathbf{I}}_{\mathbf{s}}=\sigma (\mathcal{S})\left(\frac{\mathbf{I}-\mu (\mathbf{I})}{\sigma (\mathbf{I})}\right)+\mu (\mathcal{S})$$

$$SRM-FL({\mathbf{X}}_{\mathbf{c}},{\mathbf{X}}_{\mathbf{rs}})={\mathbf{X}}_{\mathbf{s}}={\varphi }_{\gamma }({\mathbf{X}}_{\mathbf{rs}})\left(\frac{{\mathbf{X}}_{\mathbf{c}}-\mu ({\mathbf{X}}_{\mathbf{c}})}{\sigma ({\mathbf{X}}_{\mathbf{c}})}\right)+{\varphi }_{\beta }({\mathbf{X}}_{\mathbf{rs}})$$

$${\mathcal{L}}_c$$

$${\mathcal{L}}_s$$

$${\mathcal{L}}_{\varphi }$$

$${\mathbf{Z}}_{\mathbf{s}}=SRM-FL({\mathbf{Z}}_1,{\mathbf{Z}}_2)$$

$${\mathcal{L}}_{cls}$$

$$p_i^t$$

$${\alpha }_i^t$$

$${\mathcal{L}}_{ccr}=||{\mathbf{F}}_{\mathbf{s}}-\mathbf{F}|{|}_F^2$$

$${\mathcal{L}}_{pdr}={\mathrm{KLD}}^{\prime }({\mathbf{p}}^{\mathbf{s}},\mathbf{p})$$

$${\mathcal{L}}_{cons}=\frac{{\mathcal{L}}_{ccr}+{\mathcal{L}}_{pdr}}{2}$$

$${\mathcal{L}}_{total}={\mathcal{L}}_{cls}+{\mathcal{L}}_{cons}$$

Full text

I Introduction

Deep learning algorithms have demonstrated remarkable performance in pathology detection from chest X-ray (CXR) images in recent years [1, 2, 3, 4]. However, despite the promising performance of current deep learning architectures, one major bottleneck is the performance drop on unseen domain data compared to internal data due to distribution discrepancy, also known as domain shift [5, 6, 7, 8]. This phenomenon hinders the practical applicability of deep learning models for medical image analyses in real-world scenarios. A number of approaches, i.e., domain adaptation (DA) [9, 10, 11] and domain generalization (DG) [12, 13, 14, 15] methods, have been developed over time to address this domain shift.

Recent studies have demonstrated that convolutional neural networks (CNNs) exhibit a strong bias toward style rather than content for classification [16, 17]. Geirhos et al.[17] showed that CNNs are vulnerable to any changes in image statistics from those on which they are trained. As a result, the CNNs are intrinsically more susceptible to domain shift compared to the human visual system. However, Geirhos et al.[17] demonstrated that the heavy texture bias nature of the CNNs can be overcome and tuned toward content bias by utilizing a suitable dataset for training. Jackson et al.[18] showed that data augmentation via style transfer with artistic paintings could improve the model’s robustness to domain shift for computer vision tasks for natural images. Several DG methods for natural images [19, 20, 21, 22] have recently been developed to leverage neural style manipulation for improving cross-domain performances.

Motivated by the impact of neural style randomization on DG for natural images, we experiment with the style statistics of CXR images at both the image and feature levels. First, we plot the mean and standard deviation of CXR images (i.e., image-level style features) from three large-scale thoracic diseases datasets, i.e., CheXpert [23], MIMIC-CXR [24], and BRAX [25]. Inspired by the observation from the literature on natural images that feature-level style statistics can often characterize visual domains [19], we also plot the 2D t-SNE [26] visualization of feature-level style statistics (i.e., concatenation of means and standard deviations [27]) computed from the first dense block of DenseNet-121 [28] trained on the three thoracic disease datasets. The plots are illustrated in Fig. 1. We can observe from Fig. 1a that a domain distribution gap based on style statistics exists among these datasets at the image level. The t-SNE plot of the feature statistics in Fig. 1b demonstrates that the model captures this uninformative domain information of styles specific to each thoracic disease dataset, which is reflected by the separable clusters. In addition, although CheXpert and MIMIC-CXR are similar at the image level based on simple image statistics, they vary vastly at the feature level because CNNs can extract more high-level and complex features from the images. In real-world scenarios, radiologists tend to learn visual cues from CXRs, rather than uninformative textures and thus perform well across multiple domains. Therefore, in medical imaging for pathology diagnosis from CXR images, robust machine-learning models should not rely on these global image statistics. Instead, they should extract domain-invariant features that are style-invariant and content-biased (pathology traits).

However, to the best of our knowledge, no previous study has utilized the application of neural style randomization for domain-agnostic thoracic disease detection from CXR images. DG approaches for natural images based on style randomization [19, 29, 30] diversify the training data by synthesizing new styles from the existing sources available during training to simulate unseen domain data. As a result, these methods’ augmentation range is sensitive and limited to the style diversity range of the source datasets, which may result in suboptimal performance. In addition, they utilized pre-defined parameters, i.e., mean and standard deviation, as the style embeddings [19, 29, 22]. Learnable pixel-wise style embeddings may improve the performance further. Moreover, consistency regularizations between global semantic features and probability distributions of different augmented versions of the same CXR image are not well studied.

In this work, we propose a novel domain-agnostic, content-biased, and style-invariant DG framework that can be trained in single or multi-source domain settings without explicit domain labels. We employ style randomization modules (SRMs) at the image and feature levels to hierarchically augment the style while preserving the pathology-specific content traits. Utilizing the prior maximum and minimum pixel value knowledge of a normalized CXR image, we construct a set for sampling mean and standard deviation to use them for switching style characteristics at the image level. For feature-level style augmentation, we deploy a learnable SRM module that generates the pixel-wise style embeddings instead of setting pre-defined parameters, i.e., channel-wise mean and standard deviations as affine parameter values [27], for style manipulation. In addition, we employ two consistency regularization losses to enforce the consistency of semantic global feature maps and the predicted probability distributions, respectively, between the original CXR and the style-perturbed CXR. The variability of styles between different perturbations of the same CXR with content information intact makes the style attribute unreliable and makes the model focus more on pathology-specific content in making a decision. The contributions of this paper are summarized as follows:

• •

We employ a novel SRM at the image level that uses the prior knowledge of possible maximum and minimum pixel values of a CXR to create a set to randomly sample and swap style characteristics.

• •

We propose a novel feature-level SRM module that learns the pixel-level affine parameters instead of utilizing pre-defined style embeddings for creating more rich, diversified style-perturbed features on the fly.

• •

We enforce the consistency regularization losses between the global semantic feature maps and the probability distributions for CXRs with and without style characteristics perturbed to guide the model toward content-related cues.

• •

We perform extensive experiments with five large-scale datasets, i.e., CheXpert [23], MIMIC-CXR [24], Brax [25], NIH chest X-ray14 [31], and VinDr-CXR [32]. Our proposed framework demonstrates consistent improvements and outperforms competing approaches under different cross-domain settings.

II Related work

Domain Adaptation (DA): Most past DA approaches have been developed to minimize measures of distances between the source domain and the target domain features [34, 35]. Recent methods for DA can be categorized into adversarial training-based domain alignment [9, 10], Fourier transform-based methods [33], generative adversarial networks (GANs) based methods (unpaired image-to-image translation into source/target domain) [36, 37], high-level feature map alignment based on similarity measurement [11, 10], gradient reversal based alignment methods [38], and others.

Luo et al.[9] adopted class-wise adversarial training to alleviate feature differences across common categories of two datasets to mine the knowledge from the external thoracic disease dataset to improve the performance of the classifier on the internal thoracic disease dataset. However, the discriminators require explicit domain labels. Zhang et al.[10] proposed a domain, instance, and perturbation invariance learning-based framework for utilizing unsupervised domain adaptive thorax disease classification. However, this method requires access to the target domain, which is not often available. Diao et al.[37] proposed a histogram-based GAN methodology for DA that captured the variations in global intensity changes by transforming histograms along with individual pixel intensities. This method is adapted at test time, which is often inconvenient. Yang et al.[33] employed the Fourier transform to swap the low-frequency amplitude components between the source and target images while keeping the phase components unaltered to diversify the training dataset. However, this augmentation method is more image-based than goal-oriented.

A significant limitation of the DA methods is that they require explicit domain labels, i.e., source and target domains or multiple domains, which is often unavailable in most real-world medical applications due to annotation cost or privacy reasons. In addition, most real-world medical datasets usually consist of data from multiple medical sites using various scanners, i.e., ambiguous domain labels. Moreover, they require re-training during application on new target domains, which is inconvenient and impractical for deploying models in real-world medical image analysis applications.

Domain Generalization (DG): DG approaches aim to utilize the source domains only during training and generalize toward any unseen target domain without requiring prior access or re-training [39]. DG is considered more challenging than DA approaches, as the target domain samples are not exposed during the training phase. DG methods can be categorized into three broad approaches, i.e., augmentation-based methods [12, 13, 14], alignment-based methods [40, 41], and meta-learning-based methods [15, 42]. DG methods can be further divided based on the number of source domains utilized during training: multi-source DG and single-source DG. Multi-source DG assumes multiple relevant domains are available during training. Utilizing multiple source domains allows a deep learning model to discover stable patterns across source domains to better generalize to unseen domains. In contrast, the single-source DG is trained with only one source domain and is expected to generalize toward unseen domains.

Wang et al.[12] performed the mixup operation of CXRs from multiple domains to create a virtual domain and utilized ensemble learning to improve performance in unseen domains. However, this approach requires at least two available domains to perform the mixup. Zhang et al.[13] exploited a series of extensive deep-stacked augmentations during training to simulate the expected domain shift to address DG on unseen data. However, data augmentation-based methods often require carefully selecting augmentation methods to resemble the unseen domain’s appearance.

Recent literature studies have shown that CNNs demonstrate a strong bias toward style rather than content and are sensitive to any shift from image statistics that they are trained on, limiting their DG capability [16, 17]. Motivated by the impact of neural style transfer on the domain shift, a number of DG methods have been proposed based on neural style randomization. DG methods based on style randomization can be divided into two categories, i.e., feature-level and image-level randomizations.

The feature-level randomization includes the perturbation of feature statistics from intermediate convolutional layers of the backbone model. Nuriel et al.[29] applied adaptive instance normalization to swap the style statistics between activations of permuted and non-permuted samples in a mini-batch to reduce the representation of global statistics in the classifier model. However, they merely swap the style statistics, limiting the sample diversity to the available source domains. Zhou et al.[19] proposed a novel parameter-free module MixStyle that synthesizes new domains by probabilistically mixing feature statistics between instances. Nevertheless, insufficient diversity still exists due to reliance on source domains. Nam et al.[21] utilized adaptive instance normalization along with adversarial learning to reduce the intrinsic style bias. However, the style-biased network increases training cost and time. Tang et al.[30] proposed two novel modules named CrossNorm and SelfNorm for swapping feature statistics and recalibrating the statistics through attention to both diversify and bridge the gap between source and target domains. Though the SelfNorm is dynamic, the CrossNorm is limited to training data. Wang et al.[22] developed an encoder-decoder-like architecture named FSR to integrate random noise with the original style parameters to generate new style embeddings. However, they use the pre-defined parameters, i.e., channel-wise mean and standard deviation, as the style embeddings. Instead of pre-defined parameters, pixel-wise learnable style parameters may yield more diverse and hard style statistics, resulting in a more robust DG model.

The image-level style randomization-based methods try to diversify the training dataset by constructing hard stylized samples. Zhong et al.[20] proposed a novel augmentation approach, named adversarial style augmentation (AdvStyle), to generate diverse stylized images by adversarially learning the image-level style features to reduce overfitting on the source domain and be robust towards style variations of unknown domains. However, the perturbation diversity depends on the adversarial gradient. Yamashita et al.[14] introduced a form of data augmentation named STRAP, based on random style transfer with medically irrelevant style sources for improving DG performance. However, this approach requires training an additional model to perform the stylization.

Previous image-level and feature-level style randomization-based DG approaches synthesize style or swap style statistics based on the source domains available during training. Thus, the diversity of the augmented samples is limited and sensitive to the available source domains. Furthermore, the exploitation of semantic content consistency between the original and perturbed activation feature maps will allow for robust feature learning and maximize the effectiveness of style augmentation for the out-of-distribution performance of the model.

Overall, to the best of our knowledge, the application of neural style randomization for deriving a DG model for thoracic disease detection from CXRs is not well studied. A systematic exploration of the potential of extracting content-specific (pathology-dependent) but style-invariant domain-agnostic features may further improve DG performance. We have provided a summary of related research in Table I.

III Methodology

III-A Problem Definition

Domain generalization focuses on training on single or multiple source domains $\{\mathcal{D}_{S1},\mathcal{D}_{S2},\cdots\}$ and is expected to generalize toward any arbitrary unseen target domains $\{\mathcal{D}_{T1},\mathcal{D}_{T2},\cdots\}$ with different data distributions. The thoracic disease datasets have significant domain shifts, which may be caused by, but not limited to, various acquisition parameters (exposure shift), manufacturer shift (different hardware, scanning protocol), demographic variations, and so on [12, 43]. These variations introduce potential confounders and other sources of bias. Following [12], we test our proposed method on different datasets compared to the training set for cross-domain performance, referred to as the unseen domain test datasets (out-of-distribution) in this manuscript.

III-B Preliminaries

III-B1 Instance Normalization

Ulyanov et al.[44] proposed instance normalization (IN) for better stylization performance. Let the activation feature map, $\mathbf{X_{1}}\in\mathbb{R}^{D\times W\times H}$, where, $H$ and $W$ indicate spatial dimensions, and $D$ is the number of channels. IN of feature matrix $\mathbf{X_{1}}$ is formulated as,

[TABLE]

Here, $i\in\{1,2,\ldots,D\}$, $\boldsymbol{\gamma},\boldsymbol{\beta}\in\mathbb{R}^{D}$ are learnable affine vectors, and $\epsilon$ is a very small constant for stability.

III-B2 Adaptive Instance Normalization

Huang et al.[27] proposed a parameter-less module named adaptive instance normalization (AdaIN) that can change the style of a convolutional feature activation map to any given arbitrary style utilizing the convolutional feature’s channel-wise mean and standard deviation as the style statistics. For example, if we want to convert the style of feature map $\mathbf{X_{1}}$ to that of $\mathbf{X_{2}}\in\mathbb{R}^{D\times W\times H}$, the channel-wise mean and standard deviation of $\mathbf{X_{2}}$ will be computed and be used as affine parameters in instance normalization (3) for stylization.

[TABLE]

In this paper, we will leverage the AdaIN to perturb the style statistics at the image level. For feature-level style perturbation, we will learn the affine parameters for each pixel level for any arbitrary styles instead of using pre-defined channel-wise mean and standard deviation parameters.

III-B3 Kullback-Leibler Divergence loss

The Kullbeck-Leibler divergence (KLD) loss is generally applied to minimize the difference between two distributions [45]. For example, if we try to reduce the distribution gap between distributions, $\mathbf{y_{pred}}$ and $\mathbf{y_{true}}$, the loss is defined as:

[TABLE]

Where $\mathbf{y_{true}}$ is the target distribution and $\mathbf{y_{pred}}$ is generally the output from a model. In our proposed method, we will derive two predictive distributions, $\mathbf{y_{1}}$, and $\mathbf{y_{2}}$, one for images with style perturbed and the other without style perturbed. As there are no ground truth predictive distributions, to minimize the discrepancy between these predicted predictive distributions, we define the loss as:

[TABLE]

Note that * indicates that the gradient is not propagated through that parameter, as suggested by Miyato et al.[46], to avoid the model collapse issue.

III-C Image-Level Style Perturbation (SRM-IL)

Let the possible maximum and minimum value of a pixel of a normalized CXR image is $x_{max}$ and $x_{min}$, respectively. Regardless of any domain invariance, the style parameter, such as channel-wise mean and standard deviation, will be limited to this range. For any CXR image $\mathbf{I}\in\mathbb{R}^{w\times h}$, we randomly sample $\mu(\mathcal{S})$ and $\sigma(\mathcal{S})$ from the set of values $\mathcal{S}=[x_{min},x_{max}]$ and use them to transfer the style of the CXR image utilizing AdaIN [27] (4).

[TABLE]

Here, $h$ and $w$ indicate spatial dimensions. For typical normalized CXR images, the values range from 0 to 255. For other modalities, the maximum and minimum values can be set accordingly. Contrary to the previous methods, we do not utilize existing samples from the source domains for style randomization. As a result, our style diversity is not limited to source domains and can create a more diverse set of global statistics.

III-D Feature-Level Style Perturbation (SRM-FL)

In medical image segmentation, it has been demonstrated that multiple stacked augmentations or strong augmentations can effectively improve a model’s robustness to domain shift [13]. Inspired by this finding, we apply the style perturbation not only at the image level but also at the feature level. The feature-level style randomization module takes content ($\mathbf{X_{c}}\in\mathbb{R}^{D\times W\times H}$) and reference style ($\mathbf{X_{rs}}\in\mathbb{R}^{D\times W\times H}$) activation maps and transfers the style of the content activation maps to that of the style activation maps. We deploy two style learning nets, i.e., GammaNet ($\phi_{\gamma}$) and BetaNet ($\phi_{\beta}$), to learn the pixel-level affine transformation parameters for transferring style to the content activation maps. The GammaNet and BetaNet shares similar architecture, comprised of four convolutional layers. The initial and the last convolutional layers contain kernels of 1×1 with the same number of channels as the input feature map. The second and third convolutional layers have filters with 3×3 dimensions, with channel dimensions following the squeeze and excitation block structure [47], i.e., the second block reduces the channel number, while the third block expands the channel number. The style nets do not contain any non-linear activation functions or normalization layers [48]. The output from these two nets, $\phi_{\gamma}(\mathbf{X_{rs}})\in\mathbb{R}^{D\times W\times H}$, and $\phi_{\beta}(\mathbf{X_{rs}})\in\mathbb{R}^{D\times W\times H}$, are set as the affine transformation parameters in the IN function to transfer the style statistics.

[TABLE]

To train the style nets, we utilize the content and style loss [49]. The content and style loss are defined as:

[TABLE]

Here, $||\cdot||_{F}$ denotes the Frobenius norm, and $\mathcal{G}(\cdot)$ represents the Gram matrix. The $\eta$ is a hyperparameter.

III-E Domain Agnostic Network

The overall proposed framework is illustrated in Fig. 2. We adopt DenseIBN-121 [50], pre-trained on the ImageNet dataset, as the backbone of our framework. Furthermore, we replace the last fully connected (FC) layer with an FC layer with $N$ output neurons, where $N$ denotes the class number.

We take a CXR image $\mathbf{I}$ and apply image-level style perturbation to get perturbed CXR image $\mathbf{I_{s}}$. The pathology traits or content information remain the same in both images. The stylized CXR image $\mathbf{I_{s}}$ is passed through the dense block-1,2 to generate the intermediate feature space $\mathbf{Z_{1}}\in\mathbb{R}^{D\times W\times H}$. Let $\mathbf{Z_{2}}\in\mathbb{R}^{D\times W\times H}$ be the intermediate feature matrix after the dense block-1,2 for a randomly sampled stylized CXR image $\mathbf{I_{s}}^{\prime}$ from the training mini-batch. We apply the SRM-FL on $\mathbf{Z_{1}}$ and $\mathbf{Z_{2}}$ to switch the feature statistics of $\mathbf{Z_{1}}$ with that of $\mathbf{Z_{2}}$.

[TABLE]

Here, $\operatorname{SRM-FL}(\cdot)$ denotes the feature level style randomization module. Afterward, the feature space $\mathbf{Z_{s}}\in\mathbb{R}^{D\times W\times H}$ is passed to the dense block-3,4 to generate the global feature map $\mathbf{F_{s}}\in\mathbb{R}^{D^{\prime}\times W^{\prime}\times H^{\prime}}$. Here, $D^{\prime}$, $H^{\prime}$, and $W^{\prime}$ denote the feature map depth, height, and width, respectively.

Classification Loss: The feature space $\mathbf{F_{s}}$ is pooled and fed to the classifier network, which consists of one FC layer and one sigmoid layer. Let the predicted probability vector $\mathbf{p^{s}}\in\mathbb{R}^{N}$. $N$ is the total number of pathologies. The pathological ground truths are expressed as a N-dimensional label vector, $\mathbf{L}=[l_{1},\ldots,l_{i},\ldots,l_{N}]$ where $l_{i}\in\{0,1\},\ N=14$. Here, $l_{i}$ represents whether there is any pathology, i.e., 1 for presence and 0 for absence. We utilize the focal loss [51] as the classification loss due to class imbalance:

[TABLE]

Here, $i\in\{1,2,\ldots,N\}$. The $\alpha^{t}$ and $\gamma^{\prime}$ are the balancing and focusing hyper-parameter from the original paper [51].

Consistency Regularization Loss: The CXR without stylization $\mathbf{I}$ is fed through the backbone DenseIBN-121 to generate the global feature map $\mathbf{F}\in\mathbb{R}^{D^{\prime}\times W^{\prime}\times H^{\prime}}$. In neural style transfer, the deep high-level semantic feature maps are utilized for driving the stylization model toward content [49, 27]. In our method, $\mathbf{F}$ and $\mathbf{F_{s}}$ contain high-level semantic features for input $\mathbf{I}$ and $\mathbf{I_{s}}$, respectively. We hypothesize that if the model is biased toward semantic content, which is domain-irrelevant, then the predicted global feature activation map should be consistent if the model is fed with multiple copies of the same CXRs but with different style attributes. The only difference between $\mathbf{I}$ and $\mathbf{I_{s}}$ is the style perturbation, but the content attributes remain the same in both cases. A model that is biased towards pathological content should be robust to these variations in styles and generate a consistent global feature space. Inspired by the content loss from the literature [49, 27], we apply a content consistency regularization loss between these two global feature spaces to tweak the model’s sensitivity towards content.

[TABLE]

where, $||\cdot||_{F}$ denotes the Frobenius norm. In addition, we also utilize the Kullback-Leibler divergence loss for regularizing the predictive distributions from the classifier. Let $\mathbf{p}$ be the predictive distribution vector from the classifier for the CXR image without style perturbation ($\mathbf{I}$). We define the predictive distribution regularization loss, using (6), as:

[TABLE]

Thus, the total consistency regularization loss is defined as:

[TABLE]

The overall learning objective for the backbone model (without style nets) is comprised of the classification loss and consistency regularization loss and is expressed as:

[TABLE]

IV Experiments and Results

IV-A Thoracic Disease Datasets

We utilize the five large-scale thoracic disease datasets, namely, CheXpert [23], MIMIC-CXR [24], BRAX [25], NIH chest X-ray14 [31], and VinDr-CXR [32] datasets, for our experiments. The CheXpert, MIMIC-CXR, and BRAX datasets obtain 224,316 CXRs from 65,240 patients, 377,110 CXRs from 227,835 imaging studies for 65,379 patients, and 40,967 CXRs from 24,959 clinical studies, respectively. These three datasets have the same label space consisting of 14 pathologies. We follow the U-Ones policy [23] to prepare the ground truths. We use frontal CXR images and select one image per patient/study for our experiments. The VinDr-CXR dataset contains 18,000 images manually annotated by a total of 17 experienced radiologists. The NIH chest X-ray14 dataset comprises 112,120 chest X-rays from 30,805 unique patients with 14 diseases. The VinDr-CXR and NIH chest X-ray14 datasets have 10 and 7 common labels, respectively, with the CheXpert, MIMIC-CXR, and BRAX datasets. We have utilized these common labels for reporting the performance of these two datasets [12]. We use CheXpert and MIMIC-CXR datasets for training and the BRAX, NIH chest X-ray14, and VinDr-CXR datasets for unseen domain inference.

IV-B Training Scheme

We initialize the backbone weights with DenseIBN-121 [50] pre-trained on ImageNet. We resize the CXR images to 256×256 and normalize them with the mean and standard deviation of the ImageNet training set. We randomly crop 224×224 patches during training and use a centrally cropped sub-image of 224×224 dimensions for validation and inference. We provide the summary of training hyperparameter settings in Table II. We progressively add each proposed component to the base model and utilize the SRM modules only during training time. We employ the percentage area under the receiver operating characteristic curve (AUC) for performance evaluation [9, 12]. For comparison with other methods, we utilize the multi-label stratified k-fold scheme [53]111https://github.com/trent-b/iterative-stratification to split the training dataset into a 5-fold cross-validation framework to report the mean and standard deviation of the evaluation metrics across the folds. We use the paired t-test between cross-validated mean AUCs for reporting statistical significance.

IV-C Comparison with the State-of-the-art Methods

IV-C1 Multi-source Domain Generalization

We conduct evaluations and compare with state-of-the-art methods on the BRAX, NIH chest X-ray14, and VinDr-CXR test dataset (unseen domain) and demonstrate the domain generalization capacity of our proposed framework. We compare our proposed method with several existing state-of-the-art models, including the image-level style perturbation-based methods: AdvStyle [20]; feature-level style perturbation-based methods: FSR [22], SagNet [21], MixStyle [19], pAdaIN [29], and the method of Tang et al.[30]; and Fourier transformation based method: FDA [33]. We also compare our method with approaches developed explicitly for thoracic diseases, such as the task-specific adversarial learning approach from Luo et al.[9] and the multiple-domain mixup and ensemble learning approach from Wang et al.[12].

We have implemented the methods in the same domain generalization scenario using the same backbone model (DenseIBN-121) and common hyper-parameters, such as learning rate, image dimension, and so on. The hyper-parameters specific to the respective approaches are computed using the 5-fold cross-validation results. The results are summarized in Table III. We have computed both in-distribution (ID) and out-of-distribution (OOD) results and, as expected, found a performance discrepancy across all the models for all three datasets. However, our proposed method achieves the highest performance on the unseen domain test datasets compared to other state-of-the-art methods. As the results show, our proposed method achieves an average AUC of 77.32±0.35, 82.63±0.13, and 88.38±0.19 on the BRAX, NIH chest X-ray14, and VinDr-CXR datasets, respectively, in OOD configuration. Our proposed method achieves a 2.23%, 0.68%, and 0.92% relative average AUC improvement compared to the former best methods on these three datasets, respectively. We have also performed statistical comparisons using the paired t-test on the five-fold cross-validation scores, and the improvements are found to be significant ($p<0.05$). These results suggest that the proposed method has the ability to learn more robust, generalizable, pathology-specific, but domain-agnostic visual features compared to the existing approaches. For ID results, our proposed method achieves competitive results compared to the state-of-the-art approaches. Please note that ID results for the method of FSR [22], Luo et al.[9], and Wang et al.[12] are unavailable as they require multi-domain setup (for discriminators and creation of the virtual domain). Our method achieves 83.31±0.30, 85.38±0.05, and 95.15±0.21 on the BRAX, NIH chest X-ray14, and VinDr-CXR datasets, respectively, in ID configuration. Although FDA [33] achieves slightly better results on the BRAX dataset in the ID setting, they have a much larger drop in OOD performance compared to our method, demonstrating the strong domain generalization capability of our method.

IV-C2 Single-source Domain Generalization

A more restrictive setting than that of multi-source DG is that of single-source DG, in which only one source domain is available during training. We can seamlessly extend our framework to single-source DG, as our proposed framework does not explicitly require domain labels, unlike some of the other DG methods. To demonstrate the capability of single-source DG, we train our proposed framework on only the CheXpert or MIMIC-CXR dataset, evaluate its performance on the unseen domain test datasets, and compare it with the latest DG methods. The results are reported in Table IV. From Table IV, we can observe that our proposed method achieves superior results compared to existing state-of-the-art methods with statistically significant results.

IV-D Analysis of Different Components

We conduct experiments to evaluate the effect of each proposed component on the unseen domain test datasets, including image-level style perturbation (SRM-IL), feature-level style perturbation (SRM-FL), and consistency regularization loss ($\mathcal{L}_{cons}$). The results are reported in Table V. Checkmarks in the table indicate if the specific component is incorporated. As shown in Table V, although all three components improve the model performance, SRM-IL and consistency regularization loss contribute more strongly. This stronger contribution is expected since consistency loss reinforces the feature extractor to focus on the content region, and the SRM-IL module increases diversity.

We have conducted an ablation study to evaluate the impact of content consistency regularization ($\mathbf{\mathcal{L}_{ccr}}$) and predictive distribution regularization ($\mathbf{\mathcal{L}_{pdr}}$) losses on classification performance. The results are presented in Table VI. Combining both of the consistency regularization losses results in 77.32±0.35, 82.63±0.13, and 88.38±0.19 average AUC (%) score on the unseen domain test datasets, i.e., BRAX, NIH Chest X-ray14, and VinDr-CXR datasets. These scores surpass the performance level achieved by using either of the regularization loss alone. We hypothesize that while predictive distribution regularization loss enforces consistency on two probability distributions from two different style augmented versions of the same CXR image, it does not guarantee that the backbone model will focus on the same spatial region for predicting the consistent probabilities. In order to address this issue, we employ the content consistency loss in addition to the predictive distribution loss. We can also observe that the content consistency loss results in better performance compared to the predictive distribution loss when they are utilized alone.

IV-E Analysis of SRM-FL

Effect of randomization stage: Different layers of the CNN model encode distinct visual attributes, such as lower layers tend to encode low-level texture information while higher levels encode high-level semantic features [49, 54]. Thus, we must apply the SRM-FL at an adequate stage. This is because employing a SRM-FL to randomize too high-level feature statistics may alter the semantic content necessary for pathology prediction, while randomizing too low-level features may lead to sub-optimal performance. To evaluate the effect of the randomization stage, we apply the SRM-FL after dense block-1,2, and 3, one at a time, and measure the performances. The results are shown in part-1 of Table VII. Our experiment shows that the proposed framework performs best when the SRM-FL is applied after dense block-2.

Impact of learnable style embeddings: To evaluate the impact of the learnable pixel-wise affine parameters as the style embeddings compared to the pre-defined style embeddings, we replace the style nets in the SRM-FL with the AdaIN module [27] and compute the percentage AUC score. The results are reported in part-2 of Table VII. We can observe that the learnable pixel-wise affine parameters achieve better results compared to pre-defined style embedding parameters.

Analysis of the style net architecture: The style nets $\phi_{\gamma}$ and $\phi_{\beta}$ contain only convolutional (Conv) layers. We conduct experiments adding a rectified linear layer (ReLU) and a batch normalization (BN) layer after the convolutional layers, respectively. The results are reported in part-3 of Table VII. We can observe that the style nets with only convolutional layers achieve the best score. We speculate that the style parameters are set as the affine parameters, and as the affine parameters do not necessarily lie on the positive plane, the utilization of ReLU results in lower performance because it limits the learnable parameter range. In the case of BN, we speculate that the population statistics may interfere with instance statistics, as each single pair (content and style reference sample) has different styles.

Sensitivity of the hyperparameter $\boldsymbol{\eta}$: We vary the hyper-parameter $\eta$ from 1 to 1e-5 and calculate the percentage AUC performance. The results are demonstrated in Fig. 3. From Fig. 3, we can observe that increasing $\eta$ too much degrades the score as it disrupts the semantic content features. Decreasing the $\eta$ too much reduces the style diversity of the SRM-FL, resulting in performance loss. The hyper-parameter $\eta$ controls the stylization capacity of the style nets by governing the content and style loss. CXRs often have complex backgrounds, and too-hard style perturbations may increase the complexity and potential interference with diagnosing small lesions. To verify this, we plot AUC (%) scores of two different lesions, i.e., cardiomegaly and lung lesions, compared to the $\eta$ values, on the unseen domain BRAX test dataset. Note that the small lesions, i.e., masses and nodules, are included in the lung lesion category. The plots are given in Fig. 4. We can observe that for large lesions such as cardiomegaly, the AUC (%) does not vary much concerning the $\eta$ value. However, the $\eta$ value significantly impacts the detection of the lung lesion category, i.e., small lesions. Therefore, it is recommended to calculate the optimal $\eta$ value for performing the feature-level style perturbation. Another strategy that can be adopted is to increase the spatial dimension of the input image, which is often utilized in literature for thoracic disease detection to improve the detection of small lesions [1]. To demonstrate this, we have trained our model on 384×384 input images and compared the performance with the 224×224 input images under $\eta$=0.01, on the unseen domain BRAX test dataset. The results are reported in Table VIII. We can observe that the AUC (%) of the lung lesion class improves from 60.44 to 65.32.

V Discussion & Conclusion

In this paper, we propose a DG framework that extracts domain-invariant features by making a deep learning model style invariant and biased toward content. To the best of our knowledge, no previous DG study for thoracic disease detection has exploited neural style transfer to develop a domain-agnostic model. We deploy SRMs both at the image and feature level to augment the style characteristics and employ consistency regularizations to tweak the backbone model’s sensitivity towards semantic content rather than uninformative texture, i.e., styles, for pathology prediction. Extensive experiments with five large-scale thoracic disease datasets demonstrate significantly improved performance and generalizability compared to the baseline and state-of-the-art methods on the unseen domain test dataset.

One limitation of our proposed method is that we have to feed both CXR images with or without style attributes perturbed into the backbone to exploit the consistency regularization losses, which increases the training runtime and computing costs. Although our approach has high runtime and computing costs during the offline training stage, it is very efficient for online testing and would be suitable for the routine clinical workflow because the proposed mechanisms are not utilized during the inference stage. Another limitation is that for the SRM-FL, the hyperparameter $\eta$ value needs to be determined through an ablation study, while the SRM-IL and consistency regularization mechanisms are parameterless.

Future work directions can involve (1) analyzing how the proposed DG framework can complement the methods designed explicitly for improving thoracic disease detection, e.g., anatomy-aware-based methods, (2) integrating shape-based radiomic features for content regularization, and (3) utilizing local patch statistics instead of global statistics.