Real decisions are made under incomplete information. If we observe only some of the random variables we need, we can predict the others. The 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ᵢ) and iteratively refines it through attention to other distributions p(Xⱼ). 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.

Problem Setting: Conditional Marginals for Decision-Making Under Incomplete Information

Incomplete data is the rule, not the exception, in scientific and applied settings. A recommender system has a user's opinions of only a few movies. A standardized testing agency has only a few graders' opinions of each essay. The incomplete data may be acquired adaptively: a doctor does just enough tests to choose a treatment confidently. An engineer runs cheap simulations to decide which candidate molecules or hyperparameter settings are worthy of further measurements. An educational tester or job interviewer probes adaptively for strengths and weaknesses, using a limited budget. An evaluator of an AI system adaptively selects examples based on observed performance, using the results to predict performance on other examples. In each case, the decision-maker needs not a point estimate of the missing values but a predictive distribution over them given the available evidence. With such distributions, one can quantify residual uncertainty, predict the risk of a possible action, or determine whether gathering more data before acting would be worth the cost.

This paper focuses on marginal predictions. The goal is to train a neural network that can predict the marginal distribution of any single missing value conditioned on any subset of the remaining values. These conditional marginals are fitted directly by conditional maximum likelihood, a frequentist approach that requires no prior over the model parameters. Such a model is directly useful for deciding among real-world actions and for choosing what additional information (missing values) to acquire before this decision. The architecture is generic and offers fast, amortized, trainable inference of all the individual missing values in parallel. Such a model can jointly impute multiple missing values, if needed, by sampling them sequentially. Thus, the collection of conditional distributions contains as much information as a joint distribution—but in a form designed to support conditional inference.

Classical Expectation-Maximization proceeds in the opposite direction: from joint to conditional. It first obtains a joint distribution over all values by fitting a full generative model to the incomplete data, using the maximum likelihood principle or the maximum a posteriori principle if a prior is available. One then performs inference as needed to extract the desired conditional marginals from this joint distribution. Unfortunately, EM or MAP-EM requires domain expertise to specify the generative model family and any prior; can encounter local optima when fitting the model parameters; and may require intractable inference, necessitating variational or sampling-based approximations that can be slow or inaccurate. A more sophisticated variant is Bayesian posterior imputation, which integrates out the posterior uncertainty about the generative model's parameters, typically by using MCMC to (approximately) sample them jointly with the missing values. Training and inference is thus expensive to do accurately, and model specification is still challenging. This is why we propose to abandon the hypothesized generative model---even when it is correct!---and simply fit all the conditional marginals directly using a generic neural model with enough capacity to do the job.

Marformer Architecture: A Transformer for Conditional Marginal Prediction

The Marformer architecture key properties are:

  1. Observation of values. Each input to the Marformer specifies a set of random variables together with their values, some of which may be missing.
  2. Prediction of parameters. The Marformer outputs marginal distributions for all variables, conditioned on the observations in the input. More precisely, for each variable, it outputs a parameter vector that will be interpreted as specifying a predictive distribution over that variable's domain.
  3. Conditional log-likelihood maximization. The training objective is conditional log-likelihood. That is, the predicted marginal over an unknown quantity should assign a high log-probability (or log-density) to its true value. This objective encourages predictions to match the true conditional distributions, which can subsequently be used to evaluate decisions under any task-specific loss.
  4. Amortized, single-pass inference. Once trained, the Marformer predicts all missing marginals simultaneously in a single forward pass, sharing computation across the entire set of missing values. Its architecture generalizes BERT, which specifically predicts conditional marginals over all missing words in a natural-language sentence. As in BERT, each variable can attend to the evolving predictions of the other variables as their representations are transformed layer by layer.
  5. End-to-end differentiable training. Inference is fully defined by this forward pass. Thus, training the Marformer to maximize conditional log-likelihood is explicitly optimizing the parameters to make inference accurate. In contrast, classical approaches train a joint model without considering the errors that may later be introduced by approximate inference.
  6. Support for relational domains. Many machine learning methods train and test on many partially observed samples that are assumed i.i.d. However, a relational task domain may have only a single partially observed sample with very many interrelated variables. We extend the Marformer architecture to also handle such domains, using a relational attention mechanism whereby related variables know their relationships and can learn to attend to one another.

To supervise the prediction of missing data, where do we acquire the true values? The self-supervised learning recipe randomly masks out some variables, and trains a model to predict their masked values. We start with incomplete data, and predicting missing values is our actual goal. Yet we can follow the same masking recipe—randomly mask out some of the observed values and train the model to predict them from the remaining observed values. We hope that the Marformer model will generalize from predicting these masked values to predicting the originally missing values. This is reasonable to expect provided that the originally missing variables are missing completely at random (MCAR) and we mask additional variables using a missing-at-random (MAR) mechanism.

Domain 1: Graphical Models

We evaluate across three synthetic domains with missing data. Domain 1 uses Bayesian networks. The true generative distribution is a Bayesian network, and we consider a missing completely at random (MCAR) missingness mechanism. Inference under the true model proceeds by exact variable elimination or sampling-based methods. The Marformer is configured with a suitable architecture for the graph structure, and we compare predictive performance against classical methods such as CPT enumeration, likelihood weighting, and Gibbs sampling. The Marformer can match or outperform these classical methods, even when they are given the true model family and prior that generated the synthetic data. Runtime is substantially faster for the Marformer, as it avoids the iterative sampling or expensive inference required by classical methods.

Domain 2: Discretized Multivariate Gaussians

Domain 2 uses discretized multivariate Gaussians. The true generative distribution is a multivariate Gaussian, discretized into a finite number of bins per variable. We consider both MCAR and missing-at-random (MAR) missingness mechanisms. Under the true model, inference can be performed by computing conditional Gaussian distributions. The Marformer is trained to predict conditional marginals for each variable given observed values. Experimental results show that the Marformer matches or outperforms classical methods such as multivariate Gaussian imputation and conditional Gaussian prediction. At the largest training sizes, the Marformer achieves lower prediction error than baselines, while requiring orders of magnitude less computation time.

Domain 3: Structured Annotation Data

Domain 3 uses structured annotation data. The true generative distribution arises from a process where annotators label items with multiple possible categories. We construct a dataset with missing annotations, and evaluate under MCAR missingness. Classical baselines include pooled unigram, naive Bayes, and log-linear models. The Marformer outperforms the evaluated baselines at the largest training size. On this real annotation dataset, the Marformer achieves better-calibrated predictive distributions and faster inference than generative approaches that require MCMC or variational approximation.

Applying Marformer to Real Annotation Data

We also evaluate on a real annotation dataset, where the Marformer outperforms the evaluated baselines at the largest training size. The dataset consists of text annotations with multiple possible labels per item. The Marformer is configured to predict conditional marginals over the missing labels given observed annotations. Results show that the Marformer achieves better predictive performance and calibration than baselines, while being substantially faster at inference time. The method generalizes from predicting masked values at training time to predicting the actually missing values at test time, provided that the missingness mechanism is MCAR or approximately MAR.

Related Work and Limitations

Other related work on missing data imputation and conditional distribution estimation includes Bayesian approaches, variational methods, and deep learning-based imputation. The Marformer differs from these by predicting conditional marginals directly, without modeling the full joint distribution, and by requiring no domain knowledge of the data-generating process. Limitations include reliance on the MCAR assumption for the masking-based training recipe, and the fact that the method does not model dependencies between multiple missing values beyond what is captured by the chain of conditional marginals. Future work will generalize to training Marformers on MAR data using a more complex training recipe.