Marformer: A Transformer for Predicting Missing Data Distributions

AI summary

Overview

• Research area: Machine learning for incomplete data — specifically, a neural architecture (the Marformer) for predicting conditional marginal distributions over missing variables, positioned against classical approaches such as EM, MAP-EM, and Bayesian posterior imputation.
• Technical level: Advanced. The paper is notation-heavy, defining random variables, subcollections, missingness mechanisms (MCAR, MAR, IMAR, MNAR), and conditional maximum likelihood estimation, and it draws on decision theory (Bayes risk, Value of Information).
• Scope in one sentence: The paper proposes and evaluates a BERT-like Transformer that directly fits all conditional marginals of a set of random variables from partially observed data, without ever modeling the full joint distribution.

What This Paper Is About

Real decisions are made with incomplete information, and what a decision-maker needs is not a point estimate but a predictive distribution over the missing values, so that risks, and the value of collecting more data, can be quantified. Classical methods get there by first fitting a full generative joint model (via EM, MAP-EM, or Bayesian posterior imputation) and then extracting conditional marginals through inference, which requires domain expertise, can hit local optima, and may be slow or intractable. The Marformer instead fits every conditional marginal directly by conditional maximum likelihood, producing all predictions in a single forward pass and requiring no knowledge of the data-generating process.

Key Contributions

1. The Marformer architecture. A Transformer that generalizes BERT to arbitrary random variables: each variable gets a hidden-vector representation of its distribution p(X_i), and these representations are iteratively refined through attention to the other variables' distributions p(X_j). It outputs a parameter vector per variable that specifies a predictive distribution over that variable's domain.
2. Direct conditional maximum likelihood training. Rather than training a joint model and hoping approximate inference behaves, the Marformer is trained end-to-end to maximize conditional log-likelihood, so training explicitly optimizes the inference that will actually be used.
3. A masking-based self-supervised recipe for incomplete data. Starting from incomplete data (unlike standard BERT-style pretraining, which starts from complete data), the method randomly masks additional observed values and trains the model to predict them, with a theoretical justification given for when this transfers to the originally missing values.
4. Extension to relational domains and formal missingness analysis. The architecture adds a relational attention mechanism (following Shaw et al., 2018) so that a single partially observed instance with many interrelated variables can be handled, and the paper introduces the new notion of individually missing at random (IMAR), a weaker ignorability condition than MAR that suffices for predicting individual missing variables.

Main Findings

• Competitive with classical methods that know the truth. Across three synthetic domains — 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 data.
• Wins on a real dataset at scale. On a real annotation dataset, the Marformer outperforms the evaluated baselines at the largest training size.
• Substantially faster. In both synthetic and real evaluations, the Marformer is substantially faster than the evaluated generative baselines.
• Amortized single-pass prediction. Once trained, all missing marginals are predicted simultaneously in one forward pass, sharing computation across the entire set of missing values, rather than requiring per-query inference.
• Chaining recovers joint predictions. Because the model supplies all conditionals, conditional joint predictions can be recovered by chaining (as in orderless NADE and XLNet), so the collection of conditionals carries as much information as a joint distribution while staying useful for conditional inference.
• Enables decision-theoretic quantities. The predicted marginals feed directly into Bayes-optimal point prediction (Equation 1), Bayes risk of actions (Equation 3), Value of Information (Equation 4), and an ROI heuristic of (VOI − COI)/COI for active feature acquisition.
• Missingness assumptions matter. The masking recipe is justified when the original missingness is MCAR and the masking mechanism is MAR; a constructed counterexample shows that even under MAR data with MCAR masking, observability bias can corrupt training examples.
• Exploratory MAR results. Exploratory MAR evaluations appear in Appendix A, where the transfer justification does not automatically apply; the real annotation dataset is described as naturally sparse with MCAR an unverified assumption.
• Calibration is examined in §6.6, though specific calibration figures are not reported in the available text.

Note: the provided paper content is truncated before the experiment sections, so specific numerical results, dataset sizes, and named baseline systems are not reported here.

Methodology in Plain English

