PaperlayerView exact revision on arXiv

Source-linked AI summary

SAGE: Stability-Aware Graph-Based Ensemble Feature Selection for Explainable Postpartum Depression Risk Prediction

Md. Rokon Islam Emon, Syed Shariar Alam Shuvo, Shahriar Siddique Ayon, Abdullah Al Mamun, Ahnaf Atef Choudhury

arXiv:2608.22809v1cs.LG

TL;DR

Postpartum depression is a serious, multifactorial health burden, while existing prediction approaches face interpretability, feature-selection, and class-imbalance challenges. SAGE combines stability-aware graph-based feature selection, GA-ANN prediction, GAN oversampling, and local explanations, achieving its strongest reported performance with 16 features.

  • Problem

    Postpartum depression is a serious, multifactorial condition with disproportionate burden in low- and middle-income countries, motivating clinically useful prediction methods.

  • Method

    SAGE combines statistical, PCA-based, graph-interaction, and stability-weighted feature selection with GAN-enhanced GA-ANN prediction and LIME explanations.

  • Results

    87.96% accuracy, 86.32% F1-score, and 0.88 AUC were achieved by GA-ANN with SAGE-selected 16 features, the best reported feature-selection configuration.

  • Takeaways & Limitations

    Psychological and socioeconomic factors were key predictors, while LIME provided patient-specific insights for PPD risk stratification and potential clinical decision support.

  • Takeaways & Limitations

    The framework requires testing across regions and longitudinal prenatal-postnatal data to assess broader applicability and temporal risk patterns.

Abstract

from arXiv · show

Postpartum depression (PPD) poses a major burden on maternal and child health, especially in low- and middle-income countries where prevalence exceeds 19%. Despite advancements in machine learning for PPD prediction, current approaches are limited by opaque global explanations that lack clinical usefulness at the patient level, unstable feature selection, and poor generalization under class imbalance. We propose SAGE, a Stability-Aware Graph-Based Ensemble feature selection system that incorporates both local explainable AI and a genetically optimized artificial neural network (GA-ANN). Using a primary cohort of 766 postpartum women, SAGE combines information-theoretic relevance, PCA-based structure, and graph-based interactions with bootstrap stability weighting to identify robust and non-redundant predictors. The GA-ANN architecture, optimized using a genetic algorithm and enhanced with GAN based oversampling, achieved strong performance with 87.96% accuracy, 86.32% F1 score, and 0.88 AUC using only 16 features, outperforming baseline and other feature selection methods. Psychological and socioeconomic factors such as EPDS score, PHQ-9 score, feelings about motherhood, and abuse history are the main predictors, while demographic factors have less influence. The LIME-based explanations allow instance-based insight into selected features from the graph, enabling personalized risk assessment. The findings make SAGE a scalable, interpretable, and clinical tool for early identification of PPD in health-care limited resources.

I. INTRODUCTION

PPD is a serious and prevalent condition with disproportionate impacts in resource-limited settings, while existing ML methods face validation, imbalance, feature-selection stability, and interpretability limitations. SAGE addresses these challenges by combining stability-aware graph-based selection with GA-ANN, GAN oversampling, and local XAI for patient-level PPD risk stratification.

  • Motivation: Resource-limited settings show particularly high burdens, including 21.1% postnatal depression in rural areas and 25.4% prevalence in Bangladesh.Economic hardship and limited mental health services emphasize the importance of early detection in South Asian resource-limited settings.
  • Motivation: PPD reflects combined social, economic, obstetric, and psychological factors, with low social support, marital dissatisfaction, low income, and obstetric complications increasing vulnerability.Depression rates were 23.97% after complicated deliveries versus 16.71% after uncomplicated deliveries.
  • Research gap: Existing ML-based PPD studies are constrained by weak validation, class imbalance, unstable feature selection, and limited interpretability and clinical trust.Graph-based selectors may overlook statistical relevance and structural variability, while ensemble methods can produce inconsistent rankings across resamples.
  • Proposed approach: SAGE combines information-theoretic, PCA, graph-based, and stability-weighted selection with GA-ANN, GAN oversampling, and LIME explanations for robust, patient-level PPD prediction.The framework explicitly models feature interactions and selection robustness.
  • Study contributions: Analysis of 766 postpartum women in Bangladesh identifies psychological and socioeconomic factors as key predictors, while demographic factors are less influential.The study presents the integration of local XAI with graph-based feature selection for patient-level PPD risk stratification as a first contribution.

