Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

Oussama Zekri (CREST/ENSAE) · Nicolas Boulle (Imperial College London)
complex discrete structuresdiscrete diffusion modelsdiscrete generative tasksfine-tuninggenerative modelingnon-differentiable rewardsnumerical experimentsoptimization techniquespolicy gradient methodsreinforcement learning from human feedbackscalabilityscore entropy policy optimizationtheoretical justification

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.