Order-Agnostic Autoregressive Modelling with Missing Data

TL;DR

This paper introduces MO-ARM, a novel order-agnostic autoregressive model that effectively handles incomplete data, performs imputation, and actively selects missing variables, outperforming existing methods across benchmarks.

Contribution

It provides a new principled framework for training autoregressive models on incomplete data and leverages their capabilities for active missing data acquisition.

Findings

Implicit imputation performance under high missingness.

First principled training method for incomplete datasets.

Outperforms established imputation baselines on benchmarks.

Abstract

Order-Agnostic autoregressive models have demonstrated strong performance in deep generative modeling, yet their use in settings with incomplete data remains largely unexplored. In this work, we reinterpret them through the lens of missing data. First, we show that their standard training procedure on fully observed data implicitly performs imputation under a missing completely at random mechanism, resulting in robust out-of-sample imputation performance in settings with high missingness. Second, we introduce the first principled framework for training them directly on incomplete datasets under general missingness mechanisms. Third, we leverage their amortized conditional density estimation to perform active information acquisition, i.e., sequentially selecting the most informative missing variables for downstream prediction or inference. Across a suite of real-world benchmarks, our Missingness-Aware Order-Agnostic Autoregressive Model (MO-ARM) consistently outperforms established imputation baselines.