Training and Inference on Any-Order Autoregressive Models the Right Way

TL;DR

This paper improves Any-Order Autoregressive Models (AO-ARMs) by reducing model redundancy and better weighting training loss, resulting in state-of-the-art performance in arbitrary conditional inference across multiple data domains.

Contribution

It introduces a method to eliminate redundancy in AO-ARMs and emphasizes training on more frequently used conditionals, enhancing performance without losing tractability.

Findings

Achieved state-of-the-art likelihoods on Text8, CIFAR10, ImageNet32, and tabular data.

Reduced redundancy in probabilistic modeling of AO-ARMs.

Improved performance by upweighting training on frequently evaluated conditionals.

Abstract

Conditional inference on arbitrary subsets of variables is a core problem in probabilistic inference with important applications such as masked language modeling and image inpainting. In recent years, the family of Any-Order Autoregressive Models (AO-ARMs) -- closely related to popular models such as BERT and XLNet -- has shown breakthrough performance in arbitrary conditional tasks across a sweeping range of domains. But, in spite of their success, in this paper we identify significant improvements to be made to previous formulations of AO-ARMs. First, we show that AO-ARMs suffer from redundancy in their probabilistic model, i.e., they define the same distribution in multiple different ways. We alleviate this redundancy by training on a smaller set of univariate conditionals that still maintains support for efficient arbitrary conditional inference. Second, we upweight the training loss for univariate conditionals that are evaluated more frequently during inference. Our method leads to improved performance with no compromises on tractability, giving state-of-the-art likelihoods in arbitrary conditional modeling on text (Text8), image (CIFAR10, ImageNet32), and continuous tabular data domains.