🤖 machine learning

Learning What Not to Impute: An Uncertainty-Aware Diffusion Framework for Meaningful Missingness

This paper introduces Diff-Joint, an uncertainty-aware diffusion framework that addresses the selective imputation problem by distinguishing between meaningfully missing entries and observation-based gaps, thereby jointly inferring which values to preserve and which to recover to improve downstream task performance.

Original authors: Lixing Zhang, Yidong Ouyang, Weifu Li, Shixiang Zhu, Guang Cheng, Liyan Xie

Published 2026-06-04

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Lixing Zhang, Yidong Ouyang, Weifu Li, Shixiang Zhu, Guang Cheng, Liyan Xie

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

The Big Problem: The "Blank Space" Mystery

Imagine you are a detective trying to solve a case using a witness statement. The witness wrote down a list of facts, but some parts are blank.

In the world of data science, these blanks are called missing values. Traditionally, when computers see a blank, they assume it's a mistake—like a piece of paper that got torn out or a pen that ran out of ink. The computer's job is usually to guess what should have been there and fill it in. This is called imputation.

But the authors of this paper point out a crucial flaw in this logic: Not all blanks are mistakes.

Sometimes, a blank space is actually a valid answer.

The Mistake: A doctor forgot to write down a patient's blood pressure. This is a "missing value" that needs to be guessed.
The Valid Blank: A survey asks, "What is your spouse's income?" If the person is single, the answer isn't "unknown"; the answer is "Not Applicable." The blank is a meaningful part of the story.

If a computer blindly fills in the "Not Applicable" blank with a random guess (like "$50,000"), it creates a lie. It distorts the data and confuses the model.

The Solution: Diff-Joint (The Smart Detective)

The authors propose a new method called Diff-Joint. Think of it as a detective who doesn't just try to fill in the blanks, but first asks: "Is this blank a mistake, or is it a deliberate 'N/A'?"

Here is how it works, using a few metaphors:

1. The Two Masks

The method treats every missing piece of data as having two layers:

The Value Layer: What number or word should go here?
The Mask Layer: Is this blank a "Mistake" (needs filling) or a "Meaningful Blank" (should stay empty)?

2. The "Diffusion" Game (The Shredder and the Reassembler)

The core engine uses something called a Diffusion Model. Imagine you have a clear photo of a face, and you slowly turn it into static noise (like TV snow) until you can't see anything. A diffusion model learns how to reverse that process: starting with pure noise, it slowly "denoises" it back into a clear photo.

Diff-Joint does this with a twist. It doesn't just try to reconstruct the photo (the data values); it also tries to reconstruct the Mask (the decision of which parts should be blank).

3. The "Uncertainty" Test (The Confidence Check)

This is the secret sauce. The method runs a simulation where it generates many possible versions of the missing data.

Scenario A (The Mistake): If the computer is trying to guess a missing blood pressure, it might generate 10 different guesses (120, 125, 118, etc.). These guesses are all close together. The computer is confident (low uncertainty). It knows it should fill this in.
Scenario B (The Meaningful Blank): If the computer is trying to guess a "Spouse's Income" for a single person, it might generate 10 wildly different, nonsensical guesses (because there is no real answer). The guesses are all over the place. The computer is confused (high uncertainty).

The Rule: If the computer is highly confused (high uncertainty), it decides, "This blank is probably meaningful. I should leave it alone." If it is confident, it fills it in.

How It Learns (The Iterative Loop)

The method doesn't get it right the first time. It works like a sculptor refining a statue:

Guess: It starts by assuming all blanks are mistakes and fills them with random guesses.
Train: It learns from the data what the "real" patterns look like.
Test: It runs the "Uncertainty Test" again. It sees which guesses were shaky and which were solid.
Refine: It updates its map. It says, "Okay, I was wrong about this one; it's actually a meaningful blank." It keeps the blank empty.
Repeat: It does this over and over, getting better at distinguishing between "Mistakes" and "Meaningful Blanks" until the picture is clear.

The Results: Why It Matters

The authors tested this on two types of data:

Fake Data: A made-up database where they knew exactly which blanks were mistakes and which were "N/A."
Real Data: Medical records from a hospital (MIMIC-IV-ED).

The Findings:

Better Detection: Diff-Joint was much better at spotting the "Meaningful Blanks" (the "N/A"s) than other methods, which just tried to fill everything in.
Better Accuracy: Because it stopped trying to fill in the "N/A" blanks with lies, the remaining data was more accurate.
Better Predictions: When they used this cleaned-up data to predict future outcomes (like whether a patient would be admitted to the hospital), the predictions were more accurate.

The Bottom Line

Most data tools treat every missing piece as a broken puzzle piece that needs to be glued back in. Diff-Joint teaches the computer to recognize that sometimes, the puzzle piece is missing on purpose. By learning what not to impute, it preserves the true meaning of the data, leading to smarter and more honest results.

Technical Summary: Diff-Joint – An Uncertainty-Aware Diffusion Framework for Meaningful Missingness

1. Problem Definition

Standard missing value imputation methods operate under the assumption that all missing entries in a dataset correspond to unobserved regular values that must be recovered. However, in many real-world scenarios, missingness arises from two distinct sources:

Meaningfully Missing (MM): Entries that are intrinsically absent and semantically valid (e.g., "N/A" for income in a survey for an unemployed respondent, or a lab test not ordered in clinical records). Imputing these with regular values distorts the data distribution and removes semantic information.
Observation-Induced Missing: Entries that are missing due to the observation process (e.g., recording errors, data transfer failures) and should be imputed.

