📊 statistics

Missing Data Imputation under Manifold Hypothesis

This paper proposes a model-based missing data imputation method that leverages mixture variational autoencoders and latent space diffusion to sample from the conditional distribution of missing values, thereby respecting the underlying data manifold geometry while providing uncertainty quantification and efficient on-the-fly imputation.

Original authors: Zelong Bi, Amuchechukwu Ibenegbu

Published 2026-07-20

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Zelong Bi, Amuchechukwu Ibenegbu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to finish a giant, intricate mosaic, but someone has snipped out hundreds of tiny tiles. You can see the surrounding pieces, and you know the picture is supposed to be a smooth, flowing landscape, not a jagged mess of random colors. This is the daily struggle of data scientists dealing with "missing data." In the real world, sensors break, surveys get skipped, and records get lost, leaving holes in our digital maps. For decades, the standard way to fill these holes was to guess based on what was nearby, like a neighbor filling in a blank wall with whatever bricks were left in the driveway. But what if the wall isn't flat? What if the picture is actually a curved, twisting sculpture, like a rollercoaster track or a twisted ribbon? If you try to fill in the gaps on a curved track using flat, straight-line guesses, you'll end up with a piece that looks okay up close but falls right off the track when you step back.

This paper dives into a specific corner of mathematics called the "manifold hypothesis." Think of this as the idea that even though our data might look like a chaotic cloud of points in a huge, high-dimensional room (imagine a room with hundreds of directions you can move), the points actually live on a much smaller, smoother surface hidden inside that room. It's like realizing that all the ants in a chaotic pile are actually marching in a single, winding line on a piece of paper. The paper also uses "Variational Autoencoders" (VAEs), which are like smart, flexible cameras that can learn to take a complex 3D object and flatten it into a simple 2D drawing, and then rebuild the 3D object perfectly from that drawing. The big question is: if we lose some parts of the 3D object, can we use that 2D drawing to figure out exactly what the missing pieces should look like, respecting the curve of the track rather than just guessing?

The authors, Zelong Bi and Amuchechukwu Ibenegbu, propose a new way to fix these broken mosaics by treating the data like a curved surface rather than a flat sheet. They argue that the most popular methods currently used, like a tool called MissForest, are great at finding patterns but often fail to respect the "geometry" of the data. To visualize this, imagine MissForest trying to fill in a missing spot on a circle by drawing a straight line across the gap; the result might be mathematically close to the neighbors, but it breaks the circle. The authors' method, however, understands that the data lives on a curve, so it fills the gap by following the arc, keeping the shape intact.

To do this, they use a two-step magic trick. First, they use a "Mixture of VAEs" to learn the shape of the hidden surface. Imagine this as having a team of artists, each responsible for a different section of the curve (like one artist for the top of the hill, another for the bottom). They map the messy, high-dimensional data down to a simple, low-dimensional "latent space" where the curves are easy to see. Then, when a piece is missing, they don't just guess; they use a statistical procedure called "Sampling-Importance-Resampling" (SIR). Think of SIR as a game of "hot potato" where they generate thousands of possible guesses for the missing piece, but they only keep the ones that fit perfectly with the known parts of the curve, discarding the ones that look weird or out of place. This allows them to not just give one answer, but to show a range of possible answers, effectively quantifying how uncertain they are about the fill.

But they didn't stop there. They realized that sometimes the artists (the VAEs) don't have enough examples to learn every twist and turn of the curve perfectly. So, they added a "joint diffusion process." Imagine this as a magical fog that slowly clears up. You start with a completely blurry, noisy guess of the missing piece and the shape of the curve, and the model slowly "denoises" it, refining the guess step-by-step until it snaps into a sharp, perfect fit that respects both the local details and the global shape. This happens in the low-dimensional space, making it much faster and more efficient than trying to clean up the noise in the huge, high-dimensional room.

