Diffusion language models change speed based on task and language

Diffusion Language Models Change Speed Across Tasks and Languages

Computation and LanguageArtificial IntelligenceMachine Learning

Summary

Diffusion language models generate text in multiple steps, and the number of steps they need varies a lot depending on what kind of task they are doing and the language they use. The authors found that copying or lightly changing input text is faster, while creative writing takes more steps. They also saw that languages broken into many small parts need more steps per word. Interestingly, answers to problems that are correct take fewer steps than wrong ones. This means how long the model works can help judge its output.

What this means in practice

  • For language model engineers: Optimize text generation by adjusting computation based on task difficulty and output novelty using estimated decoding effort.
  • For machine translation teams: Estimate and reduce computational costs for heavily tokenized languages by measuring diffusion model decoding steps per token.

Authors

Mohsen Fayyaz, Ali Modarressi, Daniel Israel, Lucas Bandarkar, Nanyun Peng

Abstract

Diffusion language models can write several tokens in one forward pass, and they decide during generation how many passes to spend. Two outputs of the same token length can therefore vary a lot in cost. To understand the underlying factors of this cost variation, we analyze two diffusion models across 10 tasks. We find that the number of forward passes per token depends heavily on the task. Text copied or lightly edited from the input is cheap, while open-ended writing is costly. In DiffusionGemma, generating a token for story writing requires more than three times as many forward passes as for grammar correction. Two output properties explain much of this gap: how much of it is new relative to the input, and how rare its words are within the task. The same holds within a single response when thinking is involved. The final answer is cheap when it repeats the model's reasoning trace and costly when it adds something new. In multilingual generation, we find that highly tokenized languages also require more passes per token, adding to the cost they already pay with more tokens. Finally, cost tracks correctness: on the same math problem, wrong answers take more passes per token than correct ones, and weighting candidate answers by decoding speed beats majority voting by up to 4.2 points at no extra cost. Taken together, our results show that diffusion's computational advantage depends on how much of the output is already determined, while its decoding effort offers a practical signal for making inference more efficient and reliable.