Elamur Research⌕
☰
home / papers / 2026 / Energy Matching based Preference Learning for Diffusion Language Models

Energy Matching based Preference Learning for Diffusion Language Models

Paper · 2026 · Shiv Shankar

Summary

Policy-gradient RL methods such as GRPO depend on sequence likelihoods that factorize token by token. Diffusion language models (dLLMs) generate text non-autoregressively, so these likelihoods are intractable. ELBO-based approximations introduce bias, and off-policy importance weights become unstable over long reasoning chains. The paper proposes EMBR (Energy Matching Based Realignment). It treats KL-regularized RL as matching the policy to a reward-tilted, Boltzmann-style target distribution over the reference model. Because the objective needs no importance weights or on-policy rollouts, it can also learn from static reward-labelled or preference data, much like DPO.

Three variants are defined:

  • EMBR-E, the basic energy matching objective;
  • EMBR-C, a contrastive version;
  • EMBR-U, a derived upper bound suited to dLLMs.

All variants fine-tune LLaDA on GSM8K, MATH500 and HumanEval. Online, EMBR matches or beats GRPO-style methods (UniGRPO, d1, wd1). Offline, every variant outperforms ELBO-based DPO. The work gives dLLMs a practical post-training route that also works offline.

Key points

  • Base model and benchmarks: LLaDA was evaluated on GSM8K, MATH500 and HumanEval at several generation lengths, up to 512 tokens.
  • Online results: EMBR-U was generally the strongest variant. Its scores at the longest length were 83.8 on GSM8K, 39.4 on MATH500 and 48.3 on HumanEval. The corresponding d1 scores were 82.1, 40.2 and 47.1.
  • Offline data: The offline dataset was built from outputs of the d1 model sampled during its online training, then scored with normalized rewards. This was needed because the 350k-sample preference set from Zhu et al. (2025) has not been released.
  • Offline results: EMBR-U scored 68.5 on GSM8K, 37.0 on MATH500 and 39.8 on HumanEval, against DPO's 58.2, 31.6 and 35.1. VRPO numbers are not directly comparable because they used different training data. EMBR-U averaged about 2 points above the other EMBR variants.
  • Limitations: The method uses only sequence-level rewards. It was not tested on long-horizon planning tasks. Offline training still trails online training.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Fields, topics and concepts

📚 Shelves 4The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 3,320 works

Cite

Cite as Shiv Shankar. 2026. Energy Matching based Preference Learning for Diffusion Language Models. Proceedings of the 19th Conference of the European Chapter of the Association for Computational Linguistics (Volume 4: Student Research Workshop).

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.eacl-srw.57

