Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

Oussama Zekri, Nicolas Boullé
2/3/2025
stat.MLcs.AIcs.CL

Official code: ozekri/SEPO (32 stars, MIT, last commit 2025-05-19).

Abstract

Discrete diffusion models have recently gained significant attention due to their ability to process complex discrete structures for language modeling. However, fine-tuning these models with policy gradient methods, as is commonly done in Reinforcement Learning from Human Feedback (RLHF), remains a challenging task. We propose an efficient, broadly applicable, and theoretically justified policy gradient algorithm, called Score Entropy Policy Optimization (\SEPO), for fine-tuning discrete diffusion models over non-differentiable rewards. Our numerical experiments across several discrete generative tasks demonstrate the scalability and efficiency of our method. Our code is available at https://github.com/ozekri/SEPO.

View on arXivView PDF

Code Implementations(1)

ozekri/SEPOOfficial100%
MIT

Code for the paper: "Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods"

322Shell, PythonFeb 3, 20251 years agoMIT

Similar Papers

Flow Matching Policy Gradients

87% match2025

Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review

87% match1 repo1592024

A Diffusion Model Framework for Maximum Entropy Reinforcement Learning

87% match2025

EntRGi: Entropy Aware Reward Guidance for Diffusion Language Models

87% match2026

Directly Fine-Tuning Diffusion Models on Differentiable Rewards

86% match1 repo42026

Preference-Based Alignment of Discrete Diffusion Models

86% match2025

Cite this paper

@article{zekri2025finetuning,
  title  = {Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods},
  author = {Oussama Zekri and Nicolas Boullé},
  year   = {2025},
  eprint = {2502.01384},
  archivePrefix = {arXiv},
  url    = {http://arxiv.org/abs/2502.01384v3}
}

Add a badge to your README

Shows whether official code is linked for this paper.

Code of Paper badge
[![Code of Paper](https://codeofpaper.com/badge/2502.01384.svg)](https://codeofpaper.com/paper/2502.01384)

Discussion