📊 statistics

Latent Diffusion for Missing Data

This paper proposes a two-stage framework that combines a robust VAE-based imputer with a latent-space diffusion model, demonstrating that shifting generative modeling to a learned latent representation significantly improves imputation stability and quality compared to pixel-space diffusion under high levels of missing data.

Original authors: Alberte Heering Estad, Ignacio Peis, Jes Frellsen

Published 2026-05-28

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Alberte Heering Estad, Ignacio Peis, Jes Frellsen

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 have a giant photo album, but someone has spilled coffee on half the pages, blurring out important details. Your goal is to guess what the missing parts of the photos look like. This is the problem of missing data imputation.

For a long time, computers have tried to solve this by looking at the "pixels" (the tiny dots that make up the image) directly. But the paper you shared, titled "Latent Diffusion for Missing Data," argues that there's a smarter way to do this, especially when the photos are really badly damaged.

Here is the story of their discovery, explained simply:

The Problem: Trying to Fix a Blurry Photo Pixel-by-Pixel

The researchers looked at existing methods (called Pixel-Space Diffusion). Imagine trying to restore a damaged painting by looking at every single tiny dot of paint individually. If the painting is 50% covered in coffee, the computer gets confused. It sees a lot of "zero" values (the coffee spots) and tries to guess the colors based on them. As the damage gets worse, the computer starts hallucinating—making the image look grainy, noisy, and full of static, like an old TV with a bad signal.

The Solution: The "Sketch First" Approach (Latent Diffusion)

The authors, Alberte Heering, Ignacio Peis, and Jes Frellsen, proposed a new method called LDMiss. Instead of looking at every tiny dot of paint, they ask the computer to first learn the "gist" or the "sketch" of the image.

Think of it like this:

