Application of Machine Learning Algorithms for Predicting Infant Mortality in India: An Analysis of the National Family Health Survey-5, 2019–2021

Abstract

Objectives

Machine learning (ML) techniques have shown strong potential for predicting infant mortality (IM), but their application in the Indian context remains limited. This study aimed to use ML algorithms to predict IM in India using a large national survey database.

Methods

Data were analyzed from the National Family Health Survey-5, 2019–2021, a large cross-sectional survey. Random forest, decision tree, adaptive boosting, logistic regression, and naïve Bayes models were implemented using Weka version 3.8.3. Model performance was evaluated using accuracy, precision, F1-score, Matthews correlation coefficient, and area under the curve (AUC).

Results

Compared with logistic regression, which achieved 65.2% accuracy, the random forest and decision tree models showed higher predictive accuracy, at 74.1% and 73.2%, respectively. Their AUCs were 0.80 and 0.79, respectively, compared with 0.69 for logistic regression. The models identified birth order, maternal education, twin birth, wealth index, cooking fuel use, and age at first birth as the six strongest predictors of IM.

Conclusions

In this analysis, random forest and decision tree models outperformed logistic regression in predicting IM. These findings underscore the relevance of sociodemographic and economic disparities and support the use of ML algorithms for risk prediction and targeted interventions aimed at reducing IM.

I. Introduction

The infant mortality rate (IMR) is a key indicator of a nation’s overall health and socioeconomic development because it reflects the accessibility and quality of health care services, maternal health conditions, and broader socioeconomic determinants. India has made considerable progress in reducing IMR over recent decades; however, substantial disparities persist across states, rural and urban settings, and socioeconomic groups. Identifying the determinants of infant mortality and accurately predicting IMR may help policymakers implement targeted interventions to reduce preventable infant deaths.

Traditional epidemiological models have played an important role in characterizing infant mortality trends and predictors, but their predictive performance can be constrained by linear assumptions and limited modeling of variable interactions. Machine learning (ML) techniques offer an alternative approach by modeling complex, nonlinear relationships in large datasets. ML models can identify patterns, account for high-dimensional interactions among risk factors, and support more accurate prediction. With the availability of large-scale demographic and health datasets, such as the National Family Health Survey (NFHS), ML methods may improve predictive modeling of infant mortality and inform evidence-based interventions.

Several ML algorithms have shown promise for health outcome prediction, including decision trees, random forests, support vector machines, naïve Bayes, and deep learning models. These techniques have been applied to several child health and neonatal outcomes, including infant and underfive mortality, clinical outcomes in sick neonates, identification of at-risk neonates, early warning scores for neonatal complications, and low birth weight [1–11]. Despite this potential, ML-based prediction of infant mortality in India remains underexplored, with only two studies using data from the previous round of the NFHS [12,13]. A systematic comparison of ML models using the more recent NFHS-5 dataset may clarify their performance and applicability for policy planning.

This analysis included a broad feature set comprising socioeconomic, household, and maternal characteristics. By evaluating the predictive accuracy and interpretability of multiple ML models, this study aimed to identify high-risk populations and support targeted strategies to reduce infant mortality.

These findings may have implications for public health policy in India. If ML models predict infant mortality accurately, they could be integrated into health surveillance systems to support real-time risk assessment. This study contributes to the literature on AI-driven health analytics for addressing major public health challenges in developing countries.

Against this background, the specific objectives of this study were as follows:

1. 1) To use ML algorithms to predict infant mortality in India using NFHS-5 data.

2. 2) To compare the predictive performance of multiple ML algorithms for identifying infant mortality risk.

3. 3) To identify the most important socioeconomic, demographic, and maternal predictors of infant mortality.

II. Methods

1. Study Setting and Design

India, located in South Asia, is the world’s most populous country, with more than 1.4 billion people, and the seventh largest country by land area. It is a federal democratic republic comprising 28 states and 8 Union Territories. Life expectancy in India is approximately 70 years. The infant mortality rate has declined to 28 deaths per 1,000 live births, and the maternal mortality ratio is 97 deaths per 100,000 live births. These improvements reflect better health care access and increased institutional delivery, which now accounts for 90% of all births in the country. India’s total fertility rate has declined to 2.0, below the replacement level of 2.1, suggesting progress toward population stabilization.

