Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusions

TL;DR

This paper investigates masked diffusion models (MDMs), revealing their training complexity and demonstrating that adaptive inference strategies can greatly improve their performance on logic puzzles, surpassing autoregressive models in accuracy.

Contribution

The paper provides a theoretical and empirical analysis of MDMs' training challenges and introduces adaptive inference methods that significantly enhance their solving capabilities.

Findings

MDMs train on intractable subproblems compared to ARMs.

Adaptive inference improves MDM accuracy from <7% to ~90%.

Adaptive strategies outperform ARMs with more parameters.

Abstract

In recent years, masked diffusion models (MDMs) have emerged as a promising alternative approach for generative modeling over discrete domains. Compared to autoregressive models (ARMs), MDMs trade off complexity at training time with flexibility at inference time. At training time, they must learn to solve an exponentially large number of infilling problems, but at inference time, they can decode tokens in essentially arbitrary order. In this work, we closely examine these two competing effects. On the training front, we theoretically and empirically demonstrate that MDMs indeed train on computationally intractable subproblems compared to their autoregressive counterparts. On the inference front, we show that a suitable strategy for adaptively choosing the token decoding order significantly enhances the capabilities of MDMs, allowing them to sidestep hard subproblems. On logic puzzles like Sudoku, we show that adaptive inference can boost solving accuracy in pretrained MDMs from $<7$% to $\approx 90$%, even outperforming ARMs with $7\times$ as many parameters and that were explicitly trained via teacher forcing to learn the right order of decoding.