Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

2025-02-03 · Oussama Zekri, Nicolas Boullé

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.

📄 PDF Abstract BibTeX arXiv:2502.01384

Code (1)

ozekri/SEPO 공식 구현 pytorch

Tasks

Language ModelingLanguage ModellingPolicy Gradient Methods

Methods 이 논문이 사용한 방법론

Softmax The Softmax output function transforms a previous layer's output into a vector of probabilities. It is commonly used for multiclass classification. Given an input vector $x$…

Attention 설명 없음

Diffusion Diffusion models generate samples by gradually removing noise from a signal, and their training objective can be expressed as a reweighted variational lower-bound…

A Continuous-Time Reinforcement Learning Framework for Fine-Tuning Discrete Diffusion Models

2026-07-16 · Zikun Zhang, Jiayuan Sheng, David D. Yao, Wenpin Tang arxiv

We formulate reinforcement learning (RL) in continuous time with discrete state spaces and possibly arbitrary action spaces via a stochastic control approach, where the state dynamics are modeled as a controlled continuo…

Mathematical Reasoning Reinforcement Learning

Policy Gradient Guidance Enables Test Time Control

2025-10-02 · Jianing Qi, Hao Tang, Zhigang Zhu arxiv

We introduce Policy Gradient Guidance (PGG), a simple extension of classifier-free guidance from diffusion models to classical policy gradient methods. PGG augments the policy gradient with an unconditional branch and in…

Reinforcement Learning Continuous Control

Optimizing DDPM Sampling with Shortcut Fine-Tuning

2023-01-31 · Ying Fan, Kangwook Lee

In this study, we propose Shortcut Fine-Tuning (SFT), a new approach for addressing the challenge of fast sampling of pretrained Denoising Diffusion Probabilistic Models (DDPMs). SFT advocates for the fine-tuning of DDPM…

Denoising Reinforcement Learning (RL)

Diffusion Policy Policy Optimization

2024-09-01 · Allen Z. Ren, Justin Lidard, Lars L. Ankile, Anthony Simeonov 외

We introduce Diffusion Policy Policy Optimization, DPPO, an algorithmic framework including best practices for fine-tuning diffusion-based policies (e.g. Diffusion Policy) in continuous control and robot learning tasks u…

continuous-control Continuous Control Reinforcement Learning (RL)

Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function

2025-12-04 · Hyeongyu Kang, Jaewoo Lee, Woocheol Shin, Kiyoung Om 외 arxiv

Diffusion models excel at generating high-likelihood samples but often require alignment with downstream objectives. Existing fine-tuning methods for diffusion models significantly suffer from reward over-optimization, r…