Continuous diffusion language models improve reasoning and code tasks
ELF-REG: Scaling Continuous Diffusion Language Models to Reasoning Tasks
Computation and LanguageMachine Learning
Summary
This paper studies a way to improve language models that guess answers by slowly refining a noisy guess all at once, rather than token by token. The authors found a technique called ELF-REG that uses a teacher model to guide this refinement, which helps the model perform better on math problems and coding challenges. Their model outperforms similar-sized continuous diffusion models and can produce good answers even when stopped early. This approach shows promise for making faster and more accurate models for complex reasoning tasks.
What this means in practice
- •For software developers: Generate and test code snippets with improved accuracy through continuous diffusion models guided by autoregressive supervision.
- •For data science teams: Solve complex mathematical problems with enhanced diffusion language models that refine solutions in parallel.
Authors
Zeyu Michael Li, William Xingxu Chen, Bingshuo Qian, Jiayin Liu, Xiang Cheng
Abstract
Fully continuous diffusion language models (dLMs) denoise continuous representations without intermediate discretization, then decode all response tokens in parallel at the final step. Their performance on challenging reasoning tasks remains less established than that of autoregressive (AR) LLMs and masked dLMs. We scale Embedded Language Flows (ELF) to mathematical reasoning and code generation on GSM8K, MATH-500, HumanEval, and MBPP. We introduce ELF-REG, which improves learning with representation alignment and entanglement (REPA+REG), where a frozen AR teacher supervises intermediate denoiser features and supplies a global representation that is jointly denoised with the response. ELF-REG-L achieves 55.96% pass@1 on GSM8K at 64 network function evaluations (NFE), and 13.39% on MATH-500 and 22.56% on HumanEval at 128 NFE. It outperforms the evaluated comparable-scale dLMs in pass@1 on GSM8K and code, and improves MATH-500 pass@1 from 10.55% for the ELF-L baseline to 13.39% with ELF-REG-L. Without few-step training, the same task-specific checkpoints support strong low-NFE performance through early-stop, which decodes an intermediate clean prediction without completing the denoising trajectory. At 16 NFE, ELF-REG-L reaches 41.21% HumanEval pass@10, outperforming recent continuous dLMs of comparable scale.