Nearly Optimal Bayesian Inference for Structural Missingness

TL;DR

This paper tackles structural missingness by modeling the data-generating process with a Second-Order Structural Causal Model (SCM) prior and learning a joint predictive framework that preserves uncertainty. It decouples learning of the missing-value posterior from label prediction: a Prior Fitted Network (PFN) provides fast posterior predictive inferences from incomplete inputs, while a Flow Matching head builds an explicit missing-data posterior for sampling. Theoretical analysis shows posterior integration reduces MNAR and plug-in imputation biases and yields near-Bayes-optimal risk under the model, with a favorable sample-complexity bound for decoupled inference. Empirically, the PFN-Flow approach achieves state-of-the-art performance on a broad suite of classification and MNAR-imputation tasks, with substantial runtime speedups compared to strong baselines, highlighting practical benefits for uncertain, partially observed tabular data.

Abstract

Structural missingness breaks 'just impute and train': values can be undefined by causal or logical constraints, and the mask may depend on observed variables, unobserved variables (MNAR), and other missingness indicators. It simultaneously brings (i) a catch-22 situation with causal loop, prediction needs the missing features, yet inferring them depends on the missingness mechanism, (ii) under MNAR, the unseen are different, the missing part can come from a shifted distribution, and (iii) plug-in imputation, a single fill-in can lock in uncertainty and yield overconfident, biased decisions. In the Bayesian view, prediction via the posterior predictive distribution integrates over the full model posterior uncertainty, rather than relying on a single point estimate. This framework decouples (i) learning an in-model missing-value posterior from (ii) label prediction by optimizing the predictive posterior distribution, enabling posterior integration. This decoupling yields an in-model almost-free-lunch: once the posterior is learned, prediction is plug-and-play while preserving uncertainty propagation. It achieves SOTA on 43 classification and 15 imputation benchmarks, with finite-sample near Bayes-optimality guarantees under our SCM prior.

Figures (7)

• Figure 1: Overview of the training and inference framework. In pre-training, inputs $(X_{\text{train}},X_{\text{test}})$ are incomplete (with masks $M$), sampled from the causal model in Sec. \ref{['sec:causal_model']} and Fig. \ref{['fig:dfsm']}. The PFN is trained end-to-end to output the in-model posterior predictive distribution (PPD) directly from incomplete inputs (i.e., it implicitly marginalizes over missing values under the SCM prior), and predicts $y_{\text{test}}$ for the corresponding incomplete $X_{\text{test}}$. In addition, a Flow Matching head can be trained to model a PD over missing values for explicit sampling/imputation and uncertainty analysis, but label prediction uses the PFN output and does not require Monte Carlo over completions.
• Figure 2: Second-order SCMs for structural missingness. (Left) A causal view of structural missingness, capturing both $X \rightarrow M$ (structural invalidity/undefined values) and $M \rightarrow M$ (missingness propagation) dependencies. (Right) The second-order SCM generation pipeline: sample a causal graph and parameters, generate complete data, then generate masks via the missingness mechanism; the resulting prior is used to fit PFNs and to learn the missing-value posterior.
• Figure 3: Out-of-sample MAE under MNAR (mean$\pm$std over $K=10$ independently sampled MNAR mask realizations) across $N$ datasets (here $N=15$ from the cross-method intersection), with dataset-wise grouped bars, an aggregated AVG column, and an Avg Rank column on a secondary y-axis (lower is better).
• Figure 4: Imputation benchmark under MCAR missingness: out-of-sample MAE over 15 datasets, plus AVG and RANK (mean $\pm$ std over masks).
• Figure 5: Boosting baseline comparisons against PFN-Flow under MCAR. Each point corresponds to one dataset in one missingness group; the diagonal indicates parity.

Theorems & Definitions (31)

• Theorem 6.1: PFN as population risk minimization of posterior predictives
• Theorem 6.2: Consistency of conditional flow matching in zhou2025error
• Theorem 6.3: Posterior integration converges with no in-model bias
• Corollary 6.4: Strict Jensen gap of point imputation
• Corollary 6.5: Inevitable mismatch under a forced same-distribution imputer
• Remark 6.6: MNAR identifiability is model-dependent
• Theorem 6.7: Sample Complexity Advantage of Decoupled Inference (ICL Difficulty Proxy)
• proof
• proof
• Lemma 3.2: Tower property representation