WeFT: Weighted Entropy-driven Fine-Tuning for dLLMs
This addresses the problem of inconsistent generation in diffusion language models for researchers and practitioners, offering a novel fine-tuning approach with strong performance gains.
The paper tackled the challenge of applying supervised fine-tuning to diffusion language models by proposing WeFT, a weighted method based on token entropy, which achieved relative improvements of 39% to 83% over standard SFT on reasoning benchmarks.
Diffusion models have recently shown strong potential in language modeling, offering faster generation compared to traditional autoregressive approaches. However, applying supervised fine-tuning (SFT) to diffusion models remains challenging, as they lack precise probability estimates at each denoising step. While the diffusion mechanism enables the model to reason over entire sequences, it also makes the generation process less predictable and often inconsistent. This highlights the importance of controlling key tokens that guide the direction of generation. To address this issue, we propose WeFT, a weighted SFT method for diffusion language models, where tokens are assigned different weights based on their entropy. Derived from diffusion theory, WeFT delivers substantial gains: training on s1K, s1K-1.1, and 3k samples from open-r1, it achieves relative improvements of 39%, 64%, and 83% over standard SFT on four widely used reasoning benchmarks (Sudoku, Countdown, GSM8K, and MATH-500). The code and models will be made publicly available.