FairCXRnet: A Multi-Task Learning Model for Domain Adaptation in Chest X-Ray Classification for Low Resource Settings ^†

Abstract

Medical imaging analysis plays a pivotal role in modern healthcare, with physicians relying heavily on radiologists for disease diagnosis. However, many hospitals face a shortage of radiologists, leading to long queues at radiology centers and delays in diagnosis. Advances in artificial intelligence (AI) have made it possible for AI models to analyze medical images and provide insights similar to those of radiologists. Despite their successes, these models face significant challenges that hinder widespread adoption. One major issue is the inability of AI models to generalize data from new populations, as performance tends to degrade when evaluated on datasets with different or shifted distributions, a problem known as domain shift. Additionally, the large size of these models requires substantial computational resources for training and deployment. In this study, we address these challenges by investigating domain shifts using ChestXray-14 and a Nigerian chest X-ray dataset. We propose a multi-task learning (MTL) approach that jointly trains the model on both datasets for two tasks, classification and segmentation, to minimize the domain gap. Furthermore, we replace traditional convolutional layers in the backbone model (Densenet-201) architecture with depthwise separable convolutions, reducing the model’s number of parameters and computational requirements. Our proposed model demonstrated remarkable improvements in both accuracy and AUC, achieving 93% accuracy and 96% AUC when tested across both datasets, significantly outperforming traditional transfer learning methods.

1. Introduction

The use of medical imaging in contemporary healthcare systems is essential, enabling physicians to diagnose and monitor diseases with high precision. Among the various imaging modalities, chest X-rays (CXRs) are one of the most widely used tools for detecting conditions such as pneumonia, tuberculosis, and COVID-19 [1]. However, the reliance on radiologists for interpretation poses difficulties, particularly in areas where qualified medical workers are lacking [2]. This shortage often leads to delayed diagnosis, negatively impacting patient outcomes. In light of this, deep learning-based artificial intelligence (AI) methods have emerged as promising solutions to aid in automatic medical image analysis, potentially alleviating the burden on radiologists [3].

Despite the notable progress in AI-based CXR analysis, two key challenges hinder adopting these models in real-world clinical settings: domain shift and high computational complexity [4]. Domain shift refers to the significant performance degradation that occurs when a model trained in one data set is evaluated in a data set from a different population or hospital with varying imaging protocols, equipment, or patient demographics [5]. This phenomenon is particularly concerning in medical imaging analysis, where models must generalize across diverse populations and healthcare systems. For instance, a model trained on the widely used ChestXray-14 dataset may not perform well on data collected in Nigerian hospitals, where differences in image acquisition and population characteristics create a substantial domain gap. Addressing this domain shift is crucial for building robust models that can be deployed globally [6].

Additionally, state-of-the-art deep learning models often require vast computational resources for training and inference, limiting their accessibility, particularly in low-resource settings [7]. The large number of parameters in these models makes them both memory-intensive and computationally expensive, further complicating their practical deployment in resource-constrained environments [8].

To address these challenges, we propose the MTL approach that performs both classification and segmentation tasks to improve the generalizability of deep learning models across different populations. Specifically, we aim to minimize the domain gap between the ChestXray-14 dataset and a Nigerian chest X-ray dataset. By training a model on multiple related tasks (segmentation and classification) simultaneously, MTL encourages the model to learn shared representations that are more robust to domain shifts. Additionally, we replace traditional convolutional layers with depthwise separable convolutions to reduce the number of trainable parameters and computational requirements, building a model that is more efficient for deployment in low-resource environments.

Traditionally, models are trained sequentially on single tasks, such as segmentation, followed by parameter tuning for a different task, like classification, until desirable performance is achieved. This approach allows the model to handle multiple tasks but lacks efficiency and cohesion. In contrast, MTL leverages domain-specific features from related tasks to improve generalization, enabling the model to perform better across several tasks simultaneously. By allowing one task to reinforce another, MTL enhances overall model performance with a single, unified architecture.