II. RELATED WORK

Prior PPD prediction studies report strong but variable performance across datasets, modalities, and validation strategies. However, unstable feature selection, weak validation, poor imbalance handling, and limited interpretability motivate graph-oriented selection with optimized models and local explanations.

  • Predictive modeling: PPD prediction performance varies across data sources, with prepartum EPDS improving early prediction and birth and maternal weight identified as important factors.Clapp et al. analyzed 29,168 EHR cases, while Wang et al. examined 3,174 women and reported PPD trends between 30.9% and 29.1%.
  • Predictive modeling: AUC of 0.789-0.955 and accuracy of 95.0% were reported for XGBoost and neural network ensembles on cesarean delivery data, with performance varying after SMOTE.The cited studies also reported an increase in performance from AUC 0.733 to 0.897 on original data after SMOTE.
  • Feature selection: LASSO, variance inflation factor analysis, and Boruta identified compact predictor sets, including 17 predictors with XGBoost AUC of 0.849 and 11 predictors for gradient boosting.Reported predictors included gestational weight gain, in-law relationship, and sleep quality.
  • Research gaps: Most studies use sociodemographic or clinical filtering without validating feature-selection stability, limiting generalizability.This limitation is identified by a systematic review of existing PPD prediction research.
  • Interpretability and validation: Clinical deployment increasingly requires interpretable predictions, alongside external validation and diverse modalities such as plasma proteomics for pre-symptomatic risk stratification.SHAP-based methods have been widely used, while Qi et al. demonstrated externally validated models on large-scale clinical data.
  • Research gaps: Current research often uses weak validation, handles class imbalance poorly, and ignores feature interactions, motivating graph-oriented selection with optimized models and local explanations.The stated direction combines optimized models with GAN oversampling and local explanations.

III. METHODOLOGY · A. Data Collection and Preprocessing

SAGE’s methodology uses a field-collected postpartum cohort, structured preprocessing, oversampling for imbalance, and explainable AI for identifying risk and protective factors. The framework prepares clinical and categorical data for optimized postpartum depression prediction.

  • III. METHODOLOGY: The framework combines data preparation, feature engineering, optimized model development, imbalance handling through oversampling, and XAI-based factor identification.The proposed prediction framework is illustrated in Figure 1.
  • A. Data Collection and Preprocessing: The study collected primary field data from 766 postpartum women in Bangladesh between March and June 2025.Participants were recruited from urban, rural, and outpatient hospital facilities.
  • A. Data Collection and Preprocessing: Participants aged 18–41 were recruited within 24 months of childbirth under ethical guidelines, with informed consent and anonymity.The cohort was drawn from hospital settings across urban, rural, and outpatient facilities.
  • A. Data Collection and Preprocessing: The initial dataset contained 51 features and 766 patient records, while Patient ID and redundant PHQ-9 Result and EPDS Result columns were removed.Corresponding score values were retained after removing the redundant outcome columns.
  • A. Data Collection and Preprocessing: The dataset had no missing values, and entries marked as “None” were treated as valid categorical values rather than nulls.This preprocessing decision preserved “None” entries as category information.
  • A. Data Collection and Preprocessing: Four numerical features were retained, while other categorical variables were converted into numerical data using label encoding.The passage distinguishes four numerical features from the label-encoded categorical variables.

B. Stability-Aware Graph-Based Ensemble Feature Selection

