Looped Diffusion Language Models

Problem
This paper addresses the gap in the effective design of transformer architectures specifically for Masked Diffusion Models (MDMs) in language modeling. While MDMs have shown promise as alternatives to autoregressive models, the optimization of their architecture remains underexplored. This work is presented as a preprint and has not yet undergone peer review.

Method
The authors introduce LoopMDM (Looped Masked Diffusion Model), which incorporates a novel architecture that selectively loops early-middle transformer layers during training. This approach achieves a depth-scaling effect without increasing the number of parameters. The method allows for varying the number of loops during inference, enabling flexible compute scaling. The training process is designed to optimize the efficiency of the model, resulting in significant reductions in training FLOPs. The authors do not disclose specific details about the training compute or the exact architecture used, but they emphasize the importance of the looping mechanism in enhancing model performance.

Results
LoopMDM demonstrates substantial improvements over baseline MDMs across multiple pre-training corpora. Notably, it matches the performance of same-size MDMs while requiring up to 3.3 times fewer training FLOPs. On reasoning benchmarks, LoopMDM outperforms these baselines, achieving an increase of up to 8.5 points on the GSM8K dataset. Furthermore, it surpasses deeper non-looped MDMs that were trained with comparable per-step compute, indicating that the selective looping mechanism is more effective than traditional depth scaling. The authors also report that adaptively adjusting the number of loops during the sampling process yields additional gains in compute efficiency without sacrificing performance.

Limitations
The authors acknowledge that their approach may not generalize to all types of MDMs or other architectures outside the transformer framework. They do not discuss potential limitations related to the scalability of the looping mechanism in extremely large models or the impact of varying loop counts on model stability. Additionally, the paper does not provide extensive ablation studies to isolate the effects of the looping mechanism from other architectural choices.

Why it matters
The implications of this work are significant for the field of language modeling, particularly in the context of efficiency and performance trade-offs. By demonstrating that selective looping can enhance both training efficiency and model performance, this research opens avenues for further exploration of architectural innovations in MDMs. The ability to scale inference-time compute flexibly could lead to more efficient deployment of language models in resource-constrained environments. Moreover, the findings encourage a reevaluation of depth scaling strategies in transformer architectures, potentially influencing future designs in both MDMs and other model types.

Authors: Sanghyun Lee, Chunsan Hong, Seungryong Kim, Jonghyun Lee, Jongho Park, Dongmin Park
Source: arXiv cs.LG
URL: https://arxiv.org/abs/2605.26106v1
arXiv ID: 2605.26106