Lung cancer risk prediction using augmented machine learning pipelines with explainable AI

Abstract

Lung cancer remains the leading cause of cancer-related deaths worldwide, making early and precise diagnosis is critical for improving the patient survival rates. Machine learning has shown promising results in predictive analysis for lung cancer prediction. However, class imbalance in clinical datasets negatively impacts the performance of Machine Learning classifiers, leading to biased predictions and reduced accuracy. In an attempt to address this issue, various data augmentation techniques were applied alongside classification models to enhance predictive performance. This study evaluates data augmentation techniques paired with machine learning classifiers to address class imbalance in a small lung cancer dataset. A comparative analysis was conducted to assess the impact of different augmentation techniques with classification models. Experimental findings demonstrate that K-Means SMOTE, combined with a Multi-Layer Perceptron classifier, achieves the highest accuracy of 93.55% and an AUC-ROC score of 96.76%, surpassing other augmentation-classifier combinations. These results underscore the importance of selecting optimal augmentation methods to improve classification performance. Furthermore, to ensure model interpretability and transparency in medical decision-making, LIME is utilized to provide insights into model predictions. The study highlights the significance of advanced augmentation techniques in addressing data imbalance, ultimately enhancing lung cancer risk prediction through machine learning. The findings contribute to the growing field of AI-driven healthcare by emphasizing the necessity of selecting effective augmentation-classifier pairs to develop more accurate and reliable diagnostic models. Due to the dataset's high cancer prevalence (87.45%) and limited size, this work is a preliminary methodological comparison, not a clinical tool. Findings emphasize the importance of augmentation for imbalanced data and lay the groundwork for future validation with larger, representative datasets.

Keywords: SMOTE; class imbalance; explainable AI; lime; lung cancer prediction.

Figure 1

Proposed architecture of the methodology.

Figure 2

Combination of augmentation-classification methods used for lung cancer risk prediction.

Figure 3

Age and gender distribution of the patients. (a) Donut chart showing gender distribution, with 52.4% female and 47.6% male. (b) Histogram displaying age distribution of patients.

Figure 4

Feature correlation heatmap showing the relationships between different features.

Figure 5

Stacked bar chart of each feature. (a) Gender, (b) Age, (c) Smoking, ((d) Yellow Fingers, (e) Anxiety, (f) Peer Pressure, (g) Chronic Disease, (h) Fatigue, (i) Allergy, (j) Wheezing, (k) Alcohol Consuming, (l) Coughing, (m) Shortness of Breath, (n) Swallowing Difficulty, (o) Chest Pain, with the target variable lung cancer.

Figure 6

A comparative analysis of each method using AUC-ROC curve for lung cancer risk prediction. (a) SMOTE + Logistic Regression, (b) SMOTE + KNN, (c) SMOTE + XGBoost, (d) ADASYN + Decision Tree, (e) ADAYSN + Random Forest, (f) ADASYN + LightGBM, (g) SVMSMOTE + XGBoost, (h) SVMSMOTE + AdaBoost, (i) Borderline SMOTE + CatBoost, (j) SMOTENC + Logistic Regression, (k) K-Means SMOTE + Multi-Layer Perceptron, (l) SMOTE-ENN + Gradient Boosting, (m) Random Oversampling + Random Forest, (n) Random Oversampling + SVM, (o) Random Undersampling + SVM, (p) Random Undersampling + LightGBM.

Figure 7

A comparative analysis of each method using confusion matrices for lung cancer risk prediction. (a) SMOTE + Logistic Regression, (b) SMOTE + KNN, (c) SMOTE + XGBoost, (d) ADASYN + Decision Tree, (e) ADAYSN + Random Forest, (f) ADASYN + LightGBM, (g) SVMSMOTE + XGBoost, (h) SVMSMOTE + AdaBoost, (i) Borderline SMOTE + CatBoost, (j) SMOTENC + Logistic Regression, (k) K-Means SMOTE + Multi-Layer Perceptron, (l) SMOTE-ENN + Gradient Boosting, (m) Random Oversampling + Random Forest, (n) Random Oversampling + SVM, (o) Random Undersampling + SVM, (p) Random Undersampling + LightGBM.

Figure 8

A comparative analysis for a specific instance of each method using LIME for lung cancer risk prediction. (a) SMOTE + Logistic Regression, (b) SMOTE + KNN, (c) SMOTE + XGBoost, (d) ADASYN + Decision Tree, (e) ADAYSN + Random Forest, (f) ADASYN + LightGBM, (g) SVMSMOTE + XGBoost, (h) SVMSMOTE + AdaBoost, (i) Borderline SMOTE + CatBoost, (j) SMOTENC + Logistic Regression, (k) K-Means SMOTE + Multi-Layer Perceptron, (l) SMOTE-ENN + Gradient Boosting, (m) Random Oversampling + Random Forest, (n) Random Oversampling + SVM, (o) Random Undersampling + SVM, (p) Random Undersampling + LightGBM.

Figure 9

LIME instances of index (a) 15 and (b) 50 for K-Means SMOTE + Multi-Layer Perceptron combination.