I am trying to build a ML model for diagnosing a certain group of several related but distinct diseases. Using a publicly available medical dataset, I performed data cleaning and feature extraction, and now I am trying to train the actual model. (I am using CatBoost, but I think that is not particularly relevant).
My first experiment went almost too well to believe. After training the model to classify a specific disease from the group vs. healthy individuals, using a moderately-sized training and validation set (a couple 10 000s of samples), I achieved a ROC AUC of 0.99 on the hold-out test set, with a Youden's J score of about 0.9, and TPR and TNR equally around 95%.
This kind of performance with a bit of feature engineering and a relatively simple model doesn't seem realistic for the kind of problem I am trying to solve. Medical and biological data is known to be hard to classify. However, I am pretty confident the model isn't overfit, because the hold-out set was specifically engineered to avoid information leakage, it's reasonably-sized too (a couple 1000s of instances), and I set the hyperparameters of the model so that it's less prone to overfitting. I decreased the maximal depth of the trees to a mere 4 levels, and I ended up selecting the best 8 features out of an initial set of more than 100. I also tried several different random seeds to shake up the trees a bit, and the results were very similar.
However, I then tried to fit two other models for two different diseases. The results, again, were similarly good, and exceeded my expectations. However, I then plotted a chart of feature importances (CatBoost includes a built-in feature importance calculator based on SHAP values), and I made the observation that there was a single dominant feature which was always the most important for all three models.
I also tried removing this feature altogether, which lead to a significant decrease in model performance, but then the feature which was previously at the 2nd place became the universally most important one in a very similar manner.
At this point, I thought a bit about it, and I came to the conclusion that there is probably something wrong with my methodology. If each model relies on mostly the same single feature, then they are going to be highly correlated. Therefore, if one predicts a particular disease, the other two will also predict their respective two diseases with higher probability. This ultimately renders them way less useful, because they won't tell me much about a patient, except for a very general assessment of "they probably have one or more diseases" or "they are probably healthy".
I can't, however, find the error for sure. I have two theories as to what might be wrong:
- Although the diseases are distinct, they are still related, and as such, they are not mutually exclusive. In fact, a sizeable fraction of the individuals in the training dataset are diagnosed with two or more kinds of disease. This means that I am not performing true multi-class classification, strictly speaking. However, it might still be that a true multi-class classifier would be more appropriate for solving this kind of problem.
- There is considerable imbalance in both the train and the test set; only around 10 percent of patients are positive (diagnosed with at least one disease). However, class weights are input to the model prior to training, and anyway I don't see how class imbalance in itself might cause the phenomenon I am experiencing.
Do you have an idea as to what might be going on? Also, should I continue modelling by excluding the feature in question, or is it fine to keep it?
