Understanding and Accelerating the Training of Masked Diffusion Language Models

TL;DR

This paper analyzes why masked diffusion models learn slowly and proposes a simple training strategy that accelerates training by up to four times without sacrificing performance.

Contribution

It identifies the locality bias as a key factor in slow MDM training and introduces bell-shaped time sampling to significantly speed up learning.

Findings

MDMs trained with the new method reach similar NLL up to 4x faster.

Faster improvements observed in generative and zero-shot perplexity.

Enhanced downstream task performance on multiple benchmarks.

Abstract

Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models (ARMs) for language modeling. However, MDMs are known to learn substantially more slowly than ARMs, which may become problematic when scaling MDMs to larger models. Therefore, we ask the following question: how can we accelerate standard MDM training while maintaining its final performance? To this end, we first provide a detailed analysis of why MDM training is slow. We find that the main factor is the locality bias of language: the predictive information for a token is concentrated in nearby positions. We further investigate how this bias slows learning and suggest a simple yet effective remedy: bell-shaped time sampling as a training strategy. Notably, MDMs trained with our training recipe reach the same validation negative log-likelihood (NLL) up to $\sim4\times$ faster than standard training on One Billion Word Benchmark (LM1B). We also show faster improvements in generative perplexity, zero-shot perplexity, and downstream task performance on various benchmarks.