Marformer: A Transformer for Predicting Missing Data Distributions
Authors: Prabhav Singh, Xiheng Tom Wang, Haojun Shi, Jason Eisner
Organizations: School of Computing, College of Natural Sciences The University of Texas at Austin, Austin TX · Department of Computer Science Johns Hopkins University, Baltimore MD · Department of Computer Science Yale University, New Haven, CT
Real decisions are made under incomplete information. If we observe only some of the random variables we need, we can predict the others. The \textbf{conditional marginals} over the missing variables are the key ingredient for computing Bayes risk and Value of Information (VOI), the expected gain from acquiring one more observation before deciding. We present the Marformer, a Transformer trained to directly predict conditional marginals given any set of observed values. Like BERT, which is trained to predict missing words from context, the Marformer constructs a hidden-vector representation for each distribution p(Xi) and iteratively refines it through attention to other distributions p(Xj). Unlike generative approaches, the Marformer does not model the full joint distribution, requires no domain knowledge of the data-generating process, and makes all predictions in a single forward pass. We evaluate across three synthetic domains with missing data---Bayesian networks, discretized multivariate Gaussians, and structured annotation data. The Marformer can match or outperform classical missing-data methods, even when those methods are given the true model family and prior that generated the synthetic data. We also evaluate on a real annotation dataset, where the Marformer outperforms the evaluated baselines at the largest training size. In both cases, the Marformer is substantially faster than the evaluated generative baselines.
Figures & tables
Figure 1: The complete instance x at top reveals values only for attributes in r (black), where (x,r)∼P∗ . Thus, our methods receive only the incomplete instance xr at bottom. In this example, xr is used both as a test instance and as a training instance (transductive learning). At test time, a model is evaluated (top) by its ability to predict individual missing values (red) from the observed values xr (black). It predicts a distribution for each red attribute and is rewarded if this distribution assigns high probability to the true value. At training time, the Marformer (right) trains to predict masked values (blue) from seen values xs (black). Many different masks s are used. In contrast, a classical generative model (left) trains to jointly predict all observed variables (blue) from nothing, using the missing variables (white) as latent variables.
Figure 2: Marformer architecture. The annotation domain (items, judges, criteria) is used as the running example, but the architecture supports any typed entity structure.
Figure 3: Illustration of the complete-data distribution and MCAR missingness in Domain 1, using a four-variable example. (a) A Bayesian network over binary variables, with logistic conditional probabilities at each node. (b) An incomplete instance under MCAR: the observed index set R is independent of the complete data. Blue nodes show observed values; white nodes are missing. The variables labeled Yi in the figure correspond to Xi in the text. The experiments use randomly generated networks with 10 variables.
Figure 4: Learning curves for Domains 1 and 2, showing KL divergence of the true marginals from the predicted marginals ( equation 10 ), averaged over 10 distributions P∗ . Shaded regions show 95% bootstrap confidence intervals. In Domain 1, the KL divergences are so small that they are displayed on a log scale. For measuring KL, the true probabilities are computed exactly in Domain 1, but are estimated in Domain 2 using Monte Carlo integration. That is why in Domain 2, the KL divergence between two copies of the true model qϕ∗ is measured as slightly positive rather than zero (horizontal green line): the estimated probabilities are slightly different in the two copies.
Figure 5: Horizontal axis: Cross-entropy H(P∗,qϕ^) achieved by the estimated joint model (MAP-EM in Domain 1, § 4.4 ; direct MAP optimization in Domain 2, § 5.4 ). Vertical axis: Cross-entropy H(P∗,pθ^) achieved by the estimated Marformer . Different rows evaluate different Marformer sizes; different columns evaluate different ∣Train∣ sizes. Each point corresponds to a single marginal prediction P∗(Xi∣xr) on a test instance with i∈/r(m) , color-coded by the entropy H(P∗) of the true marginal distribution (aleatoric uncertainty). Points below the diagonal are predictions on which Marformer is more accurate than the joint model. Brighter points have higher entropy, due to less informative conditions xr , so they tend to have higher cross-entropy under both models. All graphs for the same domain use the same true generative distribution P∗ .
Figure 6: Training time (wall-clock seconds) on Domains 1 and 2 is essentially proportional to the training set size. Even for these small problems, Marformer is substantially faster to train. It is also faster at test time (not shown here), since inference requires only a single forward pass per test instance.
Figure 7: Domain 2 begins with a joint Gaussian distribution over latent continuous variables Z (left), and discretizes their values into fixed but irregular bins, yielding a joint distribution over observable discrete variables X (right). The positive correlation of (Z1,Z2) carries over to (X1,X2) , but the probabilities in the grid at right are also affected by the bin widths. Our experiments use 10 variables of each type, not 2 as illustrated here.
Figure 8: Log-loss (a)–(b) and MBR- L2 (c)–(d) on missing ratings in the fixed test corner as the amount of training data increases. Item-generalization results extend to Itrain=4000 for all models; in judge generalization, Marformer NT stops at Jtrain=40 because test judges are excluded from training. Shaded bands show ± one bootstrap standard error of the mean, resampling items on the item axis and judges on the judge axis. Full comparisons including the empirical baselines are reported in Appendix C .
Figure 9: Calibration of predicted categorical probabilities against empirical frequencies at Itrain=4000 and Jtrain=40 , the largest training sizes with results for all models shown. Probabilities for all C=4 Likert categories are pooled in the reliability calculation. The diagonal indicates perfect calibration; the displayed smECE widths are relplot ’s reported calibration-error widths from 500 bootstrap resamples.
Figure 10: Log-loss (a) and MBR-L2 (b) on the LLMRubric test set as a function of training items. Stan CPM is run with the HMC configuration described in Appendix C . Bands are bootstrap standard errors of the mean, obtained by resampling conversations conditional on the fitted model and observed data realization. Runtime curves appear in Appendix D .
Figure 11: Calibration plots (predicted probability vs. empirical frequency) on LLMRubric at full and minimal training size. At Itrain=175 , Marformer is well-calibrated; at Itrain=10 , calibration degrades for all methods, with MIWAE and ReMasker exhibiting overconfident predictions at the scale extremes.
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 12: Learning curves, similar to Figure 4 but where Train is MAR rather than MCAR. Despite the biased training examples, performance is qualitatively similar, though loss for all methods is higher in Domain 1.
Figure 13: Prediction-by-prediction comparison of two methods, similar to Figure 5 but where the models were trained on MAR rather than MCAR data.
Figure 14: Full log-loss curves including all baselines. Empirical baselines (Unigram Pool, Structured NB, Log-Linear) generally have higher log-loss than the Marformer and Stan variants at larger training sizes in both settings.
Figure 15: Full MBR-L2 curves including all baselines. The empirical baselines generally have higher squared error than the Marformer and Stan variants at larger training sizes in both settings.
Figure 16: Synthetic annotation runtime (minutes). Marformer costs are hypothetical estimates of a shortened validation search plus the measured refit, using the accounting in § 6.6 ; Stan costs cover its measured HMC fit.
Figure 17: Wall-clock runtime (minutes) for Marformer and Stan CPM on LLMRubric as a function of training items. Stan CPM is run with the HMC configuration described in Appendix C . At Itrain=175 , the recorded training times are approximately 22 minutes for Marformer and 162 minutes for Stan CPM. The Marformer uses best-validation checkpoint selection without refitting.
Parameter
Role
Mean
SD
5th %ile
Median
95th %ile
σu
Criterion vector scale
2.127
0.475
1.542
2.014
3.065
σv
Prototype vector scale
11.588
1.713
9.028
11.445
14.564
σuit
Criterion–prototype interaction scale
1.242
0.254
0.874
1.234
1.716
σmeas
Measurement noise
0.383
0.043
0.317
0.381
0.453
κ
Rating scale concentration
5.201
0.417
4.563
5.186
5.958
Appendix
Table 1: Posterior marginal summaries of Stan CPM scale hyperparameters estimated from LLMRubric ( Itrain=175 , weakly informative priors). Mean, SD, and percentiles are computed from HMC posterior samples.
Figure 18: Log-loss and MBR-L2 for three layerwise readout probes, evaluated on 316 missing test ratings at Itrain=500 . Layer 0 is the input representation before attention, and layer 8 is the final layer. The dashed line gives the final model’s performance; the raw Param Lens at layer 8 reproduces this output.
Figure 19: Calibration diagrams for the three readouts (rows) at layers 0, 4 and 8 (columns), on the same 316 missing test ratings. Predicted probabilities for all four rating categories are pooled. The diagonal indicates perfect calibration; titles report smECE and relplot ’s reported calibration-error width from 500 bootstrap resamples. The layer-0 Param Lens predicts 1/4 for every category due to the initialization.
Score-based diffusion models have emerged as prominent deep generative models; however, their application to tabular data remains challenging because their backbones assume fully specified inputs, whereas real-world tabular data often contain missing values. We propose AugMask, a plug-and-play training framework that adapts missing-unaware backbones to incomplete data by separating conditioning from supervision. AugMask 1) constructs numeric inputs via conditional stochastic augmentation using lightweight auxiliary models, and 2) applies denoising supervision only to observed coordinates. In effect, augmented missing entries serve as uncertain conditioning context rather than training targets. We connect this training rule to a Rao--Blackwellized objective and show that marginalizing missing entries yields a variance-weighted sensitivity penalty, discouraging over-reliance on uncertain completions. Across diverse datasets and missingness regimes, AugMask enables standard diffusion-based tabular generators to outperform specialized missing-aware baselines.
Jungkyu Kim, Taeyoung Park, Kibok Lee
Department of Statistics and Data Science, Yonsei University, South Korea.
Matrix completion underlies problems from tabular imputation to causal inference, yet existing tabular foundation models treat it as entry-by-entry prediction, repeating context for every target and discarding the matrix's two-dimensional structure. We introduce MatrixFormer, a pre-trained matrix-native transformer that predicts a full distribution for every missing entry in a single forward pass. MatrixFormer is trained entirely on synthetic low-rank and latent-factor matrices under diverse missingness patterns. Applied zero-shot and with the same model weights, MatrixFormer achieves competitive performance on causal inference panel-data tasks, language-model benchmark-score completion, tabular imputation, and recommendation systems matrix completion. These results position MatrixFormer as a general-purpose foundation model for matrix completion.
Dwaipayan Saha, Jacob Feitelberg, Kyuseong Choi +2
Prior-data fitted networks (PFNs) have recently emerged as a powerful approach for Bayesian prediction tasks, approximating the posterior predictive distribution (PPD) through in-context learning. Despite their strong empirical performance and ability to go beyond point predictions, theoretical understandings of the algorithmic capability of transformers to learn distributions in context are still lacking. Focusing on Gaussian process regression problems, we show by construction that transformers can implement a gradient descent algorithm targeting the posterior predictive mean and variance, followed by nonlinear mappings that yield binned probabilities of PPD. We study the error bounds of the approximated PPD in terms of attention depth and bin resolution. Based on these results, we further demonstrate the key role of normalization and the choice of attention depth in enabling the extrapolation abilities of transformers beyond the pretraining sample size range. We conduct simulations that corroborate our findings, providing insight into the expressivity of PFNs targeting PPDs and how architectural choices may influence generalization capabilities.
Gyeonghun Kang, Changwoo J. Lee, Xiang Cheng
Department of Statistical Science, Duke University, Durham, NC, USA · Department of Electrical and Computer Engineering, Duke University, Durham, NC, USA