Sample size calculation for training ensemble machine learning models on health data

Area of Science:

• Machine Learning
• Health Research
• Statistical Modeling

Background:

• Health research studies frequently face limitations due to small sample sizes.
• Training machine learning (ML) models necessitates substantial datasets, creating a gap in guidance for adequate sample size determination.
• Existing literature lacks comprehensive methods for calculating sample sizes for ML model development.

Purpose of the Study:

• To develop an empirically derived sample size calculator for ensemble ML models.
• To predict the necessary sample size for achieving a specific level of prognostic performance with a defined probability.
• To compare the accuracy of the developed calculator against common heuristics and statistical approaches.

Main Methods:

• Developed an empirically derived sample size calculator for ensemble ML models, including random forests, light gradient boosting machine (LGBM), and extreme gradient boosting (XGBoost).
• Defined prognostic performance as the sample area under the ROC curve (ROC-AUC) relative to the optimal model trained on the full dataset.
• Compared the calculator's accuracy against three common heuristics and one statistical approach for sample size calculation.

Main Results:

• The developed calculator demonstrated significantly better accuracy for tree-based ensemble ML models compared to other methods.
• For instance, the median relative error in sample size prediction was 25% for achieving 85% of optimal performance with 90% certainty for LGBM.
• The calculator effectively predicts sample sizes needed for desired prognostic performance levels.

Conclusions:

• The new sample size calculator provides a more accurate method for determining adequate sample sizes in health research utilizing ML.
• This tool addresses the critical need for sample size guidance in ML model development, particularly for tree-based ensemble methods.
• Improved sample size estimation can enhance the reliability and generalizability of ML models in health research.