Technical Summary: Order-Agnostic Autoregressive Modelling with Missing Data

Problem Statement

Real-world learning systems frequently operate under partial observability, where features are missing due to sensor faults, acquisition costs, or privacy constraints. Effective systems in these domains must address three core challenges: (i) learning from incomplete data, (ii) imputing missing values at inference time, and (iii) deciding which unobserved variables to acquire when observations are costly (Active Feature Acquisition).

While deep generative models theoretically offer a unified approach by modeling the full joint distribution, existing methods face significant limitations. Variational Autoencoders (VAEs) rely on approximate inference, limiting conditional accuracy. Generative Adversarial Networks (GANs) often lack tractable likelihoods and principled conditional densities. Diffusion models, while producing high-quality samples, require iterative sampling chains and struggle to condition precisely on arbitrary subsets of observed variables. Classical Autoregressive Models (ARMs) provide tractable likelihoods but rely on a fixed generation order, preventing conditioning on arbitrary subsets of variables.

Methodology: MO-ARM

The authors introduce Missingness-Aware Order-Agnostic Autoregressive Models (MO-ARM), a framework that reinterprets Order-Agnostic Autoregressive Models (OA-ARMs) through the lens of missing data.

1. Theoretical Insight: OA-ARMs and MCAR

The paper first establishes a theoretical connection between standard OA-ARM training and Missing Completely At Random (MCAR) imputation. OA-ARMs are trained on random permutations of variables, effectively treating the generation order as a latent variable. The authors demonstrate that the standard training objective, which involves predicting a variable based on a randomly selected subset of previously revealed variables, implicitly simulates an MCAR corruption process. Consequently, OA-ARMs trained solely on fully observed data naturally generalize to imputation tasks under MCAR mechanisms without modification.

2. The MO-ARM Framework

Building on this insight, the authors propose a principled framework for training OA-ARMs directly on incomplete datasets under general missingness mechanisms (spanning MCAR to Missing Not At Random, or MNAR).

• Probabilistic Formulation: The model factorizes the joint distribution of observed data ($x_{z_o}$), missing data ($x_{z_m}$), and the missingness mask ($z_o$) as:
  $$p_{\theta, \phi}(x_{z_o}, x_{z_m}, z_o) = p_\theta(x_{z_o}) p_\theta(x_{z_m} | x_{z_o}) p_\phi(z_o | x_{z_o}, x_{z_m})$$
  Here, $p_\theta$ represents the generative model (an OA-ARM), and $p_\phi$ is a learned missingness mechanism.
• Training Objective: The model maximizes the observed-data log-likelihood using a variational lower bound (ELBO).
  • For MCAR, the missingness term is independent of the data and can be omitted.
  • For MNAR, the model explicitly learns the missingness process $p_\phi$ using an element-wise Bernoulli model parameterized by a neural network. This allows the model to account for dependencies between missingness and unobserved values.
• Imputation Strategy:
  • MCAR: Imputation is performed by marginalizing over possible unmasking trajectories of the missing variables using Monte Carlo sampling.
  • MNAR: The model employs self-normalized importance weighting. Completed samples are reweighted based on the likelihood of the observed mask under the learned missingness model $p_\phi$, ensuring that completions consistent with the observed missingness pattern are prioritized.

3. Active Information Acquisition

MO-ARM leverages its amortized conditional density estimation to perform Active Feature Acquisition (AFA). By estimating the conditional mutual information $I(x_\Phi; x_{z_i} | x_{z_{<i}})$ between a target variable $x_\Phi$ and a candidate feature $x_{z_i}$ given current observations, the model can sequentially select the most informative missing variables to acquire, maximizing downstream predictive performance.

Key Contributions

1. Reinterpretation of OA-ARMs: The paper demonstrates that standard OA-ARM training implicitly performs MCAR imputation, explaining their robust out-of-sample performance in missing-data settings.
2. Principled Training on Incomplete Data: The authors introduce the first framework to train OA-ARMs directly on incomplete datasets under general missingness mechanisms (MCAR and MNAR) while preserving exact likelihoods and enabling efficient conditional sampling.
3. Unified Framework for Imputation and Acquisition: The model unifies missing data imputation and active information acquisition, utilizing the same amortized conditional distributions for both tasks.
4. Empirical Validation: The proposed MO-ARM is evaluated across real-world benchmarks, including tabular data (UCI datasets) and image inpainting (CelebA-HQ).

Experimental Results

The authors evaluate MO-ARM against established baselines, including MIWAE, not-MIWAE (for MNAR), VAEAC, DiffPuter, and classical methods like MissForest and MICE.

• Tabular Imputation: Across seven UCI datasets with varying feature types (continuous and categorical), MO-ARM consistently outperforms established generative baselines under both MCAR and MNAR missingness mechanisms. Notably, under MNAR, MO-ARM outperforms not-MIWAE, the only other baseline explicitly designed for this mechanism.
• Image Inpainting: On the CelebA-HQ dataset, MO-ARM demonstrates superior performance in preserving identity and coarse structure compared to RePaint (a diffusion-based baseline), achieving higher PSNR scores, particularly as the missing percentage increases (up to 98%).
• Active Acquisition: In sequential feature acquisition tasks, MO-ARM matches or outperforms baselines in reducing prediction error as more features are acquired, demonstrating the utility of its conditional distributions for decision-making.

Significance and Claims

The paper claims that MO-ARM constitutes a powerful and previously underappreciated tool for reasoning under incomplete data. The authors argue that order-agnostic autoregressive models offer a practical and principled alternative to latent-variable (VAE) and diffusion-based approaches for missing-data problems. Key advantages highlighted include:

• Tractable Likelihoods: Unlike GANs or approximate VAEs, MO-ARMs provide exact likelihoods.
• Arbitrary Conditioning: Unlike classical ARMs, they support conditioning on arbitrary subsets of observed variables.
• Efficiency: They can be queried repeatedly as observations are revealed, supporting sequential decision-making without the need for long iterative sampling chains required by diffusion models.

The authors acknowledge limitations, noting that exact generation remains expensive in high-dimensional settings (though mitigated by blocked unmasking) and that performance degrades as observed context becomes extremely sparse. They also note that while the framework extends to MNAR, this regime remains more challenging than MCAR and warrants further study.