The results of their experiments are quite compelling. When they tested their method on synthetic data shaped like circles, spheres, and donuts (toruses), their approach produced imputations that were far better at preserving the true shape of the data than the leading methods. While the standard method (MissForest) sometimes got a lower "error score" (RMSE) by just guessing numbers that were close to the neighbors, it often broke the geometric shape. The authors' method, however, achieved the lowest "Wasserstein distance," a fancy way of saying the overall shape and distribution of their filled-in data looked much more like the original, unbroken picture.

On real-world datasets, like records of superconductors (materials that conduct electricity with zero resistance) and energy consumption, the authors found their method was a strong competitor. It performed almost as well as MissForest when data was missing randomly, but it shined when the missing data was tricky or when a lot of it was gone. In these tough scenarios, their method stayed robust, while others started to crumble. For instance, on the superconductivity dataset, their method maintained the lowest error rates even when the missingness wasn't random, suggesting that understanding the underlying geometry helps the model guess correctly even when the data is being "hiding" in a non-random way.

However, the paper is careful to note that this isn't a magic wand for every problem. The method relies heavily on the assumption that the data actually does follow a smooth, low-dimensional curve (the manifold hypothesis). If the data is just random noise or doesn't have a clear shape, the method might struggle. They also found that on a small dataset like the wine quality data, the method didn't perform as well, likely because there wasn't enough data to teach the artists how to draw the curve in the first place. In those cases, simpler methods that just look at the nearest neighbors worked better.

Ultimately, this paper suggests that by combining the geometric understanding of manifold learning with the generative power of diffusion models, we can fill in missing data in a way that feels "right" to the shape of the universe the data comes from. It's a shift from just plugging holes with the nearest available brick to sculpting the missing piece so it fits the curve of the wall perfectly. While it requires more computing power and a specific type of data structure to work, it offers a promising path toward more accurate and trustworthy data recovery, especially in fields where the shape of the data matters as much as the numbers themselves.

Technical Summary: Missing Data Imputation Under Manifold Hypothesis

Problem Statement

Missing data imputation seeks to systematically estimate and replace missing values in a dataset. While the manifold hypothesis posits that high-dimensional data concentrate near a low-dimensional embedded manifold, mainstream imputation methods have largely failed to leverage this geometric structure. Existing state-of-the-art methods, such as MissForest, often fail to respect the underlying geometry when it exists, leading to imputations that may be point-wise accurate (low RMSE) but geometrically inconsistent with the data distribution. Furthermore, many current approaches are transductive, meaning they optimize imputations for a specific dataset without learning a reusable model, making them inefficient for handling new data or online imputation.

Methodology

The proposed framework addresses these limitations by combining manifold learning with probabilistic sampling and diffusion models. The methodology proceeds in three main stages:

1. Manifold Learning via Mixture VAEs

The approach assumes the data lies on a low-dimensional manifold of intrinsic dimension $d \ll p$ . To capture this structure, the authors utilize a mixture of Variational Autoencoders (VAEs). In this formulation:

The manifold is covered by multiple overlapping "charts."
Each chart $c \in \{1, \dots, C\}$ is associated with a local coordinate system $z \in \mathbb{R}^d$ .
The model learns an encoder $E_c$ and a decoder $D_c$ for each chart, approximating the conditional distributions $p(z|x, c)$ and $p(x|z, c)$ as Gaussians.
This multi-chart formulation allows the model to faithfully recover the underlying geometry, which single-chart VAEs often fail to do.

2. Imputation via Sampling-Importance-Resampling (SIR)

The core imputation task is to sample from the conditional distribution $p(x_{mis} | x_{obs})$ . Since the joint posterior $p(z, c | x_{obs})$ is intractable, the authors employ a Sampling-Importance-Resampling (SIR) procedure:

