TB-Bench: A Systematic Benchmark of Machine Learning and Deep Learning Methods for Second-Line TB Drug Resistance Prediction

This study presents TB-Bench, a systematic benchmark evaluating 20 machine learning and deep learning models across 14 second-line tuberculosis drugs on a large WHO dataset, revealing that traditional machine learning methods often outperform deep learning in internal tests while both struggle with cross-dataset generalization compared to catalogue-based approaches.

The Gist

Imagine you are a doctor trying to treat a patient with Tuberculosis (TB). In the past, you had to wait weeks for a lab to grow the bacteria in a dish to see which antibiotics would kill it. Today, we can read the bacteria's "instruction manual" (its DNA) in a day. The big question is: Can we use Artificial Intelligence (AI) to read that manual and instantly tell us which drugs will work?

This paper, TB-Bench, is like a massive "stress test" for 20 different AI detectives trying to solve this mystery. Here is the story of what they found, explained simply.

1. The Problem: The "Second-Line" Mystery

Think of TB treatment like a two-tiered defense system.

• First-Line Drugs: These are the standard, heavy-duty weapons. Most bacteria are scared of them.
• Second-Line Drugs: These are the special forces used when the bacteria have learned to dodge the first-line weapons. They are more expensive, have more side effects, and are harder to use.

The problem is that predicting resistance to these "Special Forces" drugs is much harder than predicting resistance to the standard ones. The AI models are good at the basics, but they stumble when the game gets complicated.

2. The Experiment: The Great AI Showdown

The researchers gathered a massive library of 50,000+ bacterial DNA samples from around the world (the "WHO Dataset"). They set up a competition with 20 different AI models:

• The "Simple" Team: Traditional Machine Learning (like XGBoost, Logistic Regression). Think of these as experienced, old-school detectives who rely on clear, logical rules and simple checklists.
• The "Complex" Team: Deep Learning (Neural Networks). Think of these as genius, futuristic AI that can find hidden patterns in massive amounts of data, like a supercomputer trying to solve a puzzle with billions of pieces.

They tested these models on 14 different second-line drugs using three different ways of looking at the DNA:

1. The Whole Genome: Reading the entire instruction manual.
2. The Coding Regions: Reading only the chapters that make proteins.
3. The "Cheat Sheet": Reading only the specific pages known to contain resistance clues.

3. The Results: The Underdog Wins (Mostly)

Here is the twist: The simple detectives often beat the supercomputers.

• The Surprise: The "Simple" models (especially one called XGBoost) were the champions. They were faster, easier to understand, and often more accurate than the complex Deep Learning models.
• The Metaphor: Imagine trying to find a specific typo in a book. The Deep Learning model tries to analyze the font, the paper texture, the author's mood, and the history of the printing press. The Simple model just looks at the letters. In this case, looking at the letters was enough. The complex models were "overthinking" it.
• The "Cheat Sheet" Works: Surprisingly, the models didn't need to read the entire DNA manual. When they only looked at the specific "cheat sheet" (known resistance genes), they performed just as well as when they read the whole book. This is great news because it means we don't need massive computing power to get good results.

4. The Big Catch: The "Foreign Student" Problem

The models did great when tested on the data they were trained on (the WHO dataset). But when the researchers tested them on a completely new set of data from a different country (China), the scores dropped significantly.

• The Analogy: Imagine a student who memorizes the answers to a specific practice test perfectly. They get 100%. But when you give them a real exam with slightly different questions, they fail.
• Why? The AI learned the "accent" of the specific groups of bacteria in the training data, rather than the universal rules of resistance. If the training data came mostly from one region, the AI gets confused when it sees bacteria from another region.

5. The Verdict: Don't Throw Out the Old Books

The study compared these AI models against TBProfiler, a tool that simply checks a pre-written list of known mutations (like a dictionary).