This study was a secondary analysis of data collected as part of the NFHS-5, 2019–2021, a nationally representative cross-sectional household survey.

2. Data Source

The NFHS-5 is a nationwide survey designed to provide reliable and up-to-date information on key indicators related to fertility, mortality, maternal and child health, nutrition, family planning, health service use, and gender-based issues such as domestic violence. It is the fifth round of the NFHS series and was coordinated by the International Institute for Population Sciences (IIPS), Mumbai, under the Ministry of Health and Family Welfare, Government of India. The NFHS-5 informs policymaking, supports monitoring of health program outcomes, and tracks progress toward national and global targets, including the Sustainable Development Goals [14].

Methodologically, NFHS-5 used a multistage stratified sampling design, with the 2011 Census as the sampling frame. In rural areas, villages were selected using probability proportional to size sampling, whereas in urban areas, census enumeration blocks served as the primary sampling units. The survey covered all 707 districts in India and collected data from 610,000 households, including interviews with women aged 15–49 years and men aged 15–54 years. Data were collected using four standardized questionnaires—Household, Woman’s, Man’s, and Biomarker—and computer-assisted personal interviewing technology was used to support accuracy and efficiency [14].

3. Study Variables

The dependent variable was infant mortality, defined as the death of a live-born infant before the first birthday. This outcome was estimated using information from the birth history section of the questionnaire administered to individual women. The independent variables included place of residence, marital status, maternal education, maternal occupation, wealth index, maternal age at first birth, child sex, birth order, number of children ever born, breastfeeding practice, source of drinking water, type of toilet facility, and type of cooking fuel. These variables were selected on the basis of existing literature on predictors of infant mortality.

4. Data Preprocessing

Data preprocessing is a necessary step in the ML pipeline because it transforms raw and potentially noisy data into a structure suitable for model training. Before model training, two key preprocessing steps were performed: handling missing data and addressing class imbalance. A total of 1,274,250 babies born to nearly 7 million women who participated in the NFHS-5 were initially included in this analysis. Missing data were found only for three variables: respondent caste, source of drinking water, and toilet facility. Records with missing values for these variables, which accounted for 7.4% of the dataset (n = 95,009), were removed, leaving 1,179,241 records for analysis. In the selected dataset, the outcome classes were highly imbalanced: 1,125,642 records (95%) were in the “No infant death” class, whereas 53,599 records (5%) were in the “Infant death” class. This imbalance could bias trained algorithms toward the majority class and increase the likelihood that new observations would be classified as no infant death.

5. Model Building

