DiRL: 확산 언어 모델을 위한 효율적인 사후 학습 프레임워크

초록

확산 언어 모델(dLLMs)은 자기회귀(AR) 모델에 대한 유망한 대안으로 부상했습니다. 최근 연구를 통해 dLLM의 사전 학습 잠재력이 검증되고 추론 속도가 가속화되었으나, dLLM의 사후 학습 분야는 여전히 미흡한 실정입니다. 기존 방법들은 계산 비효율성과 훈련-추론 간 목표 불일치로 인해 수학과 같은 복잡한 추론 과제에서 성능이 크게 제한됩니다. 이를 해결하기 위해 우리는 FlexAttention으로 가속화된 블록 단위 훈련과 LMDeploy로 최적화된 추론을 긴밀하게 통합한 효율적인 사후 학습 프레임워크인 DiRL을 제안합니다. 이 아키텍처는 간소화된 온라인 모델 업데이트 루프를 가능하게 하여 효율적인 2단계 사후 학습(지도 미세 조정 후 강화 학습)을 용이하게 합니다. 이 프레임워크를 기반으로 우리는 dLLM에 특화된 최초의 편향되지 않은 그룹 상대 정책 최적화(GRPO) 구현체인 DiPO를 제안합니다. 고품질 수학 데이터로 DiRL-8B-Instruct를 훈련하여 우리의 접근 방식을 검증했습니다. 우리 모델은 dLLM 중에서 최고 수준의 수학 성능을 달성했으며, 여러 벤치마크에서 유사한 Qwen2.5 시리즈 모델들을 능가했습니다.

English

Diffusion Language Models (dLLMs) have emerged as promising alternatives to Auto-Regressive (AR) models. While recent efforts have validated their pre-training potential and accelerated inference speeds, the post-training landscape for dLLMs remains underdeveloped. Existing methods suffer from computational inefficiency and objective mismatches between training and inference, severely limiting performance on complex reasoning tasks such as mathematics. To address this, we introduce DiRL, an efficient post-training framework that tightly integrates FlexAttention-accelerated blockwise training with LMDeploy-optimized inference. This architecture enables a streamlined online model update loop, facilitating efficient two-stage post-training (Supervised Fine-Tuning followed by Reinforcement Learning). Building on this framework, we propose DiPO, the first unbiased Group Relative Policy Optimization (GRPO) implementation tailored for dLLMs. We validate our approach by training DiRL-8B-Instruct on high-quality math data. Our model achieves state-of-the-art math performance among dLLMs and surpasses comparable models in the Qwen2.5 series on several benchmarks.

DiRL: 확산 언어 모델을 위한 효율적인 사후 학습 프레임워크

DiRL: An Efficient Post-Training Framework for Diffusion Language Models

초록

Support