The Encoder (The Artist's Sketch): First, the computer compresses the messy, coffee-stained photo into a simplified, low-resolution sketch. It ignores the tiny details and focuses on the big shapes (like "this is a nose," "this is an eye").
The Diffusion (The Restoration): The computer then tries to fix the missing parts of this sketch, not the original photo. Because the sketch is simpler and cleaner, it's much easier to guess what the missing lines should look like, even if half the paper is stained.
The Decoder (The Final Painting): Once the sketch is fixed, the computer expands it back into a full, high-quality photo.

The Experiment: A Race Against the Coffee Spills

The team tested this on the famous MNIST dataset (which is just thousands of handwritten numbers like 0, 1, 2, etc.). They simulated "missing data" by randomly turning pixels black (like coffee spots) at different rates: 10%, 30%, 50%, and even up to 80% missing.

They compared their "Sketch First" method (LDMiss) against the old "Pixel-by-Pixel" method (DDPM) and a few other competitors.

The Results: Who Won the Race?

The Old Method (Pixel-Space): When the damage was light (10% missing), it did okay. But as soon as the damage hit 20-30%, the images started to look like static noise. By 50% missing, the results were terrible. It was like trying to fix a painting by staring at the coffee stains; the computer just got overwhelmed.
The New Method (LDMiss): This method stayed calm. Even when 50% of the data was missing, the generated numbers looked sharp and clear. It only started to get a little blurry when the damage was extreme (60-80%).
The "Fill-in" Test: When they used these models to actually fill in the missing parts of the photos, LDMiss was consistently better. It didn't just look good; it actually guessed the missing numbers more accurately than the others.

Why Does This Work?

The paper suggests that the "Sketch First" approach acts like a filter. When you look at raw pixels, a missing spot is just a confusing black square. But when you look at the "latent" (sketch) version, the computer understands the meaning of the image. It knows, "This is a '3', so even if the top loop is missing, I know it should be round." This prevents the computer from amplifying the "zero" errors caused by the missing data.

The Bottom Line

The paper concludes that if you are trying to fix incomplete data, don't stare at the raw pixels. Instead, teach the computer to understand the underlying "vibe" or "structure" of the data first, and then do the fixing there.

What the paper doesn't say:

They did not test this on medical records, stock markets, or climate data (yet).
They did not claim this works for "Missing Not At Random" (where the missing data is related to the value itself, like sick people dropping out of a study). They only tested "Missing Completely At Random" (purely random coffee spills).
They only tested on handwritten numbers (MNIST), not complex photos of cats or cars.

In short: LDMiss is a more robust way to guess missing pieces of a puzzle because it focuses on the big picture rather than getting lost in the tiny, damaged details.

Technical Summary: Latent Diffusion for Missing Data

Problem Statement
Real-world datasets frequently contain missing values due to factors such as loss to follow-up in clinical studies or incomplete documentation. While probabilistic generative models (VAEs, GANs, Flows) and more recently diffusion models have shown promise for imputation, existing diffusion-based methods typically operate directly in the data (pixel) space. A critical gap exists in understanding how the choice of representation space—pixel versus latent—affects generation quality and imputation performance when training data is heavily incomplete. Specifically, it is unclear whether operating in pixel space degrades under high rates of Missing Completely At Random (MCAR) corruption, where zero-imputation of missing values may introduce artifacts that disrupt the score estimation.

Methodology
The authors propose LDMiss, a two-stage framework that shifts the diffusion process from pixel space to a learned latent representation to improve robustness against missing data.

Latent Space Learning (Stage 1): A robust Variational Autoencoder (VAE) is trained on incomplete observations. The authors utilize a $\beta$ -VAE architecture where missing entries are zero-imputed. The reconstruction loss is computed only on observed dimensions, while the KL divergence is weighted by a small $\beta$ ( $10^{-6}$ ) to ensure high-fidelity reconstructions. This stage learns a compact semantic feature space ( $z = E(x)$ ) from the incomplete data.
Latent Diffusion (Stage 2): A diffusion model is trained on the latent representations $z$ rather than the raw pixels $x$ . The forward process follows a Variance-Preserving (VP) Stochastic Differential Equation (SDE), and the reverse process is learned via a time-dependent neural network $s_\theta$ optimized using a denoising score-matching objective adapted for latent variables.
Imputation Strategy (Self-Guidance): For downstream imputation, the authors employ a "self-guidance" method. Since the concept of "missing dimensions" does not exist in the compressed latent space, the generative process is conditioned on the observed pixel data. The conditional score is approximated using Bayes' rule, decomposing it into the unconditional score and a guidance term derived from Tweedie's formula. The model decodes the latent posterior mean estimate to the pixel space to compute gradients, effectively guiding the sampling process using its own predictive capabilities.

Experimental Setup
The study conducts a controlled comparison on the MNIST dataset under MCAR corruption at varying missing rates (up to 80%). The proposed LDMiss is benchmarked against:

Pixel-space DDPM: Trained using the MissDiff approach (zero-imputation with masked loss) and the DiffPuter EM-based iterative algorithm.
Imputation Baselines: Autoencoder reconstruction and DDPM with replacement or guidance strategies.

Key Results

Sample Quality: LDMiss maintains high sample quality (measured by FID and IS) and remains stable up to 50% missingness. In contrast, pixel-space DDPM degrades progressively as missingness increases, becoming "grainy" beyond 10% missingness due to the direct processing of zero-imputed values. DiffPuter performs well at low-to-moderate missingness but loses diversity and fidelity at 60–80% missingness.
Imputation Performance: LDMiss achieves consistently superior imputation performance (MSE, FID, IS) across all training missing rates. It remains robust even when trained with 80% missing data. Pixel-space DDPM degrades as training missingness increases because its score network relies on a degraded estimation.
Stability: LDMiss exhibits lower variance across random seeds compared to pixel-space alternatives, indicating more consistent performance.

Significance and Claims
The paper claims that latent-space modeling mitigates the amplification of artifacts caused by zero-imputed inputs, which is a primary failure mode for pixel-space diffusion models trained on incomplete data. By compressing data into semantic features before the diffusion process, the model learns a more robust generative prior for incomplete-data learning.

The authors conclude that LDMiss is a strong and practically useful alternative to pixel-space diffusion for missing-data problems. They note that their findings are currently limited to the MNIST dataset under MCAR assumptions with zero-imputation and that architectural differences may contribute to the observed performance gap. Future work is suggested to test these advantages on more complex datasets, different missingness mechanisms (MAR/MNAR), and alternative imputation functions.