@inproceedings{shankar2026energy,
  title = {{Energy Matching based Preference Learning for Diffusion Language Models}},
  author = {Shiv Shankar},
  booktitle = {Proceedings of the 19th Conference of the European Chapter of the Association for Computational Linguistics (Volume 4: Student Research Workshop)},
  year = {2026},
  doi = {10.18653/v1/2026.eacl-srw.57},
  url = {https://doi.org/10.18653/v1/2026.eacl-srw.57}
}

References

45 references, 16 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Mohammad Gheshlaghi Azar, Zhaohan Daniel Guo, Bilal Piot, Remi Munos, Mark Rowland, Michal Valko, and Daniele Calandriello. 2024. A general theoretical paradigm to understand learning from human preferences. In International Conference on Artificial Intelligence and Statistics, pages 4447–4455. PMLR…
  2. Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, et al. 2006. A tutorial on energy-based learning. Predicting structured data, 1(0). find
  3. Jonas Gehring, Kunhao Zheng, Jade Copet, Vegard Mella, Quentin Carbonneaux, Taco Cohen, and Gabriel Synnaeve. 2024. Rlef: Grounding code llms in execution feedback with reinforcement learning. arXiv preprint arXiv:2410.02089.
  4. Shansan Gong, Ruixiang Zhang, Huangjie Zheng, Jiatao Gu, Navdeep Jaitly, Lingpeng Kong, and Yizhe Zhang. 2025. Diffucoder: Understanding and improving masked diffusion models for code generation. arXiv preprint arXiv:2506.20639. — [arXiv:2506.20639](https://arxiv.org/abs/2506.20639) find
  5. Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. 2025. Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948.
  6. Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. 2018. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pages 1861–1870. PMLR. find
  7. Daniel G Horvitz and Donovan J Thompson. 1952. A generalization of sampling without replacement from a finite universe. Journal of the American statistical Association, 47(260):663–685. find
  8. Edwin T Jaynes. 1979. Concentration of distributions at entropy maxima. ET Jaynes: Papers on probability, statistics and statistical physics, page 315. find
  9. Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al. 2021. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168.
  10. Chunlin Ji and Haige Shen. 2019. Stochastic variational inference via upper bound. Preprint, arXiv:1912.00650. — [arXiv:1912.00650](https://arxiv.org/abs/1912.00650) find
  11. Joost CF de Winter, Dimitra Dodou, and Yke Bauke Eisma. 2024. System 2 thinking in openai’s o1-preview model: Near-perfect performance on a mathematics exam. Computers, 13(11):278. find
  12. Xiang Deng, Yu Gu, Boyuan Zheng, Shijie Chen, Sam Stevens, Boshi Wang, Huan Sun, and Yu Su. 2023. Mind2web: Towards a generalist agent for the web. Advances in Neural Information Processing Systems, 36:28091–28114.
  13. Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. 2023. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.06767.
  14. Zeyu Jia, Alexander Rakhlin, and Tengyang Xie. 2025. Do we need to verify step by step? rethinking process supervision from a theoretical perspective. arXiv preprint arXiv:2502.10581. — [arXiv:2502.10581](https://arxiv.org/abs/2502.10581) find
  15. Muhammad Khalifa, Hady Elsahar, and Marc Dymetman. 2020. A distributional approach to controlled text generation. arXiv preprint arXiv:2012.11635. — [arXiv:2012.11635](https://arxiv.org/abs/2012.11635) find
  16. Sergey Levine. 2018. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv:1805.00909. — [arXiv:1805.00909](https://arxiv.org/abs/1805.00909) find
  17. Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, and Zhi-Quan Luo. 2023. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. arXiv preprint arXiv:2310.10505. find
  18. Kawin Ethayarajh, Winnie Xu, Niklas Muennighoff, Dan Jurafsky, and Douwe Kiela. 2024. Kto: Model alignment as prospect theoretic optimization. arXiv preprint arXiv:2402.01306. find
  19. Yannis Flet-Berliac, Nathan Grinsztajn, Florian Strub, Bill Wu, Eugene Choi, Chris Cremer, Arash Ahmadian, Yash Chandak, Mohammad Gheshlaghi Azar, Olivier Pietquin, et al. 2024. Contrastive policy gradient: Aligning llms on sequence-level scores in a supervised-friendly fashion. arXiv preprint arXiv… find
  20. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. 2023. Let’s verify step by step. arXiv preprint arXiv:2305.20050.
  21. Zicheng Lin, Tian Liang, Jiahao Xu, Xing Wang, Ruilin Luo, Chufan Shi, Siheng Li, Yujiu Yang, and Zhaopeng Tu. 2024. Critical tokens matter: Token-level contrastive estimation enhence llm’s reasoning capability. arXiv preprint arXiv:2411.19943. — [arXiv:2411.19943](https://arxiv.org/abs/2411.19943) find
  22. Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin. 2025. Understanding r1-zero-like training: A critical perspective. arXiv preprint arXiv:2503.20783.
  23. Xiaohang Tang, Rares Dolga, Sangwoong Yoon, and Ilija Bogunovic. 2025. wd1: Weighted policy optimization for reasoning in diffusion language models. Preprint, arXiv:2507.08838. — [arXiv:2507.08838](https://arxiv.org/abs/2507.08838) find
  24. Trung Quoc Luong, Xinbo Zhang, Zhanming Jie, Peng Sun, Xiaoran Jin, and Hang Li. 2024. Reft: Reasoning with reinforced fine-tuning. arXiv preprint arXiv:2401.08967, 3.
  25. Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, et al. 2025. Spg: Sandwiched policy gradient for masked diffusion language models. arXiv preprint arXiv:2510.09541. — [arXiv:2510.09541](https://arxiv.org/abs/2510.09541) find
  26. Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. 2017. Bridging the gap between value and policy based reinforcement learning. Advances in Neural Information Processing Systems, 30. find
  27. Gergely Neu, Anders Jonsson, and Vicenç Gómez. 2017. A unified view of entropy-regularized markov decision processes. arXiv preprint arXiv:1705.07798. — [arXiv:1705.07798](https://arxiv.org/abs/1705.07798) find
  28. Shen Nie, Fengqi Zhu, Zebin You, Xiaolu Zhang, Jingyang Ou, Jun Hu, Jun Zhou, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. 2025. Large language diffusion models. Preprint, arXiv:2502.09992.
  29. Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730–27744.
  30. Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning, Stefano Ermon, and Chelsea Finn. 2023. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36:53728– 53741.
  31. John Schulman, Xi Chen, and Pieter Abbeel. 2017a. Equivalence between policy gradients and soft qlearning. arXiv preprint arXiv:1704.06440. find
  32. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. 2017b. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. find
  33. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. 2022. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824–24837.
  34. Ronald J Williams and Jing Peng. 1990. An efficient gradient-based algorithm for on-line training of recurrent network trajectories. Neural computation, 2(4):490–501. find
  35. Tianhao Wu, Banghua Zhu, Ruoyu Zhang, Zhaojin Wen, Kannan Ramchandran, and Jiantao Jiao. 2023. Pairwise proximal policy optimization: Harnessing relative feedback for llm alignment. arXiv preprint arXiv:2310.00212. — [arXiv:2310.00212](https://arxiv.org/abs/2310.00212) find
  36. Ling Yang, Ye Tian, Bowen Li, Xinchen Zhang, Ke Shen, Yunhai Tong, and Mengdi Wang. 2025. Mmada: Multimodal large diffusion language models. arXiv preprint arXiv:2505.15809. — [arXiv:2505.15809](https://arxiv.org/abs/2505.15809) find
  37. Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, et al. 2025. Dapo: An open-source llm reinforcement learning system at scale. arXiv preprint arXiv:2503.14476.
  38. Huaye Zeng, Dongfu Jiang, Haozhe Wang, Ping Nie, Xiaotong Chen, and Wenhu Chen. Acecoder: Acing coder rl via automated test-case synthesis, 2025a. URL https://arxiv.org/abs/2502.01718. ↗
  39. Siyan Zhao, Devaansh Gupta, Qinqing Zheng, and Aditya Grover. 2025. d1: Scaling reasoning in diffusion large language models via reinforcement learning. arXiv preprint arXiv:2504.12216. — [arXiv:2504.12216](https://arxiv.org/abs/2504.12216) find
  40. Shiv Shankar. 2025. Padre: Pseudo-likelihood based alignment of diffusion language models. In 2nd AI for Math @ ICML 2025. find
  41. Fengqi Zhu, Rongzhen Wang, Shen Nie, Xiaolu Zhang, Chunwei Wu, Jun Hu, Jun Zhou, Jianfei Chen, Yankai Lin, Ji-Rong Wen, and Chongxuan Li. 2025. Llada 1.5: Variance-reduced preference optimization for large language diffusion models. Preprint, arXiv:2505.19223.
  42. Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, Y. K. Li, Y. Wu, et al. 2024. Deepseekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint arXiv:2402.03300.
  43. Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis Titsias. 2024. Simplified and generalized masked diffusion for discrete data. Advances in neural information processing systems, 37:103131– 103167. find
  44. Brian D. Ziebart, Andrew L. Maas, J. Andrew Bagnell, and Anind K. Dey. 2008. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence. find
  45. Richard S Sutton and Andrew G Barto. 2018. Reinforcement learning: An introduction. MIT press. find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/shankar-2026-energy-matching-based-preference-learning/