Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

TL;DR

This paper introduces SEPO, a novel policy gradient algorithm designed for fine-tuning discrete diffusion models with non-differentiable rewards, demonstrating scalability and efficiency in various generative tasks.

Contribution

We propose a theoretically justified, efficient policy gradient method, SEPO, specifically tailored for fine-tuning discrete diffusion models with non-differentiable rewards.

Findings

SEPO effectively fine-tunes discrete diffusion models.

The method scales well across multiple tasks.

Experimental results show improved performance and efficiency.

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.