RefiDiff: Progressive Refinement Diffusion for Efficient Missing Data Imputation

TL;DR

This work tackles missing data imputation in high-dimensional, mixed-type datasets under MNAR. It introduces RefiDiff, a progressive refinement framework that fuses fast local imputations with a Mamba-based diffusion denoiser to capture long-range dependencies, encoding mixed-type data into unified tokens. The approach achieves state-of-the-art performance across MCAR, MAR, and MNAR with robust out-of-sample generalization and significantly lower runtime than diffusion-only baselines, supported by theoretical guarantees on conditional sampling and KL approximation. Its practical impact lies in providing an accurate, scalable, tuning-free imputation solution suitable for real-world, heterogeneous data.

Abstract

Missing values in high-dimensional, mixed-type datasets pose significant challenges for data imputation, particularly under Missing Not At Random (MNAR) mechanisms. Existing methods struggle to integrate local and global data characteristics, limiting performance in MNAR and high-dimensional settings. We propose an innovative framework, RefiDiff, combining local machine learning predictions with a novel Mamba-based denoising network efficiently capturing long-range dependencies among features and samples with low computational complexity. RefiDiff bridges the predictive and generative paradigms of imputation, leveraging pre-refinement for initial warm-up imputations and post-refinement to polish results, enhancing stability and accuracy. By encoding mixed-type data into unified tokens, RefiDiff enables robust imputation without architectural or hyperparameter tuning. RefiDiff outperforms state-of-the-art (SOTA) methods across missing-value settings, demonstrating strong performance in MNAR settings and superior out-of-sample generalization. Extensive evaluations on nine real-world datasets demonstrate its robustness, scalability, and effectiveness in handling complex missingness patterns.

Figures (11)

• Figure 1: Overview of the proposed imputation framework. The process begins with a warm-up stage on pre-processed data, followed by a diffusion module that iteratively denoises the data. A polishing stage further enhances the imputations. Our designed denoiser $\theta_2$ will be shown in Figure \ref{['fig:denoising_network_ours']}.
• Figure 2: Illustration of the denoising network $\theta_2$, featuring a diamond-shaped structure with Mamba-based residual Up/Down blocks.
• Figure 3: Comparison between DIFFPUTER and our method on two datasets (California and Magic) under the MNAR setting. (a) and (b) show in-sample MAE and RMSE over iterations. (c) compares denoising network parameter counts. (d) presents average runtime over 10 random masks.
• Figure 4: Ablation experiments of our proposed framework in various settings and components.
• Figure 5: Critical difference diagram for numerical and categorical columns performance metrics.

Theorems & Definitions (12)

• Proposition 1: Approximation Bound for RefiDiff
• Proposition 2
• proof
• Lemma 1
• proof : Proof of Lemma \ref{['lem:measurable_mapping']}
• Proposition 3
• proof
• Proposition 4: Conditional Consistency of RefiDiff
• proof
• Proposition 5: Approximate Conditional Consistency Bound