MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization
Fuente:
arXiv
Salvato in:
| Autori principali: | , , , , , , , , , , , , |
|---|---|
| Natura: | Preprint |
| Pubblicazione: |
2025
|
| Soggetti: | |
| Accesso online: | |
| Tags: |
Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
|
| _version_ | 1866915574265675776 |
|---|---|
| author | Wang, Chenglong Gan, Yang Zhou, Hang Hu, Chi Mu, Yongyu Song, Kai Yang, Murun Li, Bei Zhang, Chunliang Liu, Tongran Zhu, Jingbo Yu, Zhengtao Xiao, Tong |
| author_facet | Wang, Chenglong Gan, Yang Zhou, Hang Hu, Chi Mu, Yongyu Song, Kai Yang, Murun Li, Bei Zhang, Chunliang Liu, Tongran Zhu, Jingbo Yu, Zhengtao Xiao, Tong |
| contents | Recent advances in diffusion language models (DLMs) have presented a promising alternative to traditional autoregressive large language models (LLMs). However, DLMs still lag behind LLMs in reasoning performance, especially as the number of denoising steps decreases. Our analysis reveals that this shortcoming arises primarily from the independent generation of masked tokens across denoising steps, which fails to capture the token correlation. In this paper, we define two types of token correlation: intra-sequence correlation and inter-sequence correlation, and demonstrate that enhancing these correlations improves reasoning performance. To this end, we propose a Multi-Reward Optimization (MRO) approach, which encourages DLMs to consider the token correlation during the denoising process. More specifically, our MRO approach leverages test-time scaling, reject sampling, and reinforcement learning to directly optimize the token correlation with multiple elaborate rewards. Additionally, we introduce group step and importance sampling strategies to mitigate reward variance and enhance sampling efficiency. Through extensive experiments, we demonstrate that MRO not only improves reasoning performance but also achieves significant sampling speedups while maintaining high performance on reasoning benchmarks. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2510_21473 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization Wang, Chenglong Gan, Yang Zhou, Hang Hu, Chi Mu, Yongyu Song, Kai Yang, Murun Li, Bei Zhang, Chunliang Liu, Tongran Zhu, Jingbo Yu, Zhengtao Xiao, Tong Computation and Language Recent advances in diffusion language models (DLMs) have presented a promising alternative to traditional autoregressive large language models (LLMs). However, DLMs still lag behind LLMs in reasoning performance, especially as the number of denoising steps decreases. Our analysis reveals that this shortcoming arises primarily from the independent generation of masked tokens across denoising steps, which fails to capture the token correlation. In this paper, we define two types of token correlation: intra-sequence correlation and inter-sequence correlation, and demonstrate that enhancing these correlations improves reasoning performance. To this end, we propose a Multi-Reward Optimization (MRO) approach, which encourages DLMs to consider the token correlation during the denoising process. More specifically, our MRO approach leverages test-time scaling, reject sampling, and reinforcement learning to directly optimize the token correlation with multiple elaborate rewards. Additionally, we introduce group step and importance sampling strategies to mitigate reward variance and enhance sampling efficiency. Through extensive experiments, we demonstrate that MRO not only improves reasoning performance but also achieves significant sampling speedups while maintaining high performance on reasoning benchmarks. |
| title | MRO: Enhancing Reasoning in Diffusion Language Models via Multi-Reward Optimization |
| topic | Computation and Language |
| url | https://arxiv.org/abs/2510.21473 |