ALoDLM: Adaptively Looped Diffusion Language Models

ALoDLM. Adaptively Looped Diffusion Language Models. By Liancheng Fang and twelve others, from the University of Illinois Chicago, Amazon AGI, and one other institution. Published on the third of October. Its code repository is worked out: about fifty source files, with a readme and a licence.

Diffusion language models generate quickly because they predict several tokens in parallel. Even so, their practical adoption is held back by a persistent quality gap against autoregressive models of comparable size. The authors attribute that gap to a mismatch between computation and difficulty. Within a partially observed sequence, some unknown tokens are easy to predict, while others need substantially more computation. Existing diffusion language models nevertheless give every unknown position the same computational depth at each denoising step. The easy token and the hard one get identical treatment.

ALoDLM replaces that uniform computation with token-adaptive latent recurrence. At each denoising step, the model refines its latent representations iteratively, and it allocates computation according to how difficult each token is. Depth stops being a constant of the architecture and becomes something each position receives in proportion to its need.

Each step sorts tokens into two groups. Tokens that are ready to commit are fed back as discrete context for the rest of the sequence. Tokens that remain unresolved keep their latent states, and refine them further through additional recurrent passes. Training then has two things to learn together: which token to predict, and how much computation to spend reaching it. To learn them jointly, the authors formulate token-wise computation schedules as latent variables and derive a conditional negative evidence lower bound, or NELBO.

ALoDLM is trained at two scales, one point seven billion parameters and eight billion parameters. Across eleven benchmarks, at both scales, it outperforms all evaluated diffusion language models, and the corresponding autoregressive baselines too, in average benchmark score. It also keeps fast parallel decoding, which gives it a strong quality-efficiency trade-off among the evaluated autoregressive and diffusion models under optimized inference engines.

Spend compute where the difficulty is. Stated that way, the principle is short. ALoDLM's contribution is applying it inside each denoising step, token by token, so that the tokens ready early are committed and fed back as context while the unresolved ones get more passes, and doing so without giving up the parallel decoding that made diffusion attractive in the first place.