Therefore, in this work, we propose a novel method for joint segmentation and classification that aligns domain differences between two datasets, improving both classification accuracy and segmentation quality. Our main technical contribution is the development of a lightweight multi-task learning model, capable of utilizing small, labeled medical image datasets for efficient classification and segmentation.

The remainder of this paper is structured as follows: Section 2 reviews related works addressing domain shifts in chest X-ray analysis and medical imaging. Section 3 outlines the methodology employed in this study, detailing the multi-task learning approach and architectural improvements. Section 4 presents the experimental results, followed by an in-depth discussion of the findings. Finally, Section 5 concludes the paper with a summary of key insights and future research directions.

2. Related Works

In recent years, the rapid advancements in AI and deep learning have enabled significant progress in medical imaging analysis, particularly in tasks such as classification, segmentation, and disease detection from chest X-rays (CXRs) [9]. However, despite the success of these models, challenges related to domain shift and computational complexity remain critical obstacles to their widespread adoption in real-world clinical settings [10]. This section provides an overview of relevant works addressing these issues, with a particular focus on methods for handling domain shifts in chest X-ray classification as well as strategies for improving computational efficiency.

2.1. Domain Shift in Medical Imaging

Domain shift refers to the degradation in model performance when it is trained on one dataset and evaluated on a different dataset with varying characteristics, such as differences in patient demographics, imaging devices, and clinical protocols [11]. This issue is especially pronounced in medical imaging, where diverse data distributions can arise due to variations in hospital equipment, image acquisition techniques, and population health conditions [12]. Several studies have explored solutions to mitigate the effects of domain shift in chest X-ray analysis.

One popular approach to address domain shift is through domain adaptation techniques [13]. For unsupervised settings where the dataset is not labeled, Unsupervised Domain Adaptation (UDA) techniques were used to mitigate the effect of domain shift. The technique performed well in a series of studies, and it performed well on the famous office-31 dataset for domain adaptation and achieved an adaptive effect in the event of large domain shifts [14]. Motiian explored supervised domain adaption techniques where both data from source and target domains are labeled, by learning a discriminative embedding subspace with mapped visual domains that are maximally segregated but semantically linked using Siamese architecture [15]. Tzeng proposed a Domain-Adversarial Neural Network (DANN) as a technique to handle domain shift, where an adversarial objective is used to align the feature distributions of the source and target domains, thereby improving model generalization across different datasets [16]. Similarly, Adversarial Discriminative Domain Adaptation (ADDA) is proposed to minimize domain discrepancies by optimizing both classification and adversarial loss functions. These approaches have shown promise, but their performance can still degrade when there is significant variability between domains [17].

Another method to tackle domain shift is transfer learning, which has been widely used in medical imaging tasks [18]. Models pre-trained on large datasets, such as ImageNet [19] or ChestXray-14 [20], are fine-tuned on a smaller, target domain dataset to adapt the learned features to the new domain. Transfer learning approaches for CXR classification were used and achieved considerable improvements, but these methods still rely heavily on large labeled datasets and do not fully eliminate the domain shift when applied to datasets with distinct characteristics, such as those from low-resource healthcare settings [21].

Despite these advances, most existing domain adaptation methods focus on aligning distributions between datasets but fail to consider the underlying task complexity and the potential benefits of incorporating additional tasks. This motivates the use of MTL as a more holistic approach to mitigate domain shift.

2.2. Multi-Task Learning in Medical Imaging

MTL has emerged as a powerful framework in deep learning, where multiple related tasks are learned jointly [22]. By sharing representations between tasks, MTL encourages the model to learn more robust and generalized features, making it less susceptible to domain-specific variations [23]. Caruana first introduced the concept of MTL, demonstrating that learning-related tasks together improve generalization and reduce overfitting [24].

In the context of medical imaging, several studies have successfully applied MTL to tasks such as disease classification and segmentation. Gai et al. [25] proposed an MTL-based architecture for simultaneous segmentation and classification of lung nodules in CT scans, showing that the model learned complementary features from both tasks, resulting in improved accuracy and generalization. Others explored the use of MTL for COVID-19 detection from chest X-rays, combining the classification of pneumonia, COVID-19, and normal cases with a segmentation task to identify affected regions in the lungs. Their results showed significant improvements over single-task learning approaches [26].

