💻 computer science

A Distributed CatBoost Approach with Weighted Aggregation for STI/HIV Risk Prediction

This paper proposes a privacy-preserving, distributed CatBoost framework with weighted aggregation that achieves high accuracy in predicting STI and HIV risks across eight countries without requiring centralized data, thereby offering a robust alternative to traditional centralized machine learning models.

Original authors: Zhaohui Tang, Yan Li, Abdulkadir Sengur, U. Rajendra Acharya

Published 2026-07-17

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zhaohui Tang, Yan Li, Abdulkadir Sengur, U. Rajendra Acharya

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

Imagine you are trying to solve a massive, global mystery: who is at risk of catching a silent, invisible enemy like HIV or other sexually transmitted infections (STIs)? These infections are tricky because they often hide without showing any symptoms, spreading quietly until it's too late. To catch them early, doctors and scientists need to predict who might get sick before they even feel sick. Usually, to make these predictions, computers (specifically, machine learning models) are fed a giant bowl of data from millions of people all mixed together. But there's a catch: mixing everyone's private medical secrets into one giant bowl raises serious privacy concerns. It's like asking everyone to hand over their diary to a central office just to find a pattern.

This is where a clever new idea comes in: instead of mixing all the data, what if we let each country keep its own diary and train its own detective? This approach, called "distributed learning," allows computers to learn from local data without ever sending the raw secrets to a central server. The paper you are about to read explores a specific, high-tech version of this idea. It uses a smart algorithm called "CatBoost" (think of it as a super-detective that is really good at finding patterns in lists of facts) and a special method called "Leave-One-Country-Out" validation. This validation method is like a final exam where the detective is trained on seven countries but tested on the eighth one it has never seen before, proving it can actually solve the mystery in new places, not just memorize old ones.

The Detective Squad and the Privacy Puzzle

In this study, the researchers faced a big challenge: they had data from 168,459 people across eight different countries, but they wanted to predict HIV and STI risks without breaking privacy rules. Instead of building one giant, centralized model that gobbles up all the data at once (which they argue is risky and might miss local nuances), they built a "distributed" team of detectives.

Here is how their system works, using a playful analogy: Imagine eight different towns, each with its own local detective agency. Instead of sending all the town's files to a central headquarters, each town trains its own team of detectives (using the CatBoost algorithm) to spot the signs of risk based on local clues like age, education, and behavior.

But here is the twist: not all towns have the same amount of evidence, and not all detective teams are equally sharp. Some towns have thousands of files; others have fewer. Some detectives are naturally better at spotting the truth. So, the researchers didn't just let every town shout their answer at the same volume. They created a "weighted aggregation" system. Think of this as a voting system where the votes of the towns with more data and sharper detectives count for more. If a town's local model is very confident and has lots of data to back it up, its prediction gets a louder voice in the final decision. If a town's model is shaky or has very little data, its voice is quieter.

To make sure the local detectives were truly reliable, the researchers didn't just hire one detective per town. They hired a whole squad of them, each trained slightly differently (using different random seeds), and then averaged their opinions. This "local ensemble" ensures that if one detective has a bad day or misses a clue, the others in the squad catch it.

The Results: A Global Team with Local Hearts

When the researchers put this distributed team to the test, the results were surprisingly strong. They used a rigorous testing method called "Leave-One-Country-Out" (LOCO). This means they trained the system on seven countries and then asked it to predict the risks for the eighth country it had never seen. They rotated this process so every country got a turn being the "test student."

For predicting HIV, the team was incredibly accurate. On average, they achieved a score called "AUC" of 0.9850 (where 1.0 is perfect) and an accuracy of 0.9537. This means the system was right about 95% of the time. For STIs, the performance was also very high, with an average AUC of 0.9396 and accuracy of 0.9117.

However, the paper also found that not every country was a perfect match for the model. For instance, when predicting STIs, the model struggled a bit more with data from India, where the accuracy dropped to 0.7775 and the AUC to 0.8275. The researchers used a visual tool called t-SNE to look at the data and found that in some countries, the patterns of "at-risk" and "not-at-risk" people were very mixed up, making them harder to separate. In other countries, like Guinea for HIV, the separation was crystal clear, leading to near-perfect scores.

The study also looked at why the model made its decisions. They found that the "clues" that mattered most changed depending on the country. In some places, education level was the biggest predictor; in others, it was marital status or specific behaviors like condom use. This proves that a "one-size-fits-all" model might miss these local differences, whereas their distributed approach successfully captured the unique flavor of each country's data.

