This paper presents a comprehensive study on the application of modern machine learning methods for predictive maintenance based on the open AI4I Predictive Maintenance dataset. The primary goal of the research is to develop and compare both binary and multiclass classification models that enable not only the prediction of machine failures but also the identification of specific failure types. Considering the strong class imbalance (failures account for approximately 3 % of the dataset), a cost-sensitive optimization approach was implemented, where false negatives (missed failures) were penalized much more heavily than false positives (false alarms).
The models investigated in this study include Logistic Regression, Random Forest, XGBoost, and LightGBM, with a particular focus on ensemble gradient boosting methods. Data preprocessing steps encompassed feature scaling, categorical encoding, and feature engineering, including the creation of new features such as Power, DeltaTemp, and a binary indicator of excessive tool wear. For binary classification, model performance was evaluated using the PR-AUC metric and Recall@Precision≥0.90, both of which are critical in safety-related domains. For multiclass classification, performance was assessed using confusion matrices and macro-averaged precision, recall, and F1-score.
The experimental results demonstrated that among the baseline models, XGBoost achieved the best performance with PR-AUC=0.64, and after hyperparameter optimization, the score improved significantly to over 0.90, confirming its strong ability to detect failures. In the multiclass setup, XGBoost again outperformed other models, providing better balance across classes, while LightGBM faced challenges due to "no further splits with positive gain" warnings under restrictive parameter settings. Feature importance analysis using SHAP values highlighted the dominant role of Tool wear, Power, and Rotational speed in predicting failures.
The practical contribution of this work lies in the demonstrated potential of machine learning models for early failure detection and proactive maintenance planning, enabling significant cost savings and operational reliability improvements. The study also emphasizes the effectiveness of combining cost-sensitive optimization with ensemble methods and establishes a foundation for future research directions, including time-to-failure regression and the use of deep neural networks.
1. Ahmed, N., Aljahdali, H. M., & Siddiqui, M. H. (2021). Predictive maintenance using machine learning for industrial machinery: A survey. Journal of King Saud University - Computer and Information Sciences, 33(10), 1189-1201. https://doi.org/10.1016/j.jksuci.2021.01.007
2. Zonta, T., da Costa, C. A., da Rosa Righi, R., de Lima, M. J., da Trindade, E. S., & Li, G. P. (2020). Predictive maintenance in the Industry 4.0: A systematic literature review. Computers & Industrial Engineering, 150. https://doi.org/10.1016/j.cie.2020.106889
3. Chen, T., & Guestrin, C. (2016). XGBoost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining (pp. 785-794). ACM. https://doi.org/10.1145/2939672.2939785
4. Ke, G., Meng, Q., Finley, T., Wang, T., Chen, W., Ma, W., Ye, Q., & Liu, T. Y. (2017). LightGBM: A highly efficient gradient boosting decision tree. Advances in Neural Information Processing Systems, 30, 3146-3154. https://papers.nips.cc/ paper/2017/hash/6449f44a102fde848669bdd9eb6b76fa-Abstract.html
5. Haque, M. A., Hasan, M. K., Islam, S., & Alhumyani, H. (2022). Handling imbalanced data in machine learning: A survey. Knowledge-Based Systems, 248, 108759. https://doi.org/10.1016/j.knosys.2022.108759
6. Lundberg, S. M., & Lee, S. I. (2017). A unified approach to interpreting model predictions. In Proceedings of the 31st International Conference on Neural Information Processing Systems (pp. 4765-4774). Curran Associates Inc. https://doi.org/10.48550/arXiv.1705.07874
7. Zhang, H., Zhao, R., Wang, J., & Li, X. (2023). Multi-failure mode prediction in predictive maintenance: A deep learning approach. Reliability Engineering & System Safety, 232, 109032. https://doi.org/10.1016/j.ress.2023.109032
8. Carvalho, T. P., Soares, F. A. A. M. N., Vita, R., Francisco, R. P., Basto, J. P., & Alcalá, S. G. S. (2019). A systematic literature review of machine learning methods applied to predictive maintenance. Computers & Industrial Engineering, 137. https://doi.org/10.1016/j.cie.2019.106024
9. Zhang, C., Xu, X., Wang, J., & Zhong, R. Y. (2022). Data-driven predictive maintenance for complex equipment: A review. Engineering Applications of Artificial Intelligence, 109, 104653. https://doi.org/10.1016/j.engappai.2021.104653
[10] Tuli, S., Casale, G., & Jennings, N. R. (2022). TranAD: Deep transformer networks for anomaly detection in multivariate time series data. Proceedings of the VLDB Endowment, 15(6), 1201-1214. https://doi.org/10.14778/3514061.3514067