Population-Structure-Aware Machine Learning for Antimicrobial Resistance Prediction from Bacterial Accessory Genomes: A Benchmark on Escherichia coli Ciprofloxacin Resistance

Abstract

Antimicrobial resistance (AMR) prediction from bacterial genomic data is frequently confounded by clonal population structure, whereby machine learning models learn strain lineage rather than the causal genetic determinants of resistance. We present a benchmarking pipeline that explicitly addresses this confound through lineage-group-aware cross-validation and kinship-based population structure correction, and we evaluate it on a real-world dataset of 1,930 Escherichia coli isolates characterized by 17,198 accessory genome features, with ciprofloxacin (CIP) susceptibility as the prediction target (428 resistant, 1,502 susceptible).

Using StratifiedGroupKFold cross-validation (grouped by 337 multilocus sequence type [MLST]-based lineages) and kinship principal components (PCs) as covariates, we compared three model classes: L2-regularized logistic regression with kinship PCs, gradient-boosted trees (LightGBM) with Bayesian hyperparameter optimization (Optuna), and an FT-Transformer deep learning architecture trained with a snapshot ensemble.

Key Results

• LightGBM achieved the best mean cross-validated performance (ROC-AUC 0.913 +/- 0.050, accuracy 0.918 +/- 0.029, F1 0.695 +/- 0.185), numerically outperforming logistic regression (AUC 0.861 +/- 0.030) and the FT-Transformer (AUC 0.868 +/- 0.069, F1 0.000 across all folds, indicating a classification-threshold failure under class imbalance rather than poor discrimination).
• DeLong's test showed that LightGBM significantly outperformed logistic regression on AUC after Bonferroni correction (aggregated mean p = 0.142; one fold individually survived correction at p = 0.0062).
• Bootstrap stability selection (50 resamples per fold) identified a consistent core of 69-89 (mean approx. 78) accessory genes robustly associated with the resistance phenotype across folds.

Enforcing true lineage-level grouping, rather than isolate-level grouping, substantially reduced apparent discrimination relative to a naively evaluated pipeline and increased fold-to-fold variance – a pattern consistent with the removal of lineage leakage rather than a weakness of the method. These results indicate that gradient-boosted tree ensembles retain a statistically confirmed advantage over a regularized linear baseline for sparse, high-dimensional, moderate-sample-size bacterial genomic prediction tasks under rigorous leakage control, and yield a tractable candidate gene set for downstream experimental validation.

Code Availability

The complete Python pipeline, including PyTorch (FT-Transformer), LightGBM, Optuna, and scikit-learn implementations, is openly available on GitHub:
🔗 https://github.com/AqilaThe…