How to Use IterativeImputer and KNNImputer For Powerful Missing Data Imputation | Towards Data Science

Despite the massive number of MOOCs and other online resources, there are still skill gaps in dealing with certain data problems. One example is properly dealing with missing data in real-world datasets. Beginner learners might take this problem lightly, and they are not to blame. Even though it is such a pressing issue, the complexity of data missingness problems has been greatly underestimated because of the availability of small, easy-to-work-with toy datasets.
As a result, many beginner data scientists don’t go beyond simple mean, median, or mode imputation. Though these methods may suffice for simple datasets, they are in no way a competent solution to handling missing data in large datasets.
Like any other stage of data science workflow, missing data imputation is an iterative process. You should be able to use multiple methods and compare their results effectively. While the basic methods may perform well, it is rarely the case, and you have to have a few backup strategies.
This tutorial will introduce two, more robust model-based imputation algorithms native to Sklearn — KNNImputer and IterativeImputer. You will learn their basic usage, tune their parameters, and finally, see how to compare their results visually.
The first step to implementing an effective strategy is to identify why the values are missing. Even though each case is unique, missingness can be grouped into 3 broad categories:
Missing Completely At Random (MCAR): this is a genuine case of data missing randomly. Examples are sudden mistakes in data entry, temporary sensor failures, or generally missing data that is not associated with any outside factor. The amount of missingness is low.
Missing At Random (MAR): this is a broader case of MCAR. Even though missing data may seem random at first glance, it will have some systematic relationship with the other observed features—for example — data missing from observational equipment during the times of scheduled maintenance breaks. The number of null values may vary.
Missing Not At Random (MNAR): missing values may exist in large amounts, and the reason for the missingness is associated with factors beyond our control or knowledge.
Identifying which category your problem falls into can help narrow down the set of solutions you can apply.
Let’s further explore these missingness types using the Diabetes dataset:
There are 5 features with different proportions of missing values. A first step in identifying the missingness type is to plot a missingness matrix. This special plot is available through the missingno package which we can import as msno:
This matrix shows how nulls are scattered across the dataset. White segments or lines represent where missing values lie. Glucose, BMI, and blood pressure columns can be considered MCAR because of two reasons:
The proportion of missing values is small.
Missing values are scattered completely randomly in the dataset.
However, a large proportion of data points are missing in both Insulin and SkinFoldThickness. So, is there some relationship between their missingness?
To answer this, MSNO provides a missingness heatmap that shows the missingness correlation:
>>> msno.heatmap(diabetes);
From the plot, we can see a strong correlation between Skin Thickness and Insulin. We can confirm this by sorting either of them:
>>> msno.matrix(diabetes.sort_values("Insulin"));
This plot proves that we can not understand why there are missing values in these two columns leading to a conclusion that they are MNAR. We also saw a weak correlation between blood pressure and skin thickness which disproves our first assumption. Rather than being MCAR, blood pressure is missing at random (MAR).
These missingness categories can be a lot to wrap your mind around. For a deeper insight, you can refer to my other article specifically on missingness types and the MSNO package:
towardsdev.com
Apart from the SimpleImputer, Sklearn provides KNNImputer class which uses the K-Nearest-Neighbors algorithm to impute numeric values. If you are not familiar with it, I recommend reading my separate article on it:
towardsdatascience.com
For reference, here is an excerpt from the article briefly touching on how the algorithm works:
“Imagine you have a variable with 2 categories which are visualized here:
Given a new, unknown sample, how do you tell which group it belongs to? Well, naturally, you would look at the surrounding points. But the result would be really dependent on how far you look. If you look at the closest 3 (solid circle), the green dot would belong to red triangles. If you look further, (the dashed circle) dot would be classified as a blue square.
kNN works the same way. Depending on the value of k, the algorithm classifies new samples by the majority vote of the nearest k neighbors in classification. For regression, which predicts the actual numerical value of a new sample, the algorithm takes the mean of the nearest k neighbors. “
KNNImputer is a slightly modified version of the algorithm where it tries to predict the value of numeric nullity by averaging the distances between its k nearest neighbors. For folks who have been using Sklearn for a time, its syntax should not be a problem:
With this imputer, the problem is choosing the right value for k. As you cannot use GridSearch to tune it, we can take a visual approach for comparison:
In the above plot, we compare the different KNN imputations for the SkinThickness feature using Probability Density plots. The closer the imputed distribution comes to the original, the better was the imputation. Here, it seems k=2 is the best choice.
Another more robust but more computationally expensive technique would be using IterativeImputer. It takes an arbitrary Sklearn estimator and tries to impute missing values by modeling other features as a function of features with missing values. Here is a more granular, step-by-step explanation of its functionality:
A regressor is passed to the transformer.The first feature (feature_1) with missing values is chosen.The data is split into train/test sets where the train set contains all the known values for feature_1, and the test set contains the missing samples.The regressor is fit on all the other variables as inputs and with feature_1 as an output.The regressor predicts the missing values.The transformer continues this process until all features are imputed.Steps 1–6 are called a single iteration round. These steps are carried out multiple times as specified by the max_iter parameter of the transformer.
A regressor is passed to the transformer.
The first feature (feature_1) with missing values is chosen.
The data is split into train/test sets where the train set contains all the known values for feature_1, and the test set contains the missing samples.
The regressor is fit on all the other variables as inputs and with feature_1 as an output.
The regressor predicts the missing values.
The transformer continues this process until all features are imputed.
Steps 1–6 are called a single iteration round. These steps are carried out multiple times as specified by the max_iter parameter of the transformer.
This means that IterativeImputer (II) predicts not one but max_iter number of possible values for a single missing sample. This has the benefit of treating each missing data point as a random variable and associate the inherent uncertainty that comes with missing values. This is also called multiple imputation, and it is the base for most other imputation techniques out there (yes, there are so many others).
When all iterations are done, II returns only the last result of the predictions because, through each iteration, the predictions get better and better (like gradient descent). The algorithm also has an early stopping feature which can terminate the iterations if there isn’t a considerable difference between rounds.
According to Sklearn, this implementation of IterativeImputer was inspired by the more popular R MICE package (Multivariate Imputation by Chained Equations). Let’s see it in action:
IterativeImputer is still an experimental feature, so don't forget to include the second line of the above snippet.
When estimator is set to None, the algorithm chooses it on its own. But, after reading this official guide on IterativeImputer by Sklearn, I learned that BayesianRidge and ExtraTreeRegressor yield the best results.
It is time to test how well the imputations work. To accomplish this, we will be predicting if a patient has diabetes or not (outcome), so it is a binary classification task. Let’s build feature/target arrays:
We will test both KNNImputer and IterativeImputer using cross-validation. For the estimators, we will use the recommended BayesianRidge and ExtraTreesRegressor:
We can see from the final results that KNNImputer with 7 neighbors is the best choice for the dataset when trained on RandomForests. Even though I mentioned that IterativeImputer would be more robust, you can never be sure. Maybe, we could have achieved better performance by tuning its parameters more.
Missing data is a problem that should be taken seriously. There is no point in spending hours learning complex ML algorithms if you can’t even fix the prerequisite data problems. Always remember that a model is only as good as the data it was trained on. This means you have to do your best in dealing with missing data points as they are ubiquitous in real-world datasets.
In this article, you learned how to deal with missingness using two model-based techniques: KNNImputer and IterativeImputer. Below are links to their documentation, an official Sklearn guide on their usage, and related sources that will help your understanding:
Imputation of missing values, Sklearn Guide
IterativeImputer documentation
KNNImputer documentation
Imputing missing values with variants of IterativeImputer
Thanks for reading!
How to Differentiate Between Scaling, Normalization, and Log Transformations
How to Perform Operator Overloading in Python OOP
Practical Sklearn Feature Selection in 3 stages