cs.LGOct 8, 2026
SaveMarformer: A Transformer for Predicting Missing Data Distributions
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
Abstract
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 and iteratively refines it through attention to other distributions . 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 at top reveals values only for attributes in (black), where . Thus, our methods receive only the incomplete instance at bottom. In this example, 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 (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 (black). Many different masks 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 is independent of the complete data. Blue nodes show observed values; white nodes are missing. The variables labeled in the figure correspond to 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 . 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 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 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 achieved by the estimated Marformer . Different rows evaluate different Marformer sizes; different columns evaluate different sizes. Each point corresponds to a single marginal prediction on a test instance with , color-coded by the entropy 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 , so they tend to have higher cross-entropy under both models. All graphs for the same domain use the same true generative distribution .
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 (left), and discretizes their values into fixed but irregular bins, yielding a joint distribution over observable discrete variables (right). The positive correlation of carries over to , 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- (c)–(d) on missing ratings in the fixed test corner as the amount of training data increases. Item-generalization results extend to for all models; in judge generalization, Marformer NT stops at 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 and , the largest training sizes with results for all models shown. Probabilities for all 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 , Marformer is well-calibrated; at , 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 , 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 |
|---|---|---|---|---|---|---|
| Criterion vector scale | 2.127 | 0.475 | 1.542 | 2.014 | 3.065 | |
| Prototype vector scale | 11.588 | 1.713 | 9.028 | 11.445 | 14.564 | |
| Criterion–prototype interaction scale | 1.242 | 0.254 | 0.874 | 1.234 | 1.716 | |
| 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 |
Table 1: Posterior marginal summaries of Stan CPM scale hyperparameters estimated from LLMRubric ( , 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 . 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 for every category due to the initialization.