💻 computer science

TITAN: A Diversity-Driven Stacked Ensemble with a Multi-Branch Neural Network for Predicting Live Birth in IVF

The paper introduces TITAN, a diversity-driven stacked ensemble model that integrates five tree-based learners and a custom multi-branch neural network with rigorous leakage detection and Bayesian optimization to achieve state-of-the-art accuracy in predicting live birth outcomes for IVF cycles using the HFEA registry.

Original authors: Bala Shiwangi¹, Shailesh D. Kamble², Brijesh Kumar¹

Published 2026-08-26

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Bala Shiwangi¹, Shailesh D. Kamble², Brijesh Kumar¹

Original paper licensed under CC BY 4.0 (https://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the quiet, high-stakes world of reproductive medicine, the journey to parenthood often begins with a single, difficult question: will this treatment cycle result in a live birth? For couples facing infertility, in vitro fertilization, or IVF, is a primary path forward, yet even after decades of medical refinement, clinics rarely see success rates exceed forty percent per attempt. The ability to predict the outcome of these cycles with accuracy would be transformative, allowing doctors to tailor treatments to individual needs and offering families clearer, more honest guidance before they begin. However, building a computer program to make these predictions is fraught with a specific, subtle danger known as data leakage. In medical records, information is often recorded in the order events happen, but sometimes data about the final result—like the weight of a baby or the date of delivery—is accidentally included in the dataset used to train the prediction model. If a computer learns from these future facts, it can appear to be a genius, achieving near-perfect scores that vanish the moment it is asked to predict a new, unseen case. The challenge, therefore, is not just to find the right patterns in the data, but to rigorously strip away any information that the doctor would not actually know at the time of the decision.

A team of researchers has addressed this challenge with a new system called TITAN, designed to predict live birth outcomes using only the information available before an embryo is transferred. The researchers started with a massive database from the United Kingdom containing records of nearly 170,000 IVF cycles. Their first and most critical step was to act as a filter, automatically scanning the data to identify and remove any columns that contained "leaky" information—details about the pregnancy or birth that would only be known after the treatment was finished. By using a smart algorithm to detect these hidden clues, they removed twenty-two variables that would have otherwise tricked the model, leaving behind thirty-eight genuine predictors such as patient age, the number of eggs retrieved, and the type of treatment used. This process ensured that the system was learning from reality, not from a shortcut.

Once the data was cleaned, the team did not rely on a single type of computer model to make the prediction. Instead, they built a diverse team of six different "learners," each with its own way of thinking about the problem. Five of these were tree-based models, which work by asking a series of yes-or-no questions to sort patients into groups, while the sixth was a complex neural network designed to mimic the way a brain processes information. To ensure these models did not all make the same mistakes, the researchers trained them using a technique called out-of-fold stacking. Imagine asking a group of experts to review a case, but instead of letting them all see the same full file at once, you give each expert a slightly different version of the evidence and ask them to predict the outcome based on what they see. This method forces each model to learn from different angles, creating a set of predictions that are complementary rather than identical.

The true innovation of TITAN lies in how it combines these six different opinions. Rather than simply taking an average of their answers, the system takes the specific confidence levels, margins of error, and uncertainty scores from each model and feeds them into a final, master model. This master model, which was fine-tuned using a sophisticated search process to find the perfect settings, acts as a judge that weighs the unique strengths of each expert. The researchers also introduced a custom-designed neural network with five parallel pathways, each structured differently to catch patterns that the others might miss. This network was trained using a method that averages its own learning steps over time, helping it to settle into a stable and reliable state rather than getting stuck in a local optimum.

When tested on a separate group of patients that the system had never seen before, TITAN achieved an accuracy of 89.19 percent, correctly predicting the outcome in nearly nine out of ten cases. This performance was significantly better than any single model used on its own and far superior to a simple guess based on the most common outcome. The system also achieved a score of 0.72 on a standard measure of how well it can distinguish between success and failure, a result that the researchers confirmed was not inflated by hidden data. In a series of tests where they removed one part of the system at a time, they found that the initial step of removing the "leaky" data was the most important factor; without it, the system's performance collapsed to the level of a random guess. The study also revealed that while the five tree-based models were quite similar in how they thought, the inclusion of the neural network added a unique perspective that improved the final result.

The findings suggest that while predicting the success of IVF is inherently difficult due to the complex biology involved, a carefully constructed system that respects the timeline of medical data can provide highly reliable forecasts. The researchers noted that the system was particularly good at identifying cycles that would not result in a live birth, correctly flagging the vast majority of unsuccessful attempts. However, predicting the successful outcomes remained more challenging, a limitation the authors attribute to the absence of certain biological details, such as the specific quality of the embryo, which are not always recorded in the initial registry data. Ultimately, TITAN offers a practical, transparent, and robust tool for clinicians, demonstrating that by combining diverse methods and rigorously cleaning the data, it is possible to build a system that offers genuine insight into the future of a treatment cycle without relying on information that does not yet exist.

Technical Summary: TITAN – A Diversity-Driven Stacked Ensemble for IVF Live Birth Prediction

Problem Statement
Predicting live birth outcomes in In Vitro Fertilization (IVF) cycles is a critical yet statistically challenging task in reproductive medicine. Existing datasets, such as the UK's Human Fertilisation and Embryology Authority (HFEA) registry, are high-dimensional but suffer from three primary methodological barriers: (1) the presence of data leakage, where post-treatment variables (e.g., birth weight, delivery date) are inadvertently included in pre-treatment models, artificially inflating performance; (2) high-dimensional, heterogeneous clinical data; and (3) severe class imbalance (approximately 3.13:1 negative to positive cases). Naïve models often fail to generalize in prospective clinical settings because they learn from information unavailable at the time of prediction.

Methodology: The TITAN Pipeline
The authors propose TITAN (Total Intelligence Through Aggregated Networks), a seven-stage end-to-end processing pipeline designed to eliminate leakage and maximize prediction diversity.

Adaptive Leakage Detection: Instead of manual review, TITAN employs an adaptive thresholding algorithm. It trains single-feature decision tree surrogates to score the predictive power of each feature. Features exceeding an adaptive threshold ( $\tau^* = 0.90$ ) are identified as "leakers" (post-outcome indicators) and removed. This process resulted in the removal of 22 post-outcome variables, retaining 38 valid pre-treatment predictors. The analysis noted that while 7 features were initially flagged as highly predictive (leaky), the adaptive thresholding strategy was necessary to remove the broader set of 22 variables that would otherwise cause data leakage.
Preprocessing & Feature Engineering:
- Missing Values: A tiered strategy applies median imputation for skewed numerical data, mean for normal data, and mode/'Unknown' for categorical data.
- Encoding: Categorical variables undergo target encoding with Laplace smoothing ( $\lambda=10$ ) to prevent leakage between folds.
- Engineering: Statistical aggregates, polynomial transformations, and pairwise interactions are generated from the top-ranked features.
- Pruning: SHAP (SHapley Additive exPlanations) values are used to prune features contributing less than 1% of the maximum importance, reducing noise.
- Scaling: RobustScaler is applied to mitigate outlier sensitivity.
Diversity-First Stacked Ensemble:
- Base Learners: Five heterogeneous tree-based models (XGBoost, LightGBM, CatBoost, Random Forest, ExtraTrees) are trained using 5-fold Out-of-Fold (OOF) stacking.
- Neural Network Component: A custom five-branch deep neural network serves as the sixth base learner. Its branches include Deep Dense, Double Residual, Wide+Attention, Hadamard Interaction, and MLP-Mixer architectures. This network is trained using Stochastic Weight Averaging (SWA) and Cosine Warm Restarts to ensure stability and distinct gradient-based inductive bias.
- Meta-Learner: The OOF predictions from all six base models are combined into a rich 21-dimensional meta-feature matrix. This matrix includes individual prediction probabilities, margins, Shannon entropy, and cross-model statistics (mean, standard deviation, range).
- Final Model: A Bayesian-tuned XGBoost model acts as the meta-learner, trained on the 21D feature set.
Optimization: Hyperparameters for the tree-based models are optimized using Optuna (TPE sampler), and the final decision threshold is optimized on a held-out validation set to balance sensitivity and specificity.

Key Contributions

Adaptive Leakage Detection: A principled, automated method to detect and remove post-outcome features without manual intervention, ensuring the model relies solely on pre-treatment data.
Diversity-First Ensemble Design: The integration of five distinct tree algorithms with a custom multi-branch neural network to maximize prediction diversity and minimize correlated errors.
Novel Neural Architecture: A five-branch neural network incorporating SWA and cosine annealing, designed to provide a gradient-based signal distinct from tree-based learners.
Rich Meta-Feature Representation: The construction of a 21-dimensional meta-feature set that captures not just probabilities, but also model confidence, margin, and entropy.
Comprehensive Ablation: A systematic study quantifying the contribution of each pipeline component, demonstrating that leakage detection and Bayesian optimization are the most critical factors.

Results
Evaluated on the HFEA registry ( $N=169,616$ cycles) with a held-out test set of 25,428 cycles:

Performance: TITAN achieved 89.19% accuracy and an AUC-ROC of 0.7202 (based on the clean, pre-treatment feature set).
Comparison: TITAN outperformed the majority-class baseline (75.81% accuracy) by 13.38 percentage points and surpassed all individual tree-based constituent models. However, Deep Learning baselines (DNN, ResNet, CNN) achieved higher AUC scores (0.889) than TITAN's 0.7202. The authors clarify that the DNN baselines achieved these higher AUCs because they were trained on a feature set that included label-encoded categorical variables (creating numerical artifacts) and potentially retained more features, whereas TITAN's lower AUC reflects a more conservative, clinically valid evaluation using strictly pre-treatment features that were SHAP-pruned and target-encoded.
Ablation Insights:
- Removing leakage detection caused the model to collapse to a dummy classifier (AUC 0.500), confirming the dataset's heavy reliance on leaked features.
- Replacing OOF stacking with simple averaging resulted in an 11.39 percentage point drop in accuracy.
- Bayesian optimization via Optuna contributed a 1.07 percentage point gain over default parameters.
Class Performance: The model demonstrated high specificity for the negative class (No Live Birth: 90% recall) but conservative sensitivity for the positive class (Live Birth: 24% recall). The authors attribute this to the absence of key post-treatment biological predictors (e.g., embryo morphokinetics) in the registry data, representing an irreducible uncertainty in pre-treatment prediction.

Significance and Claims
The paper claims that TITAN offers a practical, interpretable, and extensible framework for predicting ART outcomes. Its primary significance lies in addressing the "data leakage" problem that plagues many clinical machine learning studies, ensuring that reported performance metrics reflect true prospective capability rather than data artifacts. By prioritizing diversity in ensemble design and rigorously cleaning the data, TITAN provides a robust pathway for clinical deployment. The authors note that while the model is highly accurate for predicting unsuccessful cycles, its ability to predict live births is inherently limited by the lack of biological data available prior to embryo transfer, a constraint acknowledged as a fundamental challenge in the field. Future work is suggested to include multimodal integration with embryo imaging and prospective validation on non-UK registries.

More like this