Missingness-aware Data Imputation via AI-powered Bayesian Generative Modeling

Abstract: Missing data imputation remains a fundamental challenge in modern data science, especially when uncertainty quantification is essential. In this work, we propose MissBGM, an AI-powered missing data imputation method via Bayesian generative modeling that bridges the expressive flexibility of neural networks with the statistical rigor of Bayesian inference. Unlike existing methods that often focus on point estimates or treat the missingness mechanism implicitly, MissBGM explicitly and jointly models the data-generating and missingness mechanisms, providing principled posterior uncertainty over imputations rather than a single point estimate. We develop a stochastic optimization framework with alternating updates among missing values, model parameters, and latent variables until convergence. Our theoretical analysis shows that estimates of missing values from MissBGM converge consistently under mild assumptions. Empirically, we demonstrate that MissBGM achieves superior performance over traditional imputers and recent neural network-based methods across extensive experimental settings. These results establish MissBGM as a principled and scalable solution for modern missing data imputation. The code for MissBGM is open sourced at https://github.com/liuq-lab/MissBGM.

• The paper introduces MissBGM, which combines expressive generative AI with rigorous Bayesian inference to jointly model data and missingness for effective imputation.
• It employs a unified latent-variable framework that directly targets non-ignorable missing data (MNAR), offering accurate point estimates with posterior uncertainty.
• Empirical evaluations on synthetic and real-world datasets show lower RMSE and superior uncertainty calibration compared to classical and modern deep learning methods.

MissBGM: AI-Powered Bayesian Modeling for Missing Data Imputation

Motivation and Positioning

Handling missing data is a core issue in statistical machine learning, with critical impact in high-stakes domains such as healthcare, social science, and industrial analytics. Classical approaches (EM, likelihood, MICE) are statistically principled but limited by restrictive assumptions and poor scalability in high-dimensional, non-linear settings. Recent deep learning approaches leverage neural architectures for flexible imputation but often address only point estimates and ignore or heuristically treat the missingness mechanism, leading to unprincipled uncertainty quantification. MissBGM ("Missingness-aware Data Imputation via AI-powered Bayesian Generative Modeling") (2605.01676) advances the field by bridging expressive generative AI with rigorous Bayesian inference, explicitly co-modeling the data-generating and missingness mechanisms and providing principled posterior uncertainty over missing values.

Figure 1: MissBGM models both the data-generating and missingness mechanisms via a unified Bayesian latent-variable framework, enabling joint inference of missing values and uncertainty quantification.

Methodological Contributions

MissBGM is formulated as a joint Bayesian generative model. Each sample is generated via a latent variable $\mathbf{z}_i$ following $\pi(\mathbf{z}_i)$, with the observed data $\mathbf{x}_i$ modeled as $p_\theta(\mathbf{x}_i|\mathbf{z}_i)$ and the observation mask $\mathbf{r}_i$ as $p_\phi(\mathbf{r}_i|\mathbf{x}_i)$. Neural networks are used for the mean and variance parameterization of $p_\theta$ and for the mask logits in $p_\phi$, optionally in a Bayesian neural network (BNN) framework.

Alternating stochastic blockwise optimization is employed, iteratively updating latent variables, missing values, and model parameters on mini-batches. Posterior inference for missing entries is performed per sample, initialized via MAP estimates and refined via HMC-within-Gibbs samplers, yielding both point imputation and posterior intervals.

The model rigorously accommodates MCAR, MAR, and MNAR, but is specifically designed for the non-ignorable MNAR regime. By encoding the missingness mechanism, MissBGM targets $p(\mathbf{X}_{mis} | \mathbf{X}_{obs}, \mathbf{R})$ directly, rather than the more restrictive $p(\mathbf{X}_{mis} | \mathbf{X}_{obs})$ of MAR/MCAR.

Theoretical Guarantees