However, while MTL has demonstrated its effectiveness in medical imaging, no studies have applied this approach to address domain shifts across populations, to the best of our knowledge. By training on multiple tasks, such as classification and segmentation, MTL can potentially improve feature extraction across different domains, providing a more robust solution to the cross-population domain shift problem. This research builds on the successes of MTL by investigating its role in mitigating domain shifts between two distinct datasets: the publicly available ChestXray-14 dataset and the Nigerian chest X-ray dataset.

2.3. Lightweight and Efficient Architectures

Another key challenge in deploying deep learning models in medical imaging is the computational cost associated with training and inference. State-of-the-art deep learning models often contain millions of parameters, making them computationally expensive to train and requiring significant memory for deployment [27]. This is a critical issue, especially in low-resource settings, where access to high-performance computing infrastructure is limited [28].

Several architectural innovations have been proposed to address the computational complexity of deep learning models. Howard et al. [29] introduced MobileNet, a lightweight architecture that utilizes depthwise separable convolutions to significantly reduce the number of parameters without sacrificing performance. The concept was extended with Xception, which demonstrated the effectiveness of depthwise separable convolutions in large-scale image classification tasks [30]. These models have proven particularly useful in resource-constrained environments, as they achieve state-of-the-art performance with reduced computational requirements.

In the medical imaging domain, lightweight architectures such as MobileNet have been explored for tasks like disease detection and segmentation. Rajpurkar [20] applied MobileNet for pneumonia detection from chest X-rays and achieved competitive performance compared to large models. Additionally, the use of depthwise separable convolutions in skin cancer detection was seen in the literature, showing that such architectures could deliver strong performance while being efficient in terms of computational and memory requirements [31].

Therefore, we leverage the power of depthwise separable convolutions in our proposed multi-task learning architecture to reduce the computational burden while maintaining high accuracy and generalization across domains. By combining MTL with a lightweight architecture, we aim to create a model that is not only robust to domain shift but also practical for deployment in resource-constrained settings.

3. Methodology

This section details the methodology used to address the problem of cross-population domain shift problem in chest X-ray analysis with the MTL approach. Our goal is to develop a model capable of jointly performing classification and segmentation tasks across two datasets, namely ChestXray-14 and the Nigerian chest X-ray dataset, while also optimizing computational efficiency. To achieve this, we designed a lightweight MTL architecture using depthwise separable convolutions and Densenet-201 as the backbone model, aiming to reduce the domain gap and generalize effectively across different populations.

3.1. Dataset Description

We utilized two datasets in this study. ChestXray-14 Dataset: A publicly available dataset with 112,120 chest X-ray images labeled for 14 thoracic diseases, collected in the United States (US), and organized by the National Institutes of Health (NIH). This dataset was used as the source domain. Nigerian Chest X-ray Dataset (Nigeria CXR): Collected from Aminu Kano Teaching Hospital in Nigeria, this dataset consists of 6345 chest X-ray images labeled into four classes: Pneumonia, COVID-19, Tuberculosis (TB), and Normal. This dataset was used as the target domain.

Figure 1 shows sample images from the dataset. Both datasets were annotated by radiologists and pre-processed to ensure consistency in image size and quality. We applied standard image augmentation techniques such as rotation, flipping, and contrast adjustment to increase the diversity of the training data and improve the model’s robustness to variations in image quality and acquisition settings.

3.2. Proposed Multi-Task Learning Model

The proposed model is based on MTL architecture that performs both classification and segmentation tasks. By learning these two related tasks simultaneously, the model can share representations between tasks, improving generalization and reducing domain shift. The proposed model is well depicted in Figure 2. It shows that the model architecture consists of two main branches:

To enhance robustness against domain variations, an attention-guided fusion mechanism was utilized to exploits segmentation features that refine the classification pathway in shared feature space. The decoder output $f_{\mathrm{seg}}$ is transformed into an attention mask through a $1\times 1$ convolution and sigmoid activation, i.e.,