The dataset was randomly divided into a training set for model development and a test set for model performance evaluation using the get_train_test command in Stata version 18.0 (StataCorp LLC, College Station, TX, USA). An 80:20 split was applied, with 80% of the data allocated to training and 20% allocated to testing. The training set included 943,393 rows and 42,735 infant deaths (4.5%), and the test set included 235,848 rows and 10,864 infant deaths (4.6%). A 10-fold cross-validation analysis and a 70:30 split were also performed to assess model performance and generalizability. Class imbalance was addressed by applying the Synthetic Minority Oversampling Technique (SMOTE) to the training dataset until the minority class was balanced with the majority class. This process produced an oversampled balanced dataset with 899,703 records (50%) in the “No infant death” class and 898,849 records (50%) in the “Infant death” class for ML model development. Data preprocessing was performed in RStudio version 2024.12.1+563 (https://www.r-studio.com/), and the processed dataset was exported in CSV format.

6. Statistical Analysis

The chi-square test was used to assess associations between predictor variables and infant mortality. Variables significantly associated with infant mortality were then incorporated into a multivariable logistic regression model using Stata version 18.0. This model was compared with four ML models—decision tree, random forest, naïve Bayes, and adaptive boosting—using predictive accuracy, precision, F-measure, area under the curve (AUC), and Matthews correlation coefficient (MCC). The information gain ranking method was used to identify the most important variables for predicting infant mortality. ML models were implemented using Weka version 3.8.3 (https://ml.cms.waikato.ac.nz/weka/). Default hyperparameter settings provided by Weka were used for all algorithms, and no additional hyperparameter tuning or optimization was performed.

7. Ethics Approval

This study did not require additional ethical approval because NFHS-5 had already received ethical clearance from the Ethical Review Board of the IIPS and the ICF International Review Board before the survey was conducted. The authors were granted access to the anonymized dataset through the DHS Program, which is publicly available upon reasonable request at https://dhsprogram.com/data/available-datasets.cfm. NFHS-5 followed the principles outlined in the Declaration of Helsinki and maintained participant confidentiality and anonymity.

III. Results

1. Background Characteristics of the Study Population

Table 1 presents the sociodemographic and maternal characteristics of the study sample. Among the 1,179,241 children included in the analysis, 53,599 (4.7%) died before their first birthday. Nearly three-fourths of the children were from rural households (926,750, 71.7%). Most households had an improved drinking water source (1,089,702, 94.6%), and approximately three-fourths had an improved toilet facility (901,119, 75.7%). Approximately half of the households used unclean cooking fuel (625,185, 52.6%). Nearly half of the households belonged to the Other Backward Class category (479,060, 46.0%), and 34.4% belonged to the Scheduled Caste or Scheduled Tribe categories (n = 493,853). Most households had 4–6 members (708,758, 59.8%).

Table 1.

Socio-demographic and maternal characteristics of the study sample, National Family Health Survey-5, 2019–21 (n = 1,179,241)

Characteristic	n (%)
Infant mortality
Yes	53,599 (4.7)
No	1,125,642 (95.3)
Place of residence
Rural	926,750 (71.7)
Urban	252,491 (28.3)
Education of the mother
No education	472,310 (39.7)
Upto Primary	187,995 (15.6)
Upto Secondary	444,162 (37.3)
Higher education	74,774 (7.4)
Caste of the household
Scheduled Caste	250,438 (24.1)
Scheduled Tribe	243,415 (10.3)
Other Backward Class	479,060 (46.0)
Others	206,328 (19.6)
Source of drinking
Improved	1,089,702 (94.6)
Unimproved	89,539 (5.4)
Toilet facility
Improved	901,119 (75.7)
Unimproved	278,122 (24.3)
Marital status
Unmarried	1,373 (0.1)
Married	1,111,539 (94.4)
Widowed/Divorced/Separated	66,329 (5.5)
Wealth index
Lower	573,615 (44.2)
Middle	240,155 (20.5)
Richer	365,471 (35.3)
Any comorbidity
Yes	111,521 (10.1)
No	1,067,720 (89.9)
Cooking fuel used
Unclean	625,185 (52.6)
Clean	552,805 (47.4)
Number of household members
1–3	162,868 (13.7)
4–6	708,758 (59.8)
≥7	307,615 (26.5)
Age at first birth (yr)
<18	260,061 (23.4)
18–34	916,613 (76.4)
≥35	2,567 (0.2)
Birth order
1–2 births	818,436 (70.5)
3–4 births	287,641 (23.6)
≥5 births	73,164 (5.9)
Sex of the child
Male	618,841 (52.5)
Female	560,400 (47.5)
Twin birth
Yes	15,742 (1.4)
No	1,163,499 (98.6)
Current tobacco use
Yes	108,976 (6.1)
No	1,070,265 (93.9)
Secondhand smoking inside home
Yes	606,419 (51.2)
No	572,822 (48.8)

Numbers are unweighted counts, and percentages are weighted proportions.

2. Factors Associated with Infant Mortality

Table 2 presents the univariable associations between sociodemographic and maternal factors and infant mortality. The chi-square test showed that all examined factors were significantly associated with infant mortality (p < 0.001).

Table 2.

Socio-demographic and maternal factors associated with infant mortality, National Family Health Survey-5, 2019–21

Characteristic	Total	Univariable		Multivariable
		Infant deaths	p-value	aOR (95% CI)	p-value
Place of residence			<0.001
Rural	926,750	45,028 (4.9)		1.10 (1.06–1.12)	<0.001
Urban	252,491	8,571 (3.4)		Ref
Education of the mother			<0.001
No education	472,310	27,803 (5.9)		2.13 (2.01–2.25)	<0.001
Upto Primary	187,995	9,207 (4.9)		1.85 (1.75–1.96)	<0.001
Upto Secondary	444,162	15,065 (3.4)		1.43 (1.35–1.51)	<0.001
Higher education	74,774	1,524 (2.0)		Ref
Caste of the household			<0.001
Scheduled Caste	250,438	13,531 (5.4)		1.19 (1.15–1.22)	<0.001
Scheduled Tribe	243,415	9,987 (4.1)		0.82 (0.80–0.85)	<0.001
Other Backward Class	479,060	22,489 (4.7)		1.11 (1.08–1.14)	<0.001
Others	206,328	7,592 (3.7)		Ref
Source of drinking			<0.001
Improved	1,089,702	49,463 (4.5)		Ref
Unimproved	89,539	4,136 (4.6)		0.97 (0.94–1.01)	0.08
Toilet facility			<0.001
Improved	901,119	37,255 (4.1)		Ref
Unimproved	278,122	16,344 (5.9)		1.12 (1.09–1.14)	<0.001
Marital status			<0.001
Unmarried	1,373	43 (3.1)		Ref
Married	1,111,539	50,025 (4.5)		1.46 (1.08–1.99)	0.015
Widowed/Divorced/Separated	66,329	3,531 (5.3)		1.39 (1.02–1.89)	0.038
Wealth index			<0.001
Lower	573,615	31,981 (5.6)		1.24 (1.20–1.28)	<0.001
Middle	240,155	10,011 (4.2)		1.07 (1.04–1.11)	<0.001
Richer	365,471	11,607 (3.2)		Ref
Any comorbidity			<0.001
Yes	111,521	5,845 (5.2)		1.17 (1.13–1.20)	<0.001
No	1,067,720	47,754 (4.5)		Ref
Cooking fuel used			<0.001
Unclean	625,185	33,284 (5.3)		1.18 (1.15–1.21)	<0.001
Clean	552,805	20,248 (3.7)		Ref
Number of household members			<0.001
1–3	162,868	10,760 (6.6)		Ref
4–6	708,758	31,006 (4.4)		0.64 (0.63–0.66)	<0.001
≥7	307,615	11,833 (3.8)		0.52 (0.50–0.53)	<0.001
Age at first birth (yr)			<0.001
<18	260,061	16,125 (6.2)		1.59 (1.30–1.96)	<0.001
18–34	916,613	37,377 (4.1)		1.16 (0.94–1.42)	0.165
≥35	2,567	97 (3.8)		Ref
Birth order			<0.001
1–2 births	818,436	35,485 (4.3)		1.12 (1.09–1.14)	<0.001
3–4 births	287,641	13,305 (4.6)		Ref
≥5 births	73,164	4,809 (6.6)		1.30 (1.26–1.35)	<0.001
Sex of the child			<0.001
Male	618,841	29,650 (4.8)		1.12 (1.10–1.14)	<0.001
Female	560,400	23,949 (4.3)		Ref
Twin birth			<0.001
Yes	15,742	3,488 (22.2)		7.10 (6.82–7.28)	<0.001
No	1,163,499	50,111 (4.3)		Ref
Current tobacco use			<0.001
Yes	108,976	6,282 (5.8)		1.18 (1.14–1.21)	<0.001
No	1,070,265	47,317 (4.4)		Ref
Secondhand smoking inside home	<0.001	1.10 (1.08–1.12)	<0.001
Yes	606,419	29,302 (4.8)		Ref
No	572,822	24,297 (4.2)

Values are presented as number (%).

aOR: adjusted odds ratio, CI: confidence interval.

Table 2 also presents the strength of association between each predictor and infant mortality as adjusted odds ratios (aORs) from the multivariable logistic regression model. Compared with urban residents, mothers with higher education, and households in the richer wealth index category, respectively, rural residents (aOR: 1.10, 95% confidence interval [CI]: 1.06–1.12), mothers with no formal education (aOR: 2.13, 95% CI: 2.01–2.25), and households in the lower wealth index category (aOR: 1.24, 95% CI: 1.20–1.28) had higher odds of infant mortality. Households using unclean cooking fuel and households with unimproved toilet facilities had higher odds of infant mortality, with aORs of 1.18 (95% CI: 1.15–1.21) and 1.12 (95% CI: 1.09–1.14), respectively. Maternal age at first birth below 18 years was also significantly associated with infant mortality (aOR: 1.59, 95% CI: 1.30–1.96). Figure 1 shows the strength of association between predictors and infant mortality, expressed as aORs with 95% CIs.

Figure 1.

Coefficient plot showing the strength of association between predictors and infant mortality, expressed as adjusted odds ratios with 95% confidence intervals.

3. Comparison of ML Models

The logistic regression model predicted infant mortality with an accuracy of 65.2%, an MCC of 0.33, and an AUC of 0.69. In the 10-fold validation analysis, the random forest model achieved an accuracy of 74.1%, an MCC of 0.50, and an AUC of 0.80. The decision tree algorithm achieved an accuracy of 73.2%, an MCC of 0.49, and an AUC of 0.79. Naïve Bayes performed similarly to logistic regression, with an accuracy of 60.4%, an MCC of 0.25, and an AUC of 0.62. Model performance was similar across the 80:20 split, 70:30 split, and 10-fold validation analyses, as shown in Table 3.

Table 3.

Performance of predictive machine learning models of infant mortality as evaluated on the test data, National Family Health Survey-5, 2019–21

Training-test ratio	Measures	Logistic regression	Decision tree	Random forest	Adaptive boosting	Naïve Bayes
70:30	Accuracy	65.2	73.2	74.1	61.3	60.4
	Precision	65.2	73.2	74.2	61.3	60.4
	F-measure	65.2	73.2	74.1	61.3	60.4
	MCC	0.33	0.48	0.50	0.28	0.25
	AUC	0.69	0.79	0.80	0.66	0.62
	MAE	0.42	0.32	0.30	0.42	0.40
	RMSE	0.44	0.40	0.39	0.45	0.45
80:20	Accuracy	65.2	73.0	76.2	61.3	60.4
	Precision	65.2	73.0	74.3	61.3	60.4
	F-measure	65.2	73.0	74.2	61.3	60.4
	MCC	0.33	0.49	0.50	0.28	0.25
	AUC	0.69	0.79	0.80	0.66	0.62
	MAE	0.42	0.32	0.30	0.42	0.40
	RMSE	0.44	0.40	0.39	0.45	0.45
10-fold	Accuracy	65.2	73.2	74.2	61.3	60.4
	Precision	65.2	73.2	74.3	61.3	60.4
	F-measure	65.2	73.2	74.2	61.3	60.4
	MCC	0.33	0.49	0.50	0.28	0.25
	AUC	0.69	0.79	0.80	0.66	0.62
	MAE	0.43	0.32	0.30	0.42	0.40
	RMSE	0.45	0.40	0.39	0.45	0.45

MCC: Matthew’s correlation coefficient, AUC: area under the curve, MAE: mean absolute error, RMSE: root mean squared error.

Based on these performance metrics, random forest was the best-performing model for predicting infant mortality, followed closely by the decision tree algorithm.

4. Variable Importance

The 10 variables that contributed most to infant mortality prediction, in decreasing order of importance, were birth order, maternal education, twin birth, wealth index, cooking fuel use, maternal age at first birth, number of household members, toilet facility, place of residence, and household caste.

IV. Discussion

This study examined the use of ML methods to predict infant mortality in a large nationally representative sample of nearly 1.2 million children in India. Random forest and decision tree models predicted infant mortality more accurately than the conventional logistic regression model. The predictive models identified birth order, maternal education, twin birth, wealth index, cooking fuel use, and maternal age at first birth as the six most important predictors of infant mortality. These findings suggest that infant mortality is shaped by social, economic, demographic, and maternal health-related factors and may inform national public health policies and programs aimed at reducing IMR.

Infants are particularly susceptible to indoor air pollutants because of their developing respiratory systems and frequent proximity to cooking areas [15]. Previous evidence from India and similar settings has shown an association between household solid fuel use and the risk of stillbirth, neonatal death, and infant death [16–19]. The Million Death Study also attributed a substantial proportion of neonatal and post-neonatal deaths to solid fuel use [20,21]. Similarly, the Global Burden of Disease Study attributed more than 0.5 million deaths among children younger than 5 years to household air pollution, representing almost 11% of total under-5 mortality [22].

This study identified a significant association between younger maternal age at first birth and increased infant mortality risk in India, consistent with previous literature. Physical immaturity among younger mothers, together with poor access to antenatal, delivery, and postnatal care, lower health-seeking behavior, and reduced autonomy, may contribute to maternal and neonatal complications and mortality [23]. Evidence from India supports this interpretation [24,25]. These findings are also consistent with global evidence. In a multicountry analysis, Finlay et al. [26] reported that early maternal age was a consistent predictor of child mortality in low-resource settings, even after adjustment for confounding factors. Therefore, efforts to reduce early childbearing through delayed marriage and improved girls’ education may contribute to better child survival.

Socioeconomic factors, including wealth index, maternal education, and household caste, were significant predictors of infant mortality, consistent with evidence from low- and middle-income countries [27]. This finding indicates that reducing infant mortality requires not only improved healthcare services but also efforts to address structural inequalities in wealth, education, and social inclusion. Conditional cash transfers, promotion of girls’ education, and community outreach in marginalized caste and tribal areas may help reach underserved populations.

A notable finding was the lower odds of infant mortality among Scheduled Tribe households. This result appears to diverge from commonly reported socioeconomic disparities in India. Several factors may explain this finding. First, targeted maternal and child health interventions have been implemented in Scheduled Tribe-dominated areas, which may also explain improved maternal antenatal care indicators reported for this group in NFHS-5 compared with other caste categories [14]. Second, differences in fertility behavior, birth spacing, and community practices may influence infant survival. In addition, national-level aggregation may obscure substantial heterogeneity across states and districts. Therefore, this finding should be interpreted cautiously and examined in more granular, region-specific analyses.

ML algorithms have shown promise for predicting infant mortality, with studies demonstrating their performance across different datasets and modeling approaches. In this study, prediction accuracy and AUC were highest for the random forest model, followed by the decision tree model. These findings are consistent with those of Mfateneza et al. [4], who used ML to predict infant mortality in Rwanda and reported the highest accuracy and AUC for the random forest model (accuracy: 84.3%, AUC: 84.2%), followed by the decision tree model. Random forest has also been among the top-performing algorithms in studies predicting under-5 mortality in Nigeria, Bangladesh, and Ethiopia [6,8,28]. XGBoost, an ensemble algorithm based on decision trees, showed strong performance in predicting infant mortality in India, with an accuracy of 97.6% and an AUC of 0.974 [13]. Another nationwide analysis using the previous round of NFHS data from India reported AUC values ranging from 45% to 55% for algorithms such as logistic least absolute shrinkage and selection operator, boosted classification trees, and random forest, representing an approximately 6%–9% improvement over standard logistic regression [12].

This study had several strengths. It used a large, nationally representative survey dataset with a high response rate, standardized methodology, and validated interview tools, which supports the generalizability of the findings. Several limitations also warrant discussion. First, traditional logistic regression provides interpretable regression coefficients that indicate the strength and direction of associations, whereas ML methods do not provide the same level of interpretability. Instead, ML approaches identify predictive factors based on each variable’s relative importance in estimating mortality risk. Second, the cross-sectional design precludes causal interpretation. Longitudinal studies using ML methods are needed to further evaluate predictors of infant mortality. Third, because ML methods have been used relatively sparingly in comparable contexts, evidence for direct comparison with other countries or similar settings remains limited. Fourth, although the risk of data leakage was minimized by applying SMOTE only to the training dataset, synthetic sample bias remains possible. Therefore, the reported performance metrics should be interpreted with caution. Future studies could explore alternative approaches, such as cost-sensitive learning or ensemble techniques tailored to imbalanced data.

The use of ML algorithms to predict childhood mortality has several methodological and public health implications. Methodologically, ML techniques can model complex, nonlinear relationships and high-dimensional data more flexibly than traditional statistical methods [29,30]. This flexibility may improve predictive accuracy, support earlier identification of children at elevated risk, and guide timely targeted interventions. These models could be integrated into routine hospital health information systems to identify high-risk cases. However, model interpretability and dependence on data quality remain important concerns.

Future research should examine approaches for model optimization, interpretability, and ensemble learning to improve the accuracy of infant mortality prediction. Additional studies should assess the transferability of these predictive models to clinical and public health settings and externally validate them in diverse datasets to strengthen the generalizability of the findings.