Optimizing Machine Learning for Healthcare Applications: A Case Study on Cardiovascular Disease Prediction Through Feature Selection, Regularization, and Overfitting Reduction ^†

Abstract

The application of machine learning (ML) to medical datasets offers significant potential for improving disease prediction and patient outcomes. However, challenges such as feature redundancy, overfitting, and suboptimal model performance limit the practical effectiveness of ML algorithms. This study focuses on optimizing ML techniques for cardiovascular disease prediction using the Kaggle Cardiovascular Disease dataset. We systematically apply feature selection methods, including correlation analysis and regularization techniques (L1/L2), to identify the most relevant attributes and address multicollinearity. Advanced ensemble models such as Random Forest, XGBoost, and LightGBM are employed to mitigate overfitting and enhance predictive performance. Through hyperparameter tuning and stratified k-fold cross-validation, we ensure model robustness and generalizability. The results demonstrate that ensemble methods, particularly gradient boosting algorithms, outperform traditional models, achieving superior predictive accuracy and stability. This study highlights the importance of algorithm optimization in ML applications for healthcare, offering a replicable framework for medical datasets and paving the way for more effective diagnostic tools in cardiovascular health.

1. Introduction

Cardiovascular diseases (CVDs) remain the leading cause of mortality worldwide, accounting for nearly 20.5 million deaths annually [1]. Early and accurate prediction of CVDs can significantly enhance patient outcomes by enabling timely interventions and personalized treatment plans [2]. In recent years, machine learning (ML) has emerged as a powerful tool in healthcare, offering the ability to process complex datasets and uncover patterns that traditional statistical methods may overlook. By leveraging ML, healthcare professionals can develop predictive models to assess a patient’s risk of developing cardiovascular conditions based on diverse clinical and demographic data [3].

The Kaggle Cardiovascular Disease dataset shared by Svetlana Ulianova, a data scientist, provides a rich source of structured medical data, including features such as age, gender, cholesterol levels, and smoking status. This dataset is widely used in research to develop and evaluate predictive ML models for cardiovascular risk assessment. Its comprehensive nature makes it an excellent candidate for exploring advanced ML techniques and their impact on prediction accuracy [4].

Despite significant progress in applying ML to medical datasets, prior studies often focus on comparing algorithmic performance without fully optimizing their implementations. Key challenges such as feature redundancy, multicollinearity, and overfitting are frequently overlooked, leading to suboptimal and less generalizable models [5]. Additionally, many studies do not incorporate systematic feature selection, regularization techniques, or robust validation frameworks, which are crucial for improving model interpretability and reliability in clinical settings.

This study aims to address these gaps by systematically optimizing ML algorithms for cardiovascular disease prediction. Specifically, we achieve the following:

• Perform feature selection using correlation analysis, Recursive Feature Elimination (RFE), and mutual information classifiers.
• Apply L1/L2 regularization to address multicollinearity and enhance model interpretability.
• Mitigate overfitting through the use of Random Forest and Gradient Boosting models.
• Implement cross-validation and data augmentation to ensure robust and generalizable predictions.

Through these methods, we seek to develop a replicable framework for optimizing ML models on medical datasets, ultimately contributing to the advancement of ML applications in healthcare.

2. Related Work

The application of machine learning (ML) in healthcare has grown rapidly, driven by the availability of large medical datasets and the need for precise diagnostic tools. Studies have explored ML techniques for predicting cardiovascular diseases (CVDs), demonstrating their potential in enhancing clinical decision making [6]. Traditional algorithms like Logistic Regression, SVM, and Decision Trees have been used for CVD prediction, but they often struggle with complex, high-dimensional datasets [6]. In contrast, ensemble methods like Random Forest and Gradient Boosting have shown improved accuracy by leveraging multiple weak learners, making them effective for noisy and imbalanced medical data [7]. Neural networks and deep learning models, while capable of capturing non-linear relationships, are limited by the need for large labeled datasets [8].

Feature selection is critical in healthcare ML applications. Techniques like Recursive Feature Elimination (RFE) and correlation analysis have been shown to improve model performance and interpretability, though their application is not always systematic [9]. Overfitting remains a challenge, particularly with limited datasets, and regularization techniques like L1/L2 penalties are underutilized [10]. Robust validation frameworks, such as k-fold cross-validation, are also frequently overlooked, leading to inflated performance metrics [10]. Additionally, data augmentation for tabular medical data, such as generating synthetic samples, remains largely unexplored [11].