MissBGM achieves consistent estimation of the pseudo-true tempered conditional law governing missing values. Under regularity and correct specification, MissBGM's learned posteriors converge to the Bayes-optimal target for both full conditional distribution and posterior mean imputation. The tempering coefficient $\pi(\mathbf{z}_i)$0 controls the influence of the missingness term; ablation studies demonstrate that small nonzero $\pi(\mathbf{z}_i)$1 values yield optimal performance, while excessive weighting degrades results.

Empirical Evaluation

Synthetic Data Performance

MissBGM outperforms classical, ensemble, and state-of-the-art deep generative imputation methods across synthetic benchmarks varying sample size, dimensionality, and missing rate. It consistently attains lowest RMSE in most settings, demonstrating robust point imputation accuracy.

Figure 2: MissBGM maintains lowest RMSE versus baselines as dimensionality and missing rate vary, illustrating robustness to scale and missingness mechanism complexity.

MissBGM’s posterior uncertainty quantification is superior, with interval-width calibration and variation estimation accuracy exceeding Miceforest and ForestDiffusion. Generative AI baselines (e.g., GAIN, TabCSDI) fail to represent uncertainty, producing degenerate predictive distributions.

Real-World Data

On UCI and high-dimensional benchmarks (Wine, Breast Cancer, Concrete Strength, Gisette), MissBGM achieves uniformly lowest point imputation RMSE. On the 5000-feature Gisette dataset, three baselines fail to scale, while MissBGM demonstrates both scalability and accuracy, outperforming the best competitor by 10.2%.

Robustness and Ablation

Robustness experiments reveal MissBGM's performance remains stable across a spectrum of dimensionality ($\pi(\mathbf{z}_i)$2) and missingness rates, with only slight deterioration under extreme parameter regimes. As missing rate increases, RMSE curves exhibit U-shaped trends; moderate missing rates balance informativeness and severity, with optimal accuracy.

Figure 3: RMSE versus missing rate ($\pi(\mathbf{z}_i)$3, $\pi(\mathbf{z}_i)$4); MissBGM consistently outperforms across all rates.

Figure 4: RMSE versus missing rate ($\pi(\mathbf{z}_i)$5, $\pi(\mathbf{z}_i)$6); stability across scale.

Figure 5: RMSE versus missing rate ($\pi(\mathbf{z}_i)$7, $\pi(\mathbf{z}_i)$8); MissBGM superiority persists at large $\pi(\mathbf{z}_i)$9.

Ablation confirms that moderate tempering ($\mathbf{x}_i$0) improves performance; excessive mask modeling ($\mathbf{x}_i$1) negatively affects results.

Computational Efficiency

MissBGM’s joint modeling and uncertainty quantification do not impose prohibitive computational overhead. Running time benchmarks reveal MissBGM is comparable to GAIN and far more efficient than miceforest and diffusion-based methods.

Figure 6: Log-scale running time for imputation methods on moderate synthetic benchmark ($\mathbf{x}_i$2, $\mathbf{x}_i$3); MissBGM remains practical amid deep baselines.

Implications and Future Directions

MissBGM enables scalable, flexible, and statistically principled imputation in non-ignorable missingness settings—a foundation for downstream inference, sensitivity analysis, and causal modeling. Its methodological rigor and empirical robustness address limitations of both classical and modern AI-based approaches, establishing a template for future missing-data solutions in complex data environments.

Anticipated developments include: (1) theoretical identifiability of latent variables via nonlinear ICA advances; (2) deployment of advanced architectures (e.g., transformers, non-stationary models) for more nuanced data and mask modeling; and (3) further integration of MissBGM into causal inference, domain adaptation, and uncertainty-aware decision support pipelines.

Conclusion

MissBGM (2605.01676) exemplifies the convergence of deep generative modeling and Bayesian inference for missing data analysis, providing robust point imputation and principled uncertainty quantification, with theoretical and empirical guarantees under non-ignorable missingness. Its scalable architecture and rigorous modeling deliver improved accuracy, interpretability, and reliability in high-dimensional and heterogeneous settings, offering a definitive advance in the domain of modern missing data imputation.