• The Result: The AI models were sometimes better, but often, the simple "dictionary" approach was just as good, if not better.
• The Lesson: While AI is powerful, we can't just throw away human expertise and curated lists of known mutations. For now, the "dictionary" is still a very reliable doctor.

Summary: What Does This Mean for the Future?

1. Keep it Simple: You don't always need a super-complex AI to predict drug resistance. Simple, fast models work great and are easier to use in hospitals with limited computers.
2. Data Diversity is Key: To make these AI doctors truly reliable, we need to teach them with bacteria from everywhere in the world, not just a few places. Otherwise, they will be biased and fail when they travel.
3. The "Cheat Sheet" is Enough: We don't need to analyze the entire genome to find the answer; focusing on the known trouble spots is efficient and effective.

In a nutshell: TB-Bench tells us that while AI is a powerful tool for fighting drug-resistant TB, the best results come from combining smart, simple algorithms with diverse, global data, rather than just relying on the most complex technology available.

Technical Summary

1. Problem Statement

Drug-resistant tuberculosis (TB), particularly Multi-Drug Resistant (MDR) and Extensively Drug-Resistant (XDR) strains, poses a critical global health challenge. While Whole-Genome Sequencing (WGS) has enabled the development of Machine Learning (ML) and Deep Learning (DL) models to predict drug resistance, significant gaps remain:

• Performance Disparity: Current methods reliably predict resistance to first-line drugs but show inconsistent and lower performance for second-line drugs (e.g., Bedaquiline, Linezolid, Fluoroquinolones) compared to traditional culture-based phenotypic testing.
• Generalization Issues: Many models suffer from overfitting and fail to generalize across different datasets or geographic populations.
• Lack of Standardization: There is no unified framework to compare diverse ML/DL architectures, feature representations, and datasets, making it difficult to determine the optimal approach for clinical deployment.

2. Methodology

The authors developed TB-Bench, a comprehensive benchmarking framework to evaluate 20 distinct ML and DL models across 14 second-line drugs.

Datasets

• Primary Dataset (Training/Internal Validation): Derived from the WHO global collection (2nd edition mutation catalogue), comprising 50,801 samples. After quality control (QC) and filtering, 49,266 samples were used.
• External Validation Dataset: An independent cohort from China comprising 1,199 samples (resistant to at least one of 7 second-line drugs) to assess cross-dataset generalizability.
• Target Drugs: 14 second-line drugs including Amikacin (AMK), Bedaquiline (BDQ), Capreomycin (CAP), Cycloserine (CYC), Ethionamide (ETO), Kanamycin (KAN), Levofloxacin (LFX), Linezolid (LZD), Moxifloxacin (MFX), Ofloxacin (OFX), Para-aminosalicylic acid (PAS), and others.

Feature Representations

To account for methodological heterogeneity, models were evaluated using three distinct input feature sets:

1. Whole-Genome Variants: All SNPs and INDELs across the entire genome.
2. Coding-Region Variants: Variants restricted to coding regions.
3. Targeted Gene Variants: Variants from 73 curated Tier 1 and Tier 2 resistance-associated genes (similar to TBProfiler).

Models Evaluated

The study benchmarked 20 models categorized into:

• Traditional ML (14 models): Linear models (Logistic Regression with L1/L2/OHE), Tree-based methods (Decision Tree, Random Forest, XGBoost, Treesist), SVMs (Linear, RBF, PRPS-based), Bayesian (Bernoulli Naive Bayes), and Shallow ANN.
• Deep Learning (6 models): CNN (Label-Encoded), DeepAMR (Denoising Autoencoder), Wide & Deep Neural Networks (WDNN), Deep NN, and Single-Drug CNN (SDCNN).
• Baseline: TBProfiler (catalogue-based method) was included as a non-ML clinical standard.

Experimental Pipeline