$$A=\sigma \left({\mathrm{Conv}}_{1\times 1}\left(f_{\mathrm{seg}}\right)\right),f_{\mathrm{att}}=A\odot f_s,$$

(1)

where $f_s$ denotes the shared feature representation and ⊙ the elementwise product. The gated representation $f_{\mathrm{att}}$ emphasizes anatomically relevant regions and is concatenated with the original features to form $z=[f_s\parallel f_{\mathrm{att}}]$, which is subsequently fed into the classifier. In addition, to mitigate domain shift, we incorporate a Maximum Mean Discrepancy (MMD) regularizer that minimizes the distributional gap between source and target features in the reproducing kernel Hilbert space:

$${\mathrm{MMD}}^2(f_s^{\left(S\right)},f_s^{\left(T\right)})={\displaystyle \frac{1}{n_S^2}}\sum _{i,j}k(x_i^S,x_j^S)+{\displaystyle \frac{1}{n_T^2}}\sum _{i,j}k(x_i^T,x_j^T)-{\displaystyle \frac{2}{n_S{n}_T}}\sum _{i,j}k(x_i^S,x_j^T),$$

(2)

with $k(·,·)$ defined as a mixture of RBF kernels. The overall training objective is then expressed as

$$\mathcal{L}={\mathcal{L}}_{\mathrm{cls}}+{\lambda }_{\mathrm{seg}}{\mathcal{L}}_{\mathrm{seg}}+{\lambda }_{\mathrm{mmd}}{\mathcal{L}}_{\mathrm{MMD}},$$

(3)

ensuring joint optimization of classification accuracy, segmentation quality, and domain-invariant feature alignment.

1-Classification head: This branch is responsible for disease classification, predicting one of the four classes (Pneumonia, COVID-19, TB, Normal) in both datasets.

2-Segmentation head: The segmentation branch is designed to identify and localize regions of interest in the chest X-rays, such as areas affected by pneumonia or TB. This branch outputs a pixel-wise mask that highlights the regions of abnormality.

The two tasks share a common encoder, which extracts features from the input images. The shared encoder allows the model to learn useful representations for both tasks, encouraging it to capture more general features that are less sensitive to domain-specific variations. The training procedure is best desribed in Algorithm 1, where the proposed MTL model integrates segmentation-guided attention with distribution alignment to improve classification under cross-population settings.

Algorithm 1 Attention-Guided Fusion + RBF–MMD

Require: $f_s,f_{\mathrm{seg}},f_s^{\left(S\right)},f_s^{\left(T\right)},y_{\mathrm{cls}},y_{\mathrm{mask}},{\lambda }_{\mathrm{seg}},{\lambda }_{\mathrm{mmd}},\Sigma$   1: $A\leftarrow \sigma \left({\mathrm{Conv}}_{1\times 1}\left(f_{\mathrm{seg}}\right)\right)$   2: $f_{\mathrm{att}}\leftarrow A\odot f_s$   3: $z\leftarrow \mathrm{Concat}(f_s,f_{\mathrm{att}})$   4: $\mathrm{logits}\leftarrow {\mathrm{Head}}_{\mathrm{cls}}\left(z\right)$   5: ${\mathcal{L}}_{\mathrm{cls}}\leftarrow \mathrm{CE}/\mathrm{Focal}(\mathrm{logits},y_{\mathrm{cls}})$   6: ${\mathcal{L}}_{\mathrm{seg}}\leftarrow \mathrm{Dice}+\mathrm{CE}(\mathrm{SegPred},y_{\mathrm{mask}})$   7: $X\leftarrow \mathrm{Flatten}\left(f_s^{\left(S\right)}\right)$,   $Y\leftarrow \mathrm{Flatten}\left(f_s^{\left(T\right)}\right)$   8: ${\mathcal{L}}_{\mathrm{MMD}}\leftarrow \mathtt{MMD}\_\mathtt{RBF}(X,Y,\Sigma )$   9: $\mathcal{L}\leftarrow {\mathcal{L}}_{\mathrm{cls}}+{\lambda }_{\mathrm{seg}}{\mathcal{L}}_{\mathrm{seg}}+{\lambda }_{\mathrm{mmd}}{\mathcal{L}}_{\mathrm{MMD}}$ 10: return $\mathcal{L}$