Sampling: Joint latent samples $(z_k, c_k)$ are drawn from the training data using the learned manifold model.
Weighting: For a specific observed vector $x_{obs}$ , weights are computed as $w_k = p(x_{obs} | z_k, c_k)$ .
Resampling: A new latent pair $(z^*, c^*)$ is sampled from the set of training pairs according to the normalized weights.
Generation: The missing values are generated by sampling from $p(x_{mis} | z^*, c^*)$ , which is Gaussian due to the model's structure.
This process yields a model-based, inductive imputation that respects both geometric and distributional constraints and allows for uncertainty quantification.

3. Joint Latent Diffusion Process

To overcome the limitation of being constrained by the support of the observed training samples, the authors introduce a joint diffusion process in the latent space $(z, c)$ .

Forward Process: Noise is added independently to the continuous latent coordinates $z_t$ and the discrete chart labels $c_t$ .
Reverse Process: A neural network is trained to predict the clean latent variables $(z_0, c_0)$ from the noisy state $(z_t, c_t)$ and time step $t$ .
Sampling: The reverse process allows for the generation of high-quality latent samples at an arbitrary scale, which are then decoded to produce imputations. This enhances the imputation quality, particularly in regions of the data space not densely covered by the original training set.

Key Contributions

Model-Based Imputation via SIR: The authors propose a procedure that samples directly from $p(x_{mis} | x_{obs})$ using SIR, leveraging learned manifold structures. This approach is inductive, supports on-the-fly imputation without retraining, and provides uncertainty quantification.
Joint Latent Diffusion: A novel joint diffusion process is developed for the latent space that simultaneously updates continuous coordinates and discrete chart membership labels. This improves the generation of high-quality samples that preserve underlying geometry and distribution.
Geometry-Aware Performance: The method is designed to respect the intrinsic geometry of the data, addressing the failure of traditional methods to maintain structural consistency.

Experimental Results

The proposed method was evaluated on synthetic datasets (circles, spheres, tori) and six real-world datasets (superconductivity, energy, finance, powerplant, wine, facial expression) under Missing Completely at Random (MCAR), Missing at Random (MAR), and Missing Not at Random (MNAR) mechanisms.

Metrics: Evaluation utilized Root Mean Squared Error (RMSE) and Wasserstein distance. The authors note that while RMSE measures point-wise accuracy, Wasserstein distance better captures alignment with the underlying geometry and distribution.
Synthetic Data: The proposed method achieved the lowest Wasserstein distance on all synthetic datasets, indicating superior preservation of geometric structure compared to baselines like MissForest, MICE, and GAIN.
Real-World Data:
- The method achieved competitive performance with MissForest, often outperforming it as missing rates increased or under complex missing mechanisms (MAR/MNAR).
- On the superconductivity dataset, the method demonstrated robustness, maintaining the lowest average RMSE and Wasserstein distance across missing mechanisms.
- In downstream prediction tasks (linear regression on superconductivity data), the proposed method achieved the best predictive performance under MNAR settings and second-best under MCAR/MAR.
- The method showed limitations on the wine dataset (small sample size), where KNN performed better, suggesting that reliable manifold learning requires sufficient fully observed cases.

Significance and Claims

The paper claims that its primary significance lies in bridging the gap between manifold learning and missing data imputation. By treating imputation as a sampling problem from a learned geometric structure, the method provides a systematic way to fill missing values that are consistent with the data's intrinsic manifold.

The authors emphasize that their approach is inductive, allowing for the imputation of new data points without re-optimizing the entire procedure, a distinct advantage over transductive methods. Furthermore, the probabilistic nature of the sampling allows for uncertainty quantification of the imputations, a feature often absent in deterministic methods.

The paper modestly acknowledges limitations: the method relies on the assumption that data satisfies the manifold hypothesis, which can be difficult to verify in practice. Additionally, due to the reliance on smooth structures, the method is primarily suited for continuous variables and may not be directly applicable to discrete data without modification. Finally, the authors note that developing evaluation metrics that simultaneously capture geometric fidelity and distributional consistency remains an open direction for future work.