The authors frame prediction from incomplete data as: given whatever values you happen to have observed, output a probability distribution for each variable you have not observed. Classical statistics would build a single joint model of everything and then do algebra or simulation to extract those conditional distributions; the authors skip the joint entirely and train a neural network to output the conditionals directly, scoring it by how much probability it assigns to the true held-out values (conditional log-likelihood).

To get training signal from incomplete data, they borrow the masked-language-model trick: take the values that were observed, hide some of them at random, and ask the model to predict the hidden ones from the remaining ones. If you do this with many different random masks, the model sees a wide variety of "some values known, some unknown" situations, which is exactly the shape of the problem at test time. The paper works out precisely when this is sound: the transfer from masked values to genuinely missing values holds provided the original missingness is missing completely at random and the artificial masking is missing at random.

The architecture is BERT-like in spirit. Each variable is represented by a vector describing its predicted distribution, and layers of attention let variables look at each other's evolving predictions — so the prediction for one missing value can be informed by the current predictions for the others. A relational attention variant lets variables encode known relationships between them, which is aimed at domains where there is only one big instance rather than many i.i.d. samples.

Why This Matters

Impact on research. The paper argues that one should abandon the hypothesized generative model even when it is correct, because direct conditional fitting avoids model misspecification, bad local optima in EM, and expensive or approximate inference. It also introduces IMAR, a new ignorability condition that is weaker than MAR but sufficient for the individual-variable prediction task, plus a concrete counterexample showing that MAR data with MCAR masking is not automatically safe for training. If validated at scale, this shifts missing-data work away from bespoke generative modeling toward generic, amortized neural conditional fitting.

Real-world applications (the paper's own motivating list):

• Recommender systems, which observe a user's opinions of only a few items.
• Standardized testing and grading, where only a few graders' opinions exist per essay.
• Clinical decision-making, where a doctor runs just enough tests to choose a treatment confidently.
• Adaptive scientific and engineering pipelines, where cheap simulations decide which molecules or hyperparameter settings deserve further measurement.
• AI system evaluation and educational testing or job interviewing, where examples or probes are selected adaptively under a budget (the paper cites Li et al., 2024 for the evaluation case).

Industry relevance. Value of Information and active feature acquisition are the core of budgeted decision pipelines, and the Marformer's single-forward-pass prediction means the risk of each candidate action, and the value of acquiring each additional attribute, can be estimated in parallel rather than by Monte Carlo over chained conditionals. The paper notes that when a real-world loss is itself treated as a variable the model has been trained to predict, one call to the Marformer suffices to rank the Bayes risks of multiple candidate actions. Substantially faster prediction than the evaluated generative baselines also matters for online training settings, where the model is retrained as new observations arrive.

Future Directions

1. Training on MAR data. The authors state that a sequel paper will generalize the training recipe to MAR data — the same setting where EM works — using the Marformer itself to impute missing values, following the EM-style pattern.
2. Learning an information-acquisition policy. Also planned for a sequel: training a Marformer-based policy for incrementally acquiring conditioning variables, by reinforcement learning, so the policy can look farther into the future than the myopic ROI heuristic of (VOI − COI)/COI.
3. Jointly training the policy and the conditionals it consults, on the data the policy itself gathers (the authors cite Li and Oliva, 2020 for this direction).
4. Open questions around assumptions. The real annotation dataset has MCAR as an unverified assumption, the MAR evaluations in Appendix A are described as exploratory, and the paper leaves open how well the masking recipe extends when those mechanisms are violated.

Target Audience

Readers who will benefit most are machine learning researchers working on missing data, probabilistic inference, or transformer architectures; statisticians interested in conditional maximum likelihood and missingness mechanisms (MCAR, MAR, MNAR, and the new IMAR); and applied researchers in decision-making under uncertainty — clinical decision support, adaptive experimentation, recommender systems, and evaluation — who need calibrated per-variable predictive distributions fast. Because the paper is notation-heavy and assumes familiarity with EM, Bayes risk, and Value of Information, it is best suited to readers comfortable with graduate-level probability and machine learning.

Authors’ 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 $p(X_i)$ and iteratively refines it through attention to other distributions $p(X_j)$. 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.