The core challenge is that these two sources are not directly labeled in observed data; they both appear as the same missing token (e.g., NA). Existing methods typically treat all missing entries as recoverable, failing to distinguish between entries that should be preserved as NA and those that require imputation. This paper formalizes this as a selective imputation problem: the goal is to jointly infer which missing entries are meaningfully missing (and should be preserved) and which are observation-induced (and should be imputed).

2. Methodology: Diff-Joint

The authors propose Diff-Joint, a diffusion-based framework designed to perform selective imputation by jointly modeling tabular data and a latent missingness mask.

2.1 Joint Diffusion State

The method defines a joint diffusion state $s = (x, c)$ , where:

$x$ represents the completed tabular values.
$c$ is a binary mask indicating which entries are meaningfully missing ( $c_j=1$ ) versus observation-induced missing ( $c_j=0$ ).

2.2 Iterative Refinement Framework

Diff-Joint operates as an iterative scheme alternating between model training and latent state updates:

Initialization: Observed missing values are randomly filled to create an initial joint state, treating all missing entries as observation-induced ( $c=0$ ).
Model Update: A diffusion model is trained on the current joint state $\{(\hat{x}^{(t)}, \hat{c}^{(t)})\}$ to capture dependencies between data values and the meaningful-missingness patterns.
Latent-State Update (Conditional Sampling & Aggregation):
- Sampling: Given the observed data, the model draws $K$ conditional samples of the joint state $(x, c)$ .
- Imputation Update: For missing entries, the imputed value is updated using the mode (discrete) or mean (continuous) of the $K$ samples.
- Uncertainty-Aware Aggregation: To distinguish MM from observation-induced missingness, the method calculates an uncertainty score for each missing entry based on the variance (continuous) or entropy (discrete) of the $K$ $K$ samples.
  - Hypothesis: Observation-induced missing entries correspond to recoverable regular values and thus exhibit low uncertainty (concentrated predictive distributions). Meaningfully missing entries are not well-explained by a single regular state, leading to higher uncertainty.
- Mask Update: The method combines two signals to update the mask $\hat{c}$ $\overset{c}{^}$ :
  1. Uncertainty Signal: Entries with high uncertainty scores are clustered (via k-means) and flagged as potential MM.
  2. Sampled Mask Signal: The majority vote of the sampled masks $c^{(k)}$ .
  3. Fusion Rule: The final mask is updated using a logical OR operation: an entry is classified as MM if either the uncertainty pattern or the sampled mask supports the MM interpretation.

2.3 Identifiability

The paper provides theoretical insights (Proposition 3.1) establishing conditions under which meaningful missingness is identifiable from observed data. Specifically, identifiability holds if the observation-induced missingness mechanism is known (e.g., MCAR with known probabilities) or if the structural constraints of the data allow the separation of the two mechanisms.

3. Key Contributions

Problem Formulation: The authors formalize the selective imputation problem for tabular data, distinguishing between meaningfully missing entries (to be preserved) and observation-induced missing entries (to be imputed).
Diff-Joint Framework: They introduce a diffusion-based framework that jointly models completed values and a meaningful-missingness mask. The method iteratively refines both imputed values and MM labels through conditional sampling and uncertainty-aware aggregation.
Empirical Validation: The method is evaluated on a synthetic Bayesian-network dataset and a real-world MIMIC-IV-ED dataset. Results demonstrate that Diff-Joint effectively identifies meaningfully missing entries while maintaining competitive imputation accuracy and improving downstream task performance.

4. Experimental Results

The authors evaluated Diff-Joint against baselines including simple statistical imputation (Mean/Mode), classical iterative methods (missForest), masked autoencoders (CACTI/CMAE), and diffusion-based imputation (DiffPuter).

Synthetic Dataset (Bayesian Network):
- MM Identification: Diff-Joint is the only method explicitly identifying MM. It achieved high recall (76.99%–92.02%) across varying missing ratios (10%–40%).
- Imputation Accuracy: Diff-Joint achieved the best token-level recovery accuracy for discrete variables across all ratios. For continuous variables, it remained competitive, balancing the dual objectives of MM identification and value recovery.
- Downstream Performance: In multi-class classification tasks, Diff-Joint significantly outperformed all baselines (e.g., Macro-F1 of 75.54% vs. ~46% for the next best), demonstrating that preserving MM states improves predictive utility.
Real-World Dataset (MIMIC-IV-ED):
- MM Identification: Diff-Joint successfully identified synthetic MM labels introduced on clinical outcome variables.
- Imputation: It achieved the best accuracy for discrete token-level recovery at lower missing ratios and remained competitive for continuous imputation (often second-best with negligible difference from the best).
- Downstream Performance: Diff-Joint showed consistent gains in Macro-F1 and ROC-AUC for predicting clinical outcomes (e.g., Critical status, ICU transfer), particularly for severe outcomes where preserving the "meaningful missing" signal was crucial.

5. Significance and Claims

The paper claims that Diff-Joint addresses a critical gap in missing data handling: the inability of current methods to distinguish between "missing because it doesn't apply" and "missing because it wasn't recorded."

Semantic Preservation: By explicitly modeling the "meaningfully missing" state, the framework prevents the distortion of data distributions that occurs when NA is naively imputed with a regular value.
Downstream Utility: The authors argue that treating NA as a valid semantic state, rather than an error to be fixed, provides predictive signals that improve downstream decision-making, particularly in high-stakes domains like healthcare.
Robustness: The method remains robust across different missingness mechanisms (MCAR, MAR, MNAR) and does not require ground-truth labels for the missingness type during training.

The authors acknowledge limitations, noting that the framework assumes known candidate columns for MM and relies on a detectable uncertainty gap. They also note that the iterative diffusion procedure is computationally more expensive than non-iterative baselines, though this cost can be controlled by the number of iterations. Future work is suggested to relax identifiability conditions and improve scalability.