Why This Matters (and What It Doesn't Claim)

The main takeaway is that you can build a powerful, global health prediction tool without ever mixing everyone's private data into one big pot. The authors suggest that their method is a robust alternative to traditional centralized models, offering high accuracy while respecting privacy. They explicitly argue against the idea that you need to pool all data together to get good results, showing instead that combining smart, weighted local predictions works just as well, if not better.

However, the paper is careful not to claim this is a magic bullet for every situation. They admit that the system is more computationally demanding (it takes more computer power to train eight different teams than one big one) and that the "calibration" (how well the predicted probabilities match real-life chances) wasn't perfect for every single country. For example, while the model was great at distinguishing who was at risk, the exact probability numbers it gave for some countries (like India) were less reliable than for others (like Guinea).

Ultimately, this research suggests that by treating data like a team of local experts rather than a single giant brain, we can solve complex health puzzles with high accuracy and much greater respect for personal privacy. The authors propose this as a viable path forward for future health monitoring, provided we continue to refine how we handle the differences between diverse populations.

Technical Summary: A Distributed CatBoost Approach with Weighted Aggregation for STI/HIV Risk Prediction

Problem Statement
Sexually Transmitted Infections (STIs) and HIV remain critical global public health challenges, characterized by widespread occurrence, often asymptomatic progression, and significant social and medical implications. While machine learning (ML) has shown promise in predicting these risks using large-scale demographic and behavioral data, traditional centralized approaches face limitations. Centralized models often fail to account for the heterogeneity inherent in data distributed across different populations and countries, where class balances, feature distributions, and predictor-outcome associations vary significantly. Furthermore, centralized training raises concerns regarding user privacy and data sovereignty, particularly when dealing with sensitive health information from marginalized populations.

Methodology
The authors propose a distributed framework that utilizes the CatBoost algorithm to predict STI and HIV risks while preserving user privacy. Instead of pooling all data into a single training set, the framework operates on structured tabular data from the Demographic and Health Surveys (DHS) Program, comprising 168,459 observations across eight countries. The methodology proceeds through the following stages:

Data Preprocessing: The dataset undergoes cleaning, including the removal of missing values and highly incomplete variables. Feature selection is performed using a Random Forest (RF)-based algorithm to identify the 13 most important predictors (e.g., age, education, wealth index, condom use, sexual partner count). Robust normalization is applied to mitigate outlier effects.
Local Ensemble Training: The data is partitioned by country. Within each partition, an ensemble of multiple CatBoost models is trained using different random seeds. This local ensemble approach stabilizes predictions by averaging the probability outputs of several models, reducing variance associated with stochastic training processes.
Reliability-Aware Weighted Aggregation: To generate a global risk estimate, the framework aggregates local predictions using a weighted strategy rather than a simple average. Weights are calculated based on two factors: the number of training samples ( $n_k$ ) and the validation quality score ( $q_k$ , measured by AUC) of each partition. The unnormalized weight is defined as $w_k = n_k \cdot q_k^2$ . This ensures that partitions with larger datasets and higher predictive performance contribute more significantly to the final global probability.
Threshold Calibration and Uncertainty Estimation: The final classification does not rely on a fixed threshold (e.g., 0.5). Instead, an optimal threshold is selected via validation to maximize the F1-score, balancing sensitivity and specificity. Additionally, a predictive uncertainty module calculates the standard deviation of the probability estimates across the ensemble members, providing a confidence measure for the final prediction.
Evaluation Strategy: The model is evaluated using a Leave-One-Country-Out (LOCO) cross-validation strategy. In each fold, one country serves as the test set while the remaining seven countries' data are used to train the distributed model.

Key Contributions
The paper outlines four primary contributions:

Distributed Ensemble Architecture: The framework employs independent CatBoost ensembles within each data partition, aggregating results via probability averaging rather than merging datasets.
Partition-Specific Stabilization: It utilizes multiple CatBoost models with different seeds per partition to create robust local predictors.
Validation-Based Weighted Aggregation: A novel technique assigns weights to partitions based on data volume and validation performance, enhancing the reliability of the global prediction.
Enhanced Decision Robustness: The integration of threshold calibration and uncertainty estimation allows for adaptive decision rules and interpretable confidence metrics, addressing the specific needs of imbalanced medical datasets.

Results
The proposed framework was tested on HIV and STI prediction tasks using the LOCO methodology.

HIV Prediction: The model achieved an average Area Under the Curve (AUC) of 0.9850 and an average accuracy of 0.9537. Performance varied by country, with Guinea achieving the highest AUC (0.9929) and Ethiopia the lowest (0.9676), though all remained high.
STI Prediction: The framework obtained an average AUC of 0.9396 and an average accuracy of 0.9117. Haiti and the Dominican Republic demonstrated the strongest performance (AUC > 0.98), while India presented the most challenging case with an AUC of 0.8275, attributed to significant class overlap in the feature space as revealed by t-SNE analysis.
Comparative Performance: When compared to existing literature, including centralized models and a Federated Random Forest approach, the proposed distributed CatBoost framework demonstrated superior performance, achieving a maximum AUC of 0.99 and 95% accuracy. It outperformed previous centralized studies (which reported AUCs ranging from 0.73 to 0.89) and the Federated Random Forest approach (AUC 0.97).

Significance and Claims
The authors claim that this framework offers a viable and robust solution for STI/HIV risk prediction in distributed environments where user privacy is a primary concern. The study demonstrates that preserving local predictive structures through independent training, followed by reliability-weighted aggregation, can yield high predictability without the need for centralized data pooling.

The paper emphasizes that the approach is particularly suitable for heterogeneous health datasets where different populations exhibit varying risk factors. By incorporating uncertainty estimation and threshold calibration, the model provides not only high accuracy but also interpretable confidence levels, which are crucial for clinical decision-making and intervention strategies. The authors acknowledge that while the method is computationally more demanding than a single centralized model and showed some variation in calibration across countries, it successfully balances privacy preservation with high predictive performance, serving as a strong alternative to traditional centralized ML models.

The Detective Squad and the Privacy Puzzle

The Results: A Global Team with Local Hearts

Why This Matters (and What It Doesn't Claim)

More like this