3.3. Depthwise Separable Convolutions

To address the computational challenges associated with training deep learning models on large medical datasets, we integrated depthwise separable convolutions into the MTL architecture. Depthwise separable convolutions break down standard convolutions into two smaller operations:

• Depthwise Convolution: This operation independently applies a single filter to each input channel, reducing the computational cost.
• Pointwise Convolution: A 1 × 1 convolution is then applied to combine the outputs of the depthwise convolution.

By using depthwise separable convolutions, the model parameters are significantly reduced while maintaining its capacity to learn rich feature representations. This design choice made the model lightweight and computationally efficient, making it suitable for deployment in low-resource settings.

3.4. Model Training Strategy

The proposed MTL model training is a twofold training process: 1. Joint Training: The model was trained on both tasks (classification and segmentation) simultaneously using a multi-task loss function. The overall loss function, L_total, is a weighted sum of the classification loss and the segmentation loss:

$$L_{\mathrm{total}}=\alpha ·L_{\mathrm{classification}}+\beta ·L_{\mathrm{segmentation}}$$

where $\alpha$ and $\beta$ are hyperparameters that control the contribution of each task to the total loss.

2. Domain Alignment: To reduce the domain gap between the ChestXray-14 and the Nigeria CXR datasets, we applied feature sharing during model optimization between features extracted from the source and target domains. The model’s encoder was trained to learn domain-invariant features in the shared feature space. This helped align the feature distributions of the two datasets, improving the model’s generalization to the target domain.

All experimental results are evaluated on the standard performance metrics, such as Accuracy, Area Under the Receiver Operating Characteristic Curve (AUC), Precision, Recall, and F1-score.

The model’s performance was compared against several baseline methods, including transfer learning models and single-task learning approaches. We conducted extensive experiments to assess the model’s robustness to domain shift, particularly focusing on its ability to generalize to the Nigerian chest X-ray dataset.

3.5. Implementation Details

The model was implemented in Pytorch 2.7.1 and trained on an A100 GPU. An Adam optimizer with an initial learning rate of 0.0001 and a batch size 32 was used. Early stopping was employed to prevent overfitting, and a learning rate scheduler was utilized to reduce the learning rate when the validation loss converged gradually.

4. Results

The Results Section provides a detailed account of the performance of the proposed multi-task learning model on both classification and segmentation tasks. The evaluation focuses on key metrics such as accuracy, AUC, and F1-score, and comparisons are made against baseline models to demonstrate the effectiveness of our approach.

4.1. Experiment 1: Establishing Baselines

In this experiment, we used Densenet-201 [32] as the backbone model. The model is trained and evaluated on both ChestXray-14 and the Nigerian dataset. The goal was to establish a baseline for training using a conventional architecture.

Table 1. Experiment I results.

Datasets	Accuracy	Precision	Recall	F1-Score	AUC
ChestXray14	96.10	86.73	83.33	96.00	0.95
Nigerian-CXR	91.33	91.69	91.33	91.32	0.98

presents the result of experiment 1, which shows the traditional CNN model’s performance on both datasets. The experiment demonstrated DenseNet-201’s baseline performance when trained on the two datasets, offering a reference point for comparing with our proposed model.

4.2. Investigation of Domain Shift Between ChestXray-14 and Nigerian Dataset

In the second experiment, we trained DenseNet-201 solely on the ChestXray-14 dataset and tested it on the Nigeria CXR dataset to investigate the impact of domain shift. The goal was to observe the model’s performance when evaluated on a new population dataset without domain adaptation techniques.

Experiment 2 revealed significant performance degradation when DenseNet-201 was trained on ChestXray-14 and tested on the Nigerian dataset.

Table 2. Experiment II results.

Datasets	Accuracy	Precision	Recall	F1-Score	AUC
ChestXray14 –>Nigerian CXR	61.58	51.16	50.03	59.87	0.47

