Unleashing the Potential of Diffusion Models for Incomplete Data Imputation

2024-05-31

arXiv.org

AI Summary

To address missing value imputation under incomplete training data, this paper proposes DiffPuter—a novel diffusion-based imputation framework. DiffPuter uniquely bridges unconditional diffusion training and conditional sampling with the M-step (learning the full-variable joint distribution) and E-step (iteratively imputing missing entries via conditional expectations) of the EM algorithm, enabling end-to-end trainable progressive imputation. It requires no pretraining, auxiliary discriminators, or assumptions about missingness mechanisms, and uniformly supports joint distribution modeling under arbitrary missing data patterns. Theoretically, we establish rigorous consistency between diffusion modeling and the EM paradigm. Empirically, DiffPuter achieves state-of-the-art performance across 10 benchmark datasets, outperforming 16 baselines with average reductions of 8.10% in MAE and 5.64% in RMSE—particularly excelling in complex, non-ignorable missingness scenarios.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Reasoning under Uncertainty: Relational Probabilistic Models
• Machine Learning: Probabilistic Circuits and Graphical Models

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for incomplete, noisy, or partially observed Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

This paper introduces DiffPuter, an iterative method for missing data imputation that leverages the Expectation-Maximization (EM) algorithm and Diffusion Models. By treating missing data as hidden variables that can be updated during model training, we frame the missing data imputation task as an EM problem. During the M-step, DiffPuter employs a diffusion model to learn the joint distribution of both the observed and currently estimated missing data. In the E-step, DiffPuter re-estimates the missing data based on the conditional probability given the observed data, utilizing the diffusion model learned in the M-step. Starting with an initial imputation, DiffPuter alternates between the M-step and E-step until convergence. Through this iterative process, DiffPuter progressively refines the complete data distribution, yielding increasingly accurate estimations of the missing data. Our theoretical analysis demonstrates that the unconditional training and conditional sampling processes of the diffusion model align precisely with the objectives of the M-step and E-step, respectively. Empirical evaluations across 10 diverse datasets and comparisons with 16 different imputation methods highlight DiffPuter's superior performance. Notably, DiffPuter achieves an average improvement of 8.10% in MAE and 5.64% in RMSE compared to the most competitive existing method.

Problem

• Addressing missing data imputation using diffusion models
• Overcoming training incompleteness in generative models
• Enhancing conditional sampling for accurate value updates

Innovation

• Combines diffusion models with EM algorithm
• Uses tailored reversed sampling strategy
• Performs maximum likelihood estimation