SEPO_DNA

Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods

This repository contains the code for the SEPO algorithm presented in the paper: Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods.

SEPO (Score Entropy Policy Optimization) is an efficient, broadly applicable, and theoretically justified policy gradient algorithm 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, including applications on fine-tuning a masked diffusion language model on DNA sequences.

<p align="center"> <img src="https://github.com/ozekri/SEPO/blob/main/img/denoising_RLHF.gif" width=80% height=80% alt="Denoising RLHF process visualization"> </p>

For more details and the full implementation, please refer to the official GitHub repository.

Sample Usage: Download Checkpoint

You can download the fine-tuned models from Hugging Face directly using the huggingface_hub Python library to reproduce results:

from huggingface_hub import hf_hub_download

# Example: Download the SEPO fine-tuned model checkpoint
ckpt_path = hf_hub_download(
    repo_id="Xssama/SEPO_DNA",
    filename="finetuned_sepo_kl.ckpt",  # finetuned_sepo_kl_gf.ckpt for SEPO with gradient flow
    cache_dir="./checkpoints"  # Optional: specify your preferred local directory
)

print(f"Checkpoint downloaded to: {ckpt_path}")

Alternatively, you can use wget:

wget https://huggingface.co/Xssama/SEPO-DNA/resolve/main/finetuned_sepo_kl.ckpt -P ./checkpoints/

Citation

If you find this work useful in your research, please consider citing:

@article{zekri2025fine,
  title={Fine-Tuning Discrete Diffusion Models with Policy Gradient Methods},
  author={Zekri, Oussama and Boull{\'e}, Nicolas},
  journal={arXiv preprint arXiv:2502.01384},
  year={2025}
}