In previous work [12], ML algorithms were applied to a cardiovascular dataset in a smart home environment for medical surveillance. The study evaluated algorithms like Linear Regression, Logistic Regression, and Multi-Layer Perceptron (MLP), with MLP achieving the highest accuracy (99.80%). However, the study lacked advanced optimization techniques like feature selection and regularization, potentially leading to overfitting and suboptimal generalization.

This study builds on these findings by integrating advanced optimization strategies, including feature selection, regularization, and robust validation, to address prior limitations. Using the Kaggle Cardiovascular Disease dataset, we aim to develop a replicable methodology that improves predictive accuracy and model reliability for CVD prediction.

3. Methodology

3.1. Dataset Description

The study uses the Kaggle Cardiovascular Disease dataset (CVDd) [4], comprising 70,000 records and 11 features, including age, height, weight, gender, blood pressure, cholesterol levels, etc., essential for predicting cardiovascular disease. Preprocessing steps included the following:

• Handling Missing Values and Categorical Variables: No missing values or categorical variables were present, avoiding the need for imputation or encoding.
• Normalization/Standardization: Numerical features were scaled using Robust Scaler [13], which centers data around the median and scales it based on the interquartile range (IQR), making it less sensitive to outliers (Figure 1) compared to Z-score normalization.
• Class Imbalance Verification: After verifying the class distribution, we found that the dataset was balanced, meaning no additional techniques (e.g., SMOTE) were required to address class imbalance (Figure 2).

3.2. Feature Selection

To identify relevant features and reduce redundancy, we used correlation analysis (Figure 3).

• Correlation Analysis: Pearson correlation coefficients [14,15,16] were calculated to measure the linear relationship between each feature and the target variable.
  Features with correlation coefficients below a threshold of 0.2 were considered less relevant and removed. This threshold was chosen to strike a balance between retaining features with meaningful relationships to the target variable and eliminating those with negligible contributions, thereby reducing dimensionality while preserving predictive power. The Pearson correlation coefficient for two variables X and Y is given by the following expression:

  r_xy = (cov(X,Y))/(σx σy)

  where cov(X,Y) is the covariance between X and Y, and σx,σy are the standard deviations of X and Y, respectively.

3.3. Regularization

To address multicollinearity and prevent overfitting, we applied L1 (Lasso) and L2 (Ridge) regularization techniques [17]:

• L1 Regularization (Lasso): Lasso adds a penalty equal to the absolute value of the magnitude of coefficients to the loss function;
• L2 Regularization (Ridge): Ridge adds a penalty equal to the square of the magnitude of coefficients; this technique reduces the impact of multicollinearity by discouraging large coefficients, leading to more stable and generalizable models.

3.4. Model Optimization

To enhance performance and reduce overfitting:

• Hyperparameter Tuning: Grid Search optimized hyperparameters for each algorithm [18].
• Validation Strategy: Stratified k-Fold cross-validation (k = 10 or k = 5) preserved class distribution, reducing bias [19].
• Dataset Splitting: The dataset was split into 70% training, 20% testing, and 10% validation to ensure generalizability [20].

3.5. Machine Learning Algorithms Tested

The following classification and regression algorithms were implemented and tested:

• Classification: Logistic Regression, Random Forest, SVM, Gradient Boosting (XGBoost, LightGBM, CatBoost), k-NN.
• Regression: Linear Regression, Random Forest Regressor, SVR, Gradient Boosting Regressor (XGBoost, LightGBM, CatBoost), k-NN.

3.6. Performance Evaluation

Model performance was assessed using the following metrics:

• Classification Metrics [21]: Accuracy, Precision, Recall, F1-Score, AUC-ROC.
• Regression Metrics [22]: Mean Squared Error (MSE), Mean Absolute Error (MAE), R² (Coefficient of Determination).

4. Results and Discussion

This section presents the results of the experiments conducted on the Kaggle Cardiovascular Disease dataset using various classification and regression algorithms. The performance of each model is evaluated on both the validation and test sets, and the results are summarized in tables for clarity. A detailed analysis, general discussion, and recommendations are provided.

4.1. Classification Results

The performance of the classification algorithms is summarized in

Table 1. Performance of classification algorithms on the validation set.

Algorithm	Accuracy	Precision	Recall	F1-Score	AUC-ROC
Logistic Regression	0.712	0.714	0.712	0.712	0.775
Random Forest	0.729	0.731	0.729	0.728	0.792
XGBoost	0.733	0.734	0.733	0.732	0.795
LightGBM	0.731	0.733	0.731	0.730	0.793
CatBoost	0.729	0.731	0.729	0.729	0.793
k-NN	0.706	0.706	0.706	0.706	0.753
SVM	0.720	0.722	0.720	0.719	0.777

and

