Machine Learning on MICE-imputed data

I'm working on a project with medical data where some of it is missing. We decided to impute the data using MICE and I found enough literature about how to choose $$m$$ (the number of imputations) and $$t$$ (the number of iterations), which variables to keep for imputation and so on.

The task is to build a classification algorithm on that data.

However, it is surprisingly hard to find anything about Multiple Imputation in a Machine Learning context (where I'm not interested in means, variances, test statistics, confidence intervals, and so on but point estimates). I read Stef van Buuren's great book and found the recommended workflow: https://stefvanbuuren.name/fimd/workflow.html I checked Rubin's "Multiple Imputation for Nonresponse in Surveys" for how the pooling should be done, which will might for regression coefficients, but only if the same features are selected on each dataset and if linear/logistic regression is used in the first place.

I can conceive different ways to handle the $$m$$ differently imputed datasets:

• Use some valid way of statistical "pooling" (how exactly)
• Build $$m$$ different models and average their predictions (and cross-validate each model on each imputed dataset, and average the CV results)
• Stack the data into one dataset (cross-validate by performing the same CV split on each imputation so that no test instances leak into the training set and then combine the imputations)

I implemented the 2nd and 3rd alternative and notice that the classification metrics differ (slightly). What is the correct way to handle the different $$m$$ datasets, so that in the end we get realistic performance estimates and one model to put into production?

If the scientific interest is restricted to point estimates and if the complete-data model is linear, this analysis of the stacked imputed data will yield unbiased estimates.

But what if my models are not linear? What if I use tree-based algorithms?

I'm glad for any answer that can help me find and justify my approach.

I've thought about this, too, recently. Your options 2 and 3 are the obvious ones I did also consider. I'd love it, if there were an easier answer I have not thought about (so I look forward to answers from others).

Alternative 2

It's clear that the option you mentioned of building one model per imputed sample (and cross-validating this) is a valid approach. Of course, one downside is that you end up with a different model for each imputed dataset (just like you can end up with a different model for each cross-validation fold, but there you can refit on the complete training data in the end, while with multiple imputations that is not obviously doable). That may or may not be an issue:

1. In principle, you can just average the predictions across these different models using the usual approaches for multiple imputation (basically a simple average on an appropriate scale, if you just want predictions, if you want prediction intervals, then you need to get into things like Rubin's rule).
2. To some extent you also get the advantages of averaging several slightly different models, so this might even improve the performance of the averaged predictions a tiny bit (a lot like the approach of re-running the same model with different random number seeds and averaging, which can improve predictions - although usually not by much - it helps much more to have very diverse well-performing models).
3. It's annoying to have so many models, if you want to describe what your model looks like/what variables are selected etc., and may slow down predictions.

One way to get out of having many models (this works much more generally including when you have an ensemble) is to distill on synthetic data (e.g. you could just run your multiple imputation and create a dataset 10x or even 100x the size of your original data, predict on this artificial data and then train a model on the predictions of the ensemble for this synthetic data). Note: When you cross-validate this part, you'd also want to do the synthetic data generation separately per fold in order to not leak information (although this kind of leak is probably a small thing, as long as the synthetic data generating / imputation model does not get to see/use the prediction target on the original data). There's some literature around this, but it's still a little bit of an experimental thing.

Alternative 3

I suspect the method of just throwing all the imputations together is theoretically problematic. I'm not entirely sure how bad this is in practice though (especially if you do your cross-validation appropriately as you describe - i.e. grouped-K-fold respecting that a patient is either in the training or validation part of a fold-split and never in both).

One thing that is "wrong" about this approach is that the correlations between imputed variables are not correctly preserved, because you don't respect which variable come from the same imputation sample. This may affect different algorithm to different degrees and I speculate that it might not be that bad.

Another thing that I think goes wrong for random forest and gradient boosted trees (like XGBoost, LightGBM etc.) is that these algorithms resample observations. Ideally, they'd respect which samples are really the same record just with different imputations and either include all these samples for a record or none. This is an easy extension in a sense, but the packages I am aware of do not offer this option. This will e.g. invalidate the out-of-bag errors for RF and can generally result in overfitting (i.e. the sub-sampling of observations will not work as well as it usually would for reducing overfitting).

There may be other things that can go wrong. The good thing is that you can evaluate through cross-validation what the effects are (i.e. nothing you are doing here invalidates your evaluation through cross-validation, I think) and if they are small, perhaps it's negligible.