Five Ways to run the Advanced Imputation
A Simple Habit That Transformed My Data Cleaning Workflow.
The traditional imputation methods, imputation with the mean, median, default value, or mode, have some downsides.
First, it’s a data corruption. Also, it changes the distribution. It specifically matters for cases when you will use the cleaned dataset for machine learning later.
Today I will show you the advanced methods you can use to impute the missing values.
Types of missing data:
There are three types of missing data.
MAR (Missing at random) means that the probability depends on observed values but not on the missing value itself.
MNAR (missing not at random) means that the probability depends on the missing value itself or on unobserved factors.
MCAR (missing completely at random). That means that the probability of value being missing does not depend on any observed or unobserved value.
Depending on the type of missing data in your dataset, you will decide which advanced imputation method to use.
Now, let’s look at the methods.
Model-based imputation:
In this article, I cover mostly machine learning methods. The principles I will describe will apply to most of them.
When you train a machine learning model, you need to understand what are the missing data types. If it’s numeric, you will use the regression model, and for categorical, you will use the classification model.
Another important factor is how many values are missing, because if a column has more than 50% missing data, it is not useful for training. In some of the cases you just need to drop that column.
You cannot just apply machine learning methods everywhere and ignore the other factors. Otherwise, you just introduce the noise.
Here is the imputation pipeline:
Copy the data set. Do not word directly with the original. Use df.copy().
Flag the missing data with a new column.
Transform categorical data into integer codes. Use factorization or frequency encoding methods.
Then split the data set into train and test.
Choose the machine learning method depending on what kind of missing data you have and other factors.
Train the model.
Use the model to predict missing values.
Impute the missing data.
Always check the distribution after you see the results. Compare the before and after.
MICE:
One of the most popular imputation methods for the missing data is MICE (multiple imputation by chained equations).
This method works well if values are missing at random (MAR) and missing not at random (MNAR). It’s accurate but slow with large data sets.
Here are the most popular libraries for that method:
miceforest
statsmodels
fancyimpute
scikit-learn
Example:
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.ensemble import RandomForestRegressor
imputer = IterativeImputer(
estimator=RandomForestRegressor(n_estimators=100, random_state=42),
max_iter=10,
random_state=42
)
X_imputed = imputer.fit_transform(X)Random Forest:
Another good alternative to MICE is random forest imputation. It operates on a similar principle as MICE, and it’s good for missing at random values (MAR).
The random forest is a non-parametric method.
Example: IterativeImputer with Random Forest
import numpy as np
import pandas as pd
from sklearn.experimental import enable_iterative_imputer
from sklearn.impute import IterativeImputer
from sklearn.ensemble import RandomForestRegressor
# Create sample data with missing values
df = pd.DataFrame({
‘A’: [2, 3, np.nan, 5],
‘B’: [np.nan, 9, 16, 25]
})
# Configure IterativeImputer with RandomForest
imputer = IterativeImputer(estimator=RandomForestRegressor(), random_state=42)
df_imputed = pd.DataFrame(imputer.fit_transform(df), columns=df.columns)
print(df_imputed)Example: Missforest
import numpy as np
import pandas as pd
from missingpy import MissForest
# Example dataset with missing values
data = pd.DataFrame({
‘Age’: [25, np.nan, 30, 22, np.nan],
‘Salary’: [50000, 60000, np.nan, 52000, 58000],
‘Experience’: [2, 5, 7, np.nan, 3]
})
print(”Original data with missing values:”)
print(data)
# Initialize MissForest imputer
imputer = MissForest(max_iter=10, n_estimators=100, random_state=42)
# Perform imputation
imputed_data = imputer.fit_transform(data)
# Convert back to DataFrame
imputed_df = pd.DataFrame(imputed_data, columns=data.columns)
print(”\nData after MissForest imputation:”)
print(imputed_df)KNN Imputation:
Use this machine learning method when the data is too noisy.
KNN is a nonparametric method, so it makes no assumptions about the data distribution. It can work well for complex, non-normal, or nonlinear datasets, but it can oversmooth the distribution and lose information.
Use the fancyimpute and scikit-learn libraries.
Example:
from sklearn.impute import KNNImputer
from sklearn.preprocessing import StandardScaler
# KNN is sensitive to feature scale
scaler = StandardScaler()
df_scaled = scaler.fit_transform(df)
imputer = KNNImputer(n_neighbors=5, weights=’distance’)
df_imputed = scaler.inverse_transform(
imputer.fit_transform(df_scaled)
)Deep learning methods:
Generative Adversarial Imputation Networks (GAIN)
A Generative Adversarial Network (GAN) is a deep learning model that generates realistic synthetic data by putting two neural networks against each other in a competitive training process: the generator and the discriminator.
Generative Adversarial Imputation Networks (GAIN) are specialized extensions of GANs designed for missing data imputation. The generator fills in missing values, and the discriminator learns how to distinguish observed entries from imputed ones.
GAIN imputations preserve the statistical properties of the dataset better than traditional methods.
DMGAN:
DMGAN (Dynamic Multiple GAN) is a deep learning–based method for missing data imputation that uses multiple generative adversarial networks (GANs) to capture both sample distribution and attribute characteristics.
It is effective for high-dimensional and complex datasets where missingness patterns vary.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# Generator
class Generator(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(Generator, self).__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim),
nn.Sigmoid()
)
def forward(self, x):
return self.model(x)
# Discriminator
class Discriminator(nn.Module):
def __init__(self, input_dim, hidden_dim):
super(Discriminator, self).__init__()
self.model = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, 1),
nn.Sigmoid()
)
def forward(self, x):
return self.model(x)
# Example setup
input_dim = 10 # number of features
hidden_dim = 64
generator = Generator(input_dim, hidden_dim, input_dim)
discriminator = Discriminator(input_dim, hidden_dim)
criterion = nn.BCELoss()
optimizer_G = optim.Adam(generator.parameters(), lr=0.001)
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.001)
# Dummy batch with missing values replaced by 0
real_data = torch.rand((32, input_dim))
mask = (torch.rand_like(real_data) > 0.2).float() # 20% missing
incomplete_data = real_data * mask
# Training loop (simplified)
for epoch in range(100):
# Generate imputed data
gen_data = generator(incomplete_data)
filled_data = incomplete_data + (1 - mask) * gen_data
# Train discriminator
optimizer_D.zero_grad()
real_labels = torch.ones((32, 1))
fake_labels = torch.zeros((32, 1))
loss_D = criterion(discriminator(real_data), real_labels) + \
criterion(discriminator(filled_data.detach()), fake_labels)
loss_D.backward()
optimizer_D.step()
# Train generator
optimizer_G.zero_grad()
loss_G = criterion(discriminator(filled_data), real_labels)
loss_G.backward()
optimizer_G.step()
print(”Training complete. Imputed values generated.”)Diffusion models:
Autoencoders:
An autoencoder is a neural network designed to learn efficient representations of data. It compresses inputs into a lower-dimensional space and then reconstructs them back.
It consists of an encoder that maps input data to a latent representation and a decoder that reconstructs the original input from this compressed form.
Example:
import numpy as np
import tensorflow as tf
from tensorflow import keras
MISSING = -1.0
def masked_mse(y_true, y_pred):
mask = tf.cast(tf.not_equal(y_true, MISSING), tf.float32)
squared = tf.square((y_true - y_pred) * mask)
return tf.reduce_sum(squared) / tf.reduce_sum(mask)
input_dim = 8
model = keras.Sequential([
keras.layers.Input(shape=(input_dim,)),
keras.layers.Dense(32, activation=”relu”),
keras.layers.Dense(16, activation=”relu”),
keras.layers.Dense(32, activation=”relu”),
keras.layers.Dense(input_dim),
])
model.compile(optimizer=”adam”, loss=masked_mse)
# X_train is complete; during inference, missing entries are set to MISSING.
# model.predict(X_incomplete) returns the full reconstruction.Variational Autoencoders (VAE):
A Variational Autoencoder (VAE) is a generative model that learns a probabilistic latent space.
VAEs encode data into distributions (mean and variance) instead of fixed points in standard autoencoders.
The encoder maps incomplete data into a latent distribution, and then the decoder reconstructs missing values by sampling from it.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# VAE definition
class VAE(nn.Module):
def __init__(self, input_dim, hidden_dim, latent_dim):
super(VAE, self).__init__()
# Encoder
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc_mu = nn.Linear(hidden_dim, latent_dim)
self.fc_logvar = nn.Linear(hidden_dim, latent_dim)
# Decoder
self.fc2 = nn.Linear(latent_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, input_dim)
def encode(self, x):
h = torch.relu(self.fc1(x))
return self.fc_mu(h), self.fc_logvar(h)
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z):
h = torch.relu(self.fc2(z))
return torch.sigmoid(self.fc3(h))
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
return self.decode(z), mu, logvar
# Setup
input_dim, hidden_dim, latent_dim = 10, 64, 16
model = VAE(input_dim, hidden_dim, latent_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)
# Loss function (reconstruction + KL divergence)
def loss_fn(recon_x, x, mu, logvar, mask):
BCE = nn.MSELoss()(recon_x * mask, x * mask) # only compare observed entries
KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
return BCE + KLD
# Dummy data with missing values
data = torch.rand((32, input_dim))
mask = (torch.rand_like(data) > 0.2).float() # 20% missing
incomplete = data * mask
# Training loop
for epoch in range(100):
recon, mu, logvar = model(incomplete)
loss = loss_fn(recon, data, mu, logvar, mask)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Imputation
imputed_data = incomplete + (1 - mask) * model(incomplete)[0].detach()
print(”Imputed data:\n”, imputed_data)Denoising autoencoder imputation:
This method learns the underlying structure of the data.
The encoder maps the data to the higher-dimensional space, and then it runs the transformation. Then it uses a decoder to map it back to the original dimension and remove noise.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# Autoencoder definition
class DenoisingAutoencoder(nn.Module):
def __init__(self, input_dim, hidden_dim):
super(DenoisingAutoencoder, self).__init__()
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.Linear(hidden_dim, input_dim),
nn.Sigmoid()
)
def forward(self, x):
return self.decoder(self.encoder(x))
# Setup
input_dim = 10
hidden_dim = 64
model = DenoisingAutoencoder(input_dim, hidden_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Dummy data with missing values
data = torch.rand((32, input_dim))
mask = (torch.rand_like(data) > 0.2).float() # 20% missing
incomplete = data * mask
# Training loop
for epoch in range(100):
noisy_input = incomplete + torch.randn_like(incomplete) * 0.1 # add noise
output = model(noisy_input)
loss = criterion(output * mask, data * mask) # only compare observed entries
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Imputation
imputed_data = incomplete + (1 - mask) * model(incomplete).detach()
print(”Imputed data:\n”, imputed_data)Time series imputation:
BRITS:
Bidirectional Recurrent Imputation for Time Series is a deep learning framework that uses bidirectional recurrent networks to impute missing values in sequential data.
It processes time series both forward and backward and captures temporal dependencies from past and future contexts at the same time.
BRITS produces more accurate reconstructions than unidirectional models.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# Simplified BRITS-style model
class BRITS(nn.Module):
def __init__(self, input_dim, hidden_dim):
super(BRITS, self).__init__()
self.rnn_forward = nn.GRU(input_dim, hidden_dim, batch_first=True)
self.rnn_backward = nn.GRU(input_dim, hidden_dim, batch_first=True)
self.fc = nn.Linear(hidden_dim * 2, input_dim)
def forward(self, x):
# Forward RNN
out_f, _ = self.rnn_forward(x)
# Backward RNN (reverse sequence)
out_b, _ = self.rnn_backward(torch.flip(x, [1]))
out_b = torch.flip(out_b, [1])
# Concatenate forward & backward hidden states
out = torch.cat([out_f, out_b], dim=-1)
return self.fc(out)
# Setup
input_dim, hidden_dim = 5, 32
model = BRITS(input_dim, hidden_dim)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Dummy time series with missing values
data = torch.rand((32, 10, input_dim)) # batch, seq length, features
mask = (torch.rand_like(data[...,0]) > 0.2) # 20% missing
incomplete = data.clone()
incomplete[~mask] = 0.0
# Training loop
for epoch in range(50):
output = model(incomplete)
loss = criterion(output[mask], data[mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Imputation
imputed_data = incomplete.clone()
imputed_data[~mask] = model(incomplete)[~mask].detach()
print(”Imputed time series:\n”, imputed_data)SAITS :
Self-attention-based imputation methods use transformer-style architectures to fill in missing values.
They model complex dependencies across features and time steps, and use attention mechanisms to capture both temporal patterns and feature correlations instead of relying on simple statistical assumptions.
Example:
from pypots.imputation import SAITS
import numpy as np
# X: (n_samples, n_steps, n_features) with np.nan for missing values
X = np.random.randn(100, 24, 5)
mask = np.random.rand(100, 24, 5) < 0.2
X[mask] = np.nan
model = SAITS(
n_steps=24,
n_features=5,
n_layers=2,
d_model=64,
d_inner=128,
n_heads=4,
d_k=16,
d_v=16,
dropout=0.1,
epochs=50,
)
model.fit(X)
X_imputed = model.impute(X)ImputeGAP Time Series Imputation:
ImputeGAP is a Python library designed for time series imputation that uses over 40 state-of-the-art algorithms and tools to deal with missing values in both univariate and multivariate datasets.
Example:
from imputegap.recovery.manager import TimeSeries
from imputegap.recovery.imputation import Imputation
from imputegap.recovery.contamination import GenGap
from imputegap.tools import utils
ts = TimeSeries()
ts.load_series(utils.search_path(”eeg-alcohol”), normalizer=”z_score”)
ts_m = GenGap.mcar(ts.data)
imputer = Imputation.MatrixCompletion.CDRec(ts_m)
imputer.impute()
imputer.score(ts.data, imputer.recov_data)
ts.print_results(imputer.metrics)GPT4TS:
GPT4TS is a general-purpose time series analysis framework. It is used for imputation of missing values in sequential data.
It reuses pretrained transformers (from NLP or vision domains) and then adapts them for time series imputation.
GPT4TS can fill gaps in time series more realistically and consistently than traditional statistical or task-specific approaches.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# Simplified GPT4TS-style model
class GPT4TS(nn.Module):
def __init__(self, input_dim, model_dim, n_heads, n_layers):
super(GPT4TS, self).__init__()
self.embedding = nn.Linear(input_dim, model_dim)
encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=n_heads)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
self.output_layer = nn.Linear(model_dim, input_dim)
def forward(self, x, mask=None):
x = self.embedding(x)
x = self.transformer(x, src_key_padding_mask=mask)
return self.output_layer(x)
# Setup
input_dim, model_dim, n_heads, n_layers = 5, 64, 4, 2
model = GPT4TS(input_dim, model_dim, n_heads, n_layers)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Dummy time series with missing values
data = torch.rand((32, 10, input_dim)) # batch, sequence length, features
mask = (torch.rand_like(data[...,0]) > 0.2) # 20% missing
incomplete = data.clone()
incomplete[~mask] = 0.0 # replace missing with 0
# Training loop
for epoch in range(50):
output = model(incomplete, mask=~mask)
loss = criterion(output[mask], data[mask]) # only compare observed entries
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Imputation
imputed_data = incomplete.clone()
imputed_data[~mask] = model(incomplete)[~mask].detach()
print(”Imputed time series:\n”, imputed_data)NuwaTS:
NuwaTS is a transformer-based framework for time series imputation and generation, based on Microsoft’s NUWA model for visual synthesis.
It treats time series as structured sequences and uses masked modeling to fill in missing values. The framework adapts transformer architectures for sequential numeric data.
Example:
import torch
import torch.nn as nn
import torch.optim as optim
# Simplified NuwaTS-style transformer
class NuwaTS(nn.Module):
def __init__(self, input_dim, model_dim, n_heads, n_layers):
super(NuwaTS, self).__init__()
self.embedding = nn.Linear(input_dim, model_dim)
encoder_layer = nn.TransformerEncoderLayer(d_model=model_dim, nhead=n_heads)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
self.output_layer = nn.Linear(model_dim, input_dim)
def forward(self, x, mask=None):
x = self.embedding(x)
x = self.transformer(x, src_key_padding_mask=mask)
return self.output_layer(x)
# Setup
input_dim, model_dim, n_heads, n_layers = 5, 64, 4, 2
model = NuwaTS(input_dim, model_dim, n_heads, n_layers)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
# Dummy time series with missing values
data = torch.rand((32, 10, input_dim)) # batch, sequence length, features
mask = (torch.rand_like(data[...,0]) > 0.2) # 20% missing
incomplete = data.clone()
incomplete[~mask] = 0.0
# Training loop
for epoch in range(50):
output = model(incomplete, mask=~mask)
loss = criterion(output[mask], data[mask])
optimizer.zero_grad()
loss.backward()
optimizer.step()
# Imputation
imputed_data = incomplete.clone()
imputed_data[~mask] = model(incomplete)[~mask].detach()
print(”Imputed time series:\n”, imputed_data)What approach to choose:
What imputation methods should you choose?
Let’s start with the size of missingness:
If you have a high missingness ratio, you can use deep learning methods: ReMasker, DSAN, and DiffPuter.
If you have low amounts of missing values, you can use missForest or iterative machine learning methods: MICE, etc.
If data is skewed, use the KNN instead.
Best Practices:
Evaluate imputation fairness. Use group-level metrics for that.
Document and justify your imputation choices.
Make sure that you avoid leakage: do not run the imputer on a full data set before splitting it into the train and test parts.
For data missing not at random, use deep generative models: GAIN, MIWAE, and HyperImpute.
Automation:
Instead of doing this every time using different methods, you can use AutoML. These Python libraries can help you automate the imputation process.
Here are some of them:
HyperImpute
Feature-Engine
Follow me to deal with missing data like a pro.
What imputation method do you use most?
Let me know in the comments 👇