Table 2. Performance of classification algorithms on the testing set.

Algorithm	Accuracy	Precision	Recall	F1-Score	AUC-ROC
Logistic Regression	0.716	0.718	0.716	0.715	0.782
Random Forest	0.731	0.733	0.731	0.730	0.797
XGBoost	0.735	0.737	0.735	0.735	0.800
LightGBM	0.732	0.733	0.732	0.731	0.798
CatBoost	0.730	0.731	0.730	0.730	0.797
k-NN	0.703	0.703	0.703	0.703	0.759
SVM	0.721	0.723	0.721	0.720	0.784

. Key metrics include Accuracy, Precision, Recall, F1-Score, and AUC-ROC.

• Best Performing Models: XGBoost and LightGBM achieved the highest AUC-ROC scores (0.800 and 0.798, respectively), excelling in capturing complex data relationships.
• Moderate Performers: Random Forest and CatBoost delivered competitive AUC-ROC scores (~0.797), showing robustness and reduced overfitting.
• Lower Performers: Logistic Regression and SVM performed moderately, with AUC-ROC scores of 0.782 and 0.784, respectively, but were outperformed by ensemble methods.
• Weakest Performer: k-NN underperformed (AUC-ROC: 0.759), likely due to sensitivity to data scaling and inability to capture complex patterns.

4.2. Regression Results

The performance of the regression algorithms is summarized in

Table 3. Performance of regression algorithms on the validation set.

Algorithm	MSE	MAE	R2
Linear Regression	0.219	0.441	0.123
Random Forest Regressor	0.182	0.366	0.274
SVR	0.193	0.398	0.229
XGBoost Regressor	0.181	0.368	0.274
LightGBM Regressor	0.181	0.366	0.275
CatBoost Regressor	0.181	0.364	0.276
k-NN Regressor	0.185	0.364	0.258

and

Table 4. Performance of regression algorithms on the testing set.

Algorithm	MSE	MAE	R2
Linear Regression	0.220	0.442	0.120
Random Forest Regressor	0.180	0.365	0.280
SVR	0.192	0.396	0.233
XGBoost Regressor	0.180	0.368	0.279
LightGBM Regressor	0.180	0.365	0.280
CatBoost Regressor	0.180	0.363	0.282
k-NN Regressor	0.183	0.362	0.267

. Key metrics include Mean Squared Error (MSE), Mean Absolute Error (MAE), and R².

• Best Performing Models: CatBoost and LightGBM achieved the highest R² values (0.282 and 0.280, respectively), excelling in modeling complex relationships and minimizing errors.
• Moderate Performers: XGBoost and Random Forest Regressor performed well, with R² values around 0.279 and 0.280, respectively, proving robust for regression tasks.
• Lower Performers: SVR and k-NN Regressor delivered moderate results (R²: 0.233 and 0.267, respectively), limited by their inability to handle non-linear relationships effectively.
• Weakest Performer: Linear Regression performed poorly (R²: 0.120), as expected due to its simplicity and inability to capture non-linear patterns.

4.3. General Discussion

The experimental results highlight the following key insights:

• Superiority of Ensemble Methods: Ensemble methods (XGBoost, LightGBM, CatBoost, Random Forest) outperformed traditional algorithms (Logistic Regression, Linear Regression, k-NN) due to their ability to capture complex, non-linear relationships.
• Importance of Hyperparameter Tuning: Grid Search significantly improved performance, as seen with XGBoost, LightGBM, and CatBoost.
• Robustness of Gradient Boosting: XGBoost and LightGBM excelled in handling high-dimensional data and minimizing overfitting.
• Limitations of Traditional Algorithms: Linear Regression and k-NN underperformed due to simplicity and sensitivity to data scaling, respectively.
• Balanced Dataset: The dataset was balanced, ensuring stable and reliable results.

4.4. Recommendations

Based on the findings, the following recommendations are proposed:

• Adopt Ensemble Methods: Prioritize XGBoost, LightGBM, and CatBoost for CVD prediction due to their superior performance.
• Invest in Hyperparameter Optimization: Use Grid Search or Bayesian Optimization to maximize model performance.
• Explore Advanced Techniques: Investigate deep learning or hybrid models for further accuracy improvements.
• Focus on Interpretability: Use SHAP to explain predictions and build trust among healthcare professionals.
• Validate on External Datasets: Ensure generalizability by testing on external datasets.

5. Conclusions and Future Work

This study highlights the superiority of ensemble methods, particularly XGBoost and LightGBM, for CVD prediction. Future work should focus on external validation, integrating unstructured data, and developing explainable AI techniques to enhance clinical applicability and improve patient outcomes.