dllm-reasoning/d1

Official Implementation for the paper "d1: Scaling Reasoning in Diffusion Large Language Models via Reinforcement Learning"

What it solves

It addresses the challenge of scaling reasoning capabilities in pre-trained diffusion Large Language Models (dLLMs), which typically struggle with complex reasoning tasks compared to standard autoregressive models.

How it works

The project implements a two-stage approach:

  1. Masked SFT: It uses completion-only masked Supervised Fine-Tuning (SFT) to initialize the model.
  2. diffu-GRPO: A novel policy gradient method based on Group Relative Policy Optimization (GRPO) specifically designed for masked dLLMs. This method features efficient log probability estimation to enable reinforcement learning (RL) for reasoning scaling.

Who it’s for

Researchers and developers working on diffusion-based language models and those interested in applying reinforcement learning to improve the reasoning performance of non-autoregressive LLMs.

Highlights

  • Novel RL Method: Introduces diffu-GRPO for efficient policy gradient updates in diffusion LLMs.
  • Two-Stage Pipeline: Combines masked SFT with RL to scale reasoning.
  • Complete Tooling: Provides implementation for SFT, RL training, and evaluation scripts.
  • Integration: Follows a Transformers-like interface for training and data collation.

Related

  • Project
    om-ai-lab/VLM-R1

    VLM-R1 is a framework for training R1-style Large Vision-Language Models using GRPO to improve reasoning and generalization in tasks like visual grounding, math, and object detection.

  • Project
    X-GenGroup/Flow-Factory

    A unified framework for online RL and offline fine-tuning of diffusion and flow-matching models across image, video, and audio-video modalities.

  • Project
    hiyouga/EasyR1

    An efficient and scalable reinforcement learning training framework that extends support to vision language models and multiple RL algorithms like GRPO.

  • Project
    ZHZisZZ/dllm

    A library that unifies the training, inference, and evaluation of diffusion language models, providing scalable pipelines and recipes for converting AR models into diffusion models.