• Preprocessing: Standardized pipeline using fastp, BWA, Picard, freebayes, and SnpEff. Variants were filtered for depth (DP ≥ 10) and allele frequency (AF ≥ 0.75).
• Splitting: 80/20 stratified split (Training/Validation vs. Held-out Test) to preserve class imbalance.
• Evaluation Metrics: Primary metric was PRAUC (Area Under the Precision-Recall Curve) due to class imbalance; secondary metrics included F1-score and Youden's J statistic for threshold optimization.
• Analysis: SHAP (SHapley Additive exPlanations) was used for feature importance analysis, and Jaccard similarity was computed to analyze cross-resistance patterns.

3. Key Contributions

1. Systematic Benchmarking: The first large-scale, standardized comparison of 20 ML/DL models for 14 second-line TB drugs using a unified framework.
2. Feature Set Analysis: Demonstrated that biologically curated, smaller feature sets (Tier 1 & 2 genes) perform comparably to whole-genome variants, reducing computational complexity without sacrificing accuracy.
3. Generalization Assessment: Rigorously tested models on an external dataset, revealing significant performance drops that highlight current limitations in cross-dataset generalization.
4. Open Source Framework: Released the TB-Bench source code and pipelines on GitHub to facilitate reproducibility and future method development.

4. Key Results

Model Performance (Internal Validation)

• Simple ML > Complex DL: Traditional ML models, particularly XGBoost and Logistic Regression (LR), consistently outperformed complex Deep Learning models.
  • XGBoost achieved the highest PRAUC scores (46%–93%) for 10 of the 14 drugs.
  • DL models (e.g., WDNN, CNN) showed comparable or slightly lower performance, suggesting that the signal for resistance in TB is dominated by high-effect variants rather than complex non-linear interactions.
• Feature Efficiency: Models trained on Tier 1 & 2 gene variants (approx. 19x fewer features than whole-genome) achieved performance comparable to whole-genome models. In some cases (e.g., Linezolid), curated features even outperformed whole-genome inputs.
• Class Imbalance Impact: Performance was heavily correlated with the prevalence of resistant isolates. Drugs with <10% resistance (e.g., Cycloserine, Linezolid, PAS) showed poor PRAUC (<60%) across all models.

Generalization (External Validation)

• Significant Performance Drop: When tested on the independent Chinese dataset, all models (both ML and DL) showed a substantial decline in performance (PRAUC < 75% for all drugs).
• No Model Superiority: Neither ML nor DL models significantly outperformed the random baseline for drugs like BDQ and LZD in the external set.
• Catalogue Baseline: TBProfiler performed comparably to, or better than, learned models on the external dataset for several drugs, reinforcing the value of expert-curated mutation catalogues.

Biological Insights

• Cross-Resistance: Models often identified variants associated with resistance to other drugs, indicating they learn shared markers of multi-drug resistance (MDR) rather than drug-specific causal mechanisms.
• Sampling Bias: Metadata analysis revealed that poor generalization for BDQ and LZD was likely due to training data originating from only 1–2 specific projects, leading to overfitting on study-specific or geographic-specific features rather than global biological determinants.

5. Significance and Conclusion

• Clinical Utility: The study suggests that simple, interpretable ML models (like XGBoost) trained on curated gene sets are currently the most practical approach for second-line drug resistance prediction, especially in low-resource settings where computational efficiency is key.
• Data Quality over Model Complexity: The results indicate that increasing model complexity (DL) does not yield better results if the underlying data lacks geographic and genetic diversity. The primary bottleneck is data generalizability, not algorithmic sophistication.
• Future Directions: To improve clinical utility, future efforts must focus on:
  • Expanding WGS cohorts to include diverse geographic and genetic backgrounds.
  • Integrating non-coding regulatory elements and other omics data (e.g., transcriptomics).
  • Moving beyond binary classification to quantitative prediction of Minimum Inhibitory Concentration (MIC).

In summary, TB-Bench establishes that while ML/DL offers promise, current models for second-line TB drugs are limited by data bias and class imbalance. Simple, well-validated models often rival complex deep learning architectures, and expert-curated catalogues remain a robust clinical standard.