summarizes our quantitative findings, including the Accuracy and average AUC values for experiment 2. The findings following the model’s evaluation on Nigeria CXR show reduced performance on all the performance metrics, indicating that the model did not generalize well to new datasets. This suggests that the training data from the source domain does not sufficiently represent the distribution of data from the target domain.

4.3. Proposed FairCXRnet Results

In the experiment 3, we trained the proposed MTL model on the ChestXray-14 and Nigeria CXR datasets for joint classification and segmentation. The model was tested on the Nigerian dataset to evaluate its ability to generalize and reduce domain shift.

Figure 3 shows the performance in terms of accuracy of the proposed model compared with the transfer learning/fine tuning and continual learning technique, and

Table 3. Proposed FairCXRnet results.

Techniques	Accuracy	Precision	Recall	F1-Score	AUC
Transfer learning and fine-tuning [33]	88.36	82.88	79.16	88.47	0.87
Continual learning [34]	72.41	74.66	72.02	71.97	0.71
Proposed FairCXRnet	93.26	89.88	91.76	93.07	0.96

summarizes the findings.

From Table 3, it could be inferred that the proposed FairCXRnet model outperformed the transfer learning approach using Densenet-201 across all metrics, demonstrating its effectiveness in reducing domain shift and improving overall performance across tasks. Continual learning (CL) also shows promising performance, striking a good balance between precision and recall. However, it suffers from catastrophic forgetting: as it continues learning new tasks, it tends to forget the old task. Additionally, while Densenet has over 20 M parameters, and a 77 MB size on disk, the integration of depthwise separable convolutions in our model reduces the total parameters to just 197,848, and model size to 31 MB, achieving 59% parameter reduction with more generalization capability, making it suitable for low-resource settings.

Discussion

The experimental results demonstrate the effectiveness of the proposed FairCXRnet model in addressing the challenges of cross-population domain shift and computational efficiency in chest X-ray analysis. In the first experiment, DenseNet201, trained on both the ChestXray-14 and Nigerian datasets, exhibited reasonable performance. However, when the model was trained solely on ChestXray-14 and tested on the Nigerian dataset (experiment 2), we observed a significant drop in accuracy and AUC, confirming the presence of domain shift. This performance degradation stresses the difficulty in applying models trained on one population to data from a different region or demographic without adaptation.

In contrast, the third experiment, where the proposed FairCXRnet model was trained on both datasets, demonstrated notable improvements. The FairCXRnet model achieved an accuracy of 93% and an AUC of 96% on the Nigerian dataset, outperforming the DenseNet-201 baseline by a substantial margin, further validating the benefit of jointly learning-related tasks. The use of depthwise separable convolutions not only reduced the model’s parameter count but also reduced training time significantly from 8 h to 2 h, improving computational efficiency and making the model more suitable for deployment in resource-constrained settings. These results suggest that multi-task learning, combined with domain alignment strategies, is a promising approach for improving generalization in medical imaging, particularly in diverse and underrepresented populations.

5. Conclusions and Future Work

In this paper, we addressed the challenges of domain shift and computational efficiency in chest X-ray analysis by proposing an MTL model that jointly performs classification and segmentation. By leveraging the ChestXray-14 and Nigerian chest X-ray datasets, we demonstrated that our model successfully reduces the domain gap and improves performance on unseen data from a different population. The experimental results show that our MTL approach outperformed traditional single-task models, achieving higher classification accuracy and better segmentation quality while significantly reducing the number of model parameters through the use of depthwise separable convolutions. This makes the model not only effective in mitigating domain shift but also computationally efficient, which is crucial for deployment in resource-limited settings such as hospitals with minimal computing infrastructure.

Our findings suggest that multi-task learning, combined with domain alignment strategies, holds promise for improving generalization in medical imaging applications across diverse populations. The results also highlight the importance of considering domain-specific variations when developing AI models for healthcare, ensuring that they are adaptable and robust across different regions and patient demographics.

While our proposed model achieved significant improvements, there are several avenues for future research. First, we plan to explore more advanced domain adaptation techniques, such as self-supervised learning, to further improve the model’s generalization to new datasets. Additionally, incorporating larger and more diverse datasets could help to better understand the effects of demographic variations on model performance.