Discovering Non-monotonic Autoregressive Orderings with Variational Inference

TL;DR

This paper introduces an unsupervised, parallelizable method for discovering high-quality autoregressive orderings in language modeling using variational inference and policy gradients, outperforming fixed order baselines.

Contribution

It proposes a novel variational inference framework with a Transformer encoder for learning permutation-based orderings without supervision.

Findings

Discovered orderings are competitive with or better than fixed orders.

Method is context-aware and scalable due to parallelizable design.

Achieved effective learning of sequence orderings from data alone.

Abstract

The predominant approach for language modeling is to process sequences from left to right, but this eliminates a source of information: the order by which the sequence was generated. One strategy to recover this information is to decode both the content and ordering of tokens. Existing approaches supervise content and ordering by designing problem-specific loss functions and pre-training with an ordering pre-selected. Other recent works use iterative search to discover problem-specific orderings for training, but suffer from high time complexity and cannot be efficiently parallelized. We address these limitations with an unsupervised parallelizable learner that discovers high-quality generation orders purely from training data -- no domain knowledge required. The learner contains an encoder network and decoder language model that perform variational inference with autoregressive orders (represented as permutation matrices) as latent variables. The corresponding ELBO is not differentiable, so we develop a practical algorithm for end-to-end optimization using policy gradients. We implement the encoder as a Transformer with non-causal attention that outputs permutations in one forward pass. Permutations then serve as target generation orders for training an insertion-based Transformer language model. Empirical results in language modeling tasks demonstrate that our method is context-aware and discovers orderings that are competitive with or even better than fixed orders.