SAGE integrates information gain, PCA structure, graph centrality, and bootstrap stability into an ensemble feature-selection framework. It produces robust rankings that prioritize psychological and socioeconomic predictors while reducing redundancy for downstream prediction.

  • Framework: SAGE combines statistical relevance, structural variance, and graph-based interaction information, weighting the ensemble for feature-selection stability.The framework integrates information gain, PCA, and graph centrality into a single selection method.
  • Stability and graph construction: The feature graph uses Pearson correlations with edges formed when |ρij| ≥0.5, while B = 100 bootstrap resamples produce stability-adjusted scores.Stability is computed as S∗j = µj/(σj + ϵ), with ϵ = 10^-6; normalized component scores use equal weights α = β = γ = 1/3.
  • Implementation: O(B · n · d + d^3) is SAGE’s computational complexity; for n = 766 and d = 48, it completes in under one second on standard hardware.The implementation uses Python with scikit-learn and NetworkX.
  • Feature ranking: 0.89 is the EPDS Score importance and 0.84 is the PHQ-9 Score importance, making them the two dominant predictors in the SAGE ranking.Other key factors include feeling about motherhood (0.62), occupation (0.51), education (0.50), abuse (0.44), and anger after childbirth (0.45).
  • Selected feature set and evaluation: 17 features including the target remained after removing the least important features from 766 observations.Evaluation used stratified 5-fold cross-validation, with 80% of each fold for training, 20% for testing, and 20% of training reserved for validation.

C. Model Selection and Hyperparameter Tuning

Multiple grid-searched baseline models and an ANN were evaluated, with ANN performance further improved through genetic-algorithm optimization. CTGAN oversampling addressed class imbalance within training folds while preserving untouched test sets for unbiased evaluation.

  • Model evaluation: LR, SVC, KNN, and RF were optimized via grid search alongside an ANN with two hidden layers, ReLU activation, dropout 0.3, and Adam learning rate 0.001.The ANN used 64 and 32 neurons in its hidden layers and was configured for binary classification.
  • Genetic-algorithm optimization: ANN outperformed all baseline models on the full feature set, with performance further improved through Genetic Algorithm optimization.The GA used 20 individuals over 50 generations, tournament selection of size 3, single-point crossover Pc = 0.8, and Gaussian mutation Pm = 0.2.
  • Genetic-algorithm optimization: The GA searched n ∈ {16, 32, 64, 128}, η ∈ [10^-4, 10^-2], d ∈ [0.1, 0.5], b ∈ {16, 32, 64}, and e ∈ {50, 100, 200}.Mean 5-fold CV F1-score served as the fitness function.
  • Class-imbalance handling: CTGAN oversampling was applied exclusively within each training fold, leaving the test set untouched for unbiased evaluation.CTGAN training used 300 epochs, 256 × 128 generator/discriminator hidden dimensions, batch size 500, dropout 0.3, and mode-specific normalization for categorical variables.

D. Explainable AI Integration

SAGE integrates LIME to improve interpretability through instance-level explanations of case-specific feature contributions. LIME locally approximates the complex model with an interpretable surrogate that balances prediction fidelity and interpretability.

  • Explainable AI Integration: LIME improves interpretability by explaining model predictions.It is used to provide instance-level explanations in the study.
  • Explainable AI Integration: Instance-level LIME explanations identify case-specific feature contributions.This supports examining why the model produces a prediction for an individual case.
  • Explainable AI Integration: LIME locally approximates the complex model with an interpretable surrogate model.The surrogate minimizes a loss function balancing fidelity to the original model and interpretability.

E. Evaluation Metrics

Models were evaluated using standard classification metrics derived from the confusion matrix, including accuracy, precision, recall, F1-score, and AUC-ROC.

  • Evaluation Metrics: Accuracy, precision, recall, F1-score, and AUC-ROC were computed from true positives, true negatives, false positives, and false negatives.AUC-ROC measures the ability to distinguish between classes across different decision thresholds, with values closer to 1 indicating stronger discrimination.

IV. RESULTS ANALYSIS AND DISCUSSION · A. Performance Evaluation of Baseline Models · B. Feature Selection and Oversampling for Performance Enhancement

Under stratified five-fold cross-validation, GA-ANN performed best without oversampling, while SAGE feature selection with GAN oversampling produced the strongest overall results using 16 features.

  • IV. RESULTS ANALYSIS AND DISCUSSION: Stratified 5-fold cross-validation with random seed 42 supported robust and reproducible evaluation of all models.Implementations used Python, TensorFlow 2.13, Scikit-learn 1.3, DEAP 1.4, and SDV 1.2 (CTGAN).
  • A. Performance Evaluation of Baseline Models: 79.28% accuracy, 76.52% F1-score, and 0.79 AUC made GA-ANN the best model with the full feature set and no oversampling.ANN achieved 77.67% accuracy, 74.48% F1-score, and 0.78 AUC, while RFC achieved 75.58% accuracy, 70.72% F1-score, and 0.76 AUC.
  • A. Performance Evaluation of Baseline Models: 73.43% accuracy made LR the strongest baseline model, whereas SVC and KNN showed lower performance.These results highlighted the superiority of GA-based deep learning approaches among the evaluated models.
  • B. Feature Selection and Oversampling for Performance Enhancement: 87.96% accuracy, 86.32% F1-score, and 0.88 AUC were achieved by GA-ANN with SAGE-based selection, GAN oversampling, and 16 features.SAGE-based feature selection produced the best results among the evaluated feature-selection methods.
  • B. Feature Selection and Oversampling for Performance Enhancement: 87.42% accuracy and 0.87 AUC were achieved by GA-ANN using PCA-IG with 15 features.PCA-IG also performed strongly under GAN-based oversampling.
  • B. Feature Selection and Oversampling for Performance Enhancement: GA-ANN consistently outperformed the other models, while SAGE provided the most effective feature representation under GAN oversampling.The comparison covered different feature-selection methods and their associated feature counts.
  • B. Feature Selection and Oversampling for Performance Enhancement: 86.18% accuracy and 0.86 AUC were achieved by GA-ANN using RFE + Tree-based methods with 19 features.This performance was slightly lower but remained competitive with the other feature-selection methods.

C. LIME-Based Insights into Model Decisions

LIME explanations show how SAGE-selected features drive individual GA-ANN predictions, distinguishing protective socioeconomic factors from psychosocial and lifestyle risks. SAGE also achieves 0.88 AUC with 16 features while providing stability and interpretability.

  • LIME-Based Local Explanation: 69% Class 0 probability versus 31% Class 1 reflects protective socioeconomic factors outweighing psychosocial stressors in one LIME explanation.Protective factors included relationship with in-laws (0.94), income (0.72), and age of older children (0.70); risks included poor sleep (-1.74) and abuse (-0.97).
  • LIME-Based Local Explanation: Class 1 explanations identify prior depression, education, income, PHQ-9, EPDS, and poor sleep as risk-driving features.Total children, husband’s education, and post-childbirth occupation contributed toward Class 0 in the cited LIME plot.
  • Comparative Interpretation: 0.88 AUC with only 16 features demonstrates SAGE’s competitive predictive performance relative to Zhang et al.’s 0.849 AUC using 17 predictors.SAGE also addresses selection stability through bootstrap-weighted ensemble selection, unlike the cited LASSO-Boruta study.

V. CONCLUSION AND FUTURE WORK

SAGE is presented as a stability-weighted ensemble feature-selection framework combining a GA-optimized ANN and local XAI for PPD prediction. Future work should evaluate regional and longitudinal generalization, incorporate multimodal data, and extend the system to other perinatal mental health problems and clinical systems.

  • Conclusion: SAGE combines multiple feature-selection methods with stability weighting, a GA-optimized ANN, and local XAI to identify reliable predictors from high-dimensional PPD data.The framework is designed to address class imbalance and improve accuracy over traditional methods.
  • Conclusion: Psychological and sociodemographic factors are identified as the main risk factors for postpartum depression.
  • Future Work: Future studies should test SAGE across South Asia and Africa and use longitudinal prenatal and postnatal analyses to assess temporal risk.Regional testing can assess whether similar patterns exist, while longitudinal analysis can improve understanding of temporal risk.
  • Future Work: Multimodal EHR, wearable, and clinical-text data could enhance predictive performance while retaining stability-aware feature selection.
  • Future Work: The system should be extended to other perinatal mental health problems and clinical systems.
Loading 2608.22809v1…