DARE: Diffusion Large Language Models Alignment and Reinforcement Executor

1. 문제 정의 (Problem)

확산 기반 대규모 언어 모델 (dLLMs, Diffusion LLMs) 은 autoregressive(자기회귀) 모델의 대안으로 등장하여 병렬 생성과 이터레이티브 디노이징 (iterative denoising) 을 가능하게 했지만, 그 오픈소스 생태계는 심각한 단편화 (fragmentation) 상태에 있습니다.

• 연구 반복의 저해: 각 논문마다 별도의 모델 포크, 롤아웃 (rollout) 구현, 보상 인터페이스, 평가 스크립트를 제공하여 새로운 모델이나 목적 함수를 통합하려면 주변 인프라를 다시 구현해야 합니다.
• 비교의 어려움: 알고리즘 간의 차이와 실행/평가 방식의 차이가 혼재되어 있어, 공정한 비교가 어렵습니다.
• 기존 LLM 프레임워크의 부적합: 기존 LLM 강화학습 (RL) 프레임워크는 좌우 방향 (left-to-right) 생성과 단일 디코딩 경로의 로그 확률을 가정합니다. 반면, dLLM 은 확산 과정 (forward/reverse process), 디노이징 상태의 가능도 (likelihood) surrogate, 모델 가족별 특화 롤아웃 백엔드가 필요하므로 기존 프레임워크를 직접 재사용할 수 없습니다.

2. 방법론 (Methodology)

이러한 문제를 해결하기 위해 저자들은 **DARE (dLLMs Alignment and Reinforcement Executor)**라는 통합 프레임워크를 제안했습니다. DARE 는 verl(학습용) 과 OpenCompass(평가용) 를 기반으로 하되, dLLM 특유의 실행 계층을 추가하여 구축되었습니다.

핵심 아키텍처 및 기능

1. 통합 실행 스택 (Unified Execution Stack):

   • 모델 지원: 마스킹 확산 언어 모델 (MDLMs, 예: LLaDA, Dream) 과 블록 확산 언어 모델 (BDLMs, 예: SDAR, LLaDA2.x) 을 모두 지원합니다.
   • 학습 파이프라인: 지도 미세조정 (SFT), 매개변수 효율적 미세조정 (PEFT), 선호도 최적화 (Preference Optimization), 그리고 dLLM 특화 강화학습 (RL) 알고리즘을 하나의 공유 파이프라인으로 통합합니다.
   • 추상화: Worker(롤아웃, 액터, 보상 모델 등), Dataflow(데이터 흐름), Workflow(최적화 루프) 의 3 가지 추상화를 통해 알고리즘별 로직과 공유 워크플로우를 분리합니다.

2. 시스템 최적화 (Acceleration Backends):

   • 학습 측 최적화:
     • MDLM 의 경우: 가변 길이 배치 처리를 위해 FlashAttention-varlen 을 사용하여 패딩 오버헤드를 줄이고 학습 처리량을 향상시킵니다.
     • BDLM 의 경우: 반자기회귀 구조를 위해 FlexAttention 을 통합하고 커널 최적화를 적용합니다.
   • 롤아웃 측 최적화 (Rollout-side):
     • MDLM: 롤아웃 (샘플링) 과 액터 업데이트 시 사용하는 어텐션 백엔드를 분리합니다. 롤아웃에는 KV 캐시 재사용이 용이한 Fast-dLLM 과 FlashAttention 을 사용하여 속도를 극대화합니다.
     • BDLM: LMDeploy 및 SGLang 을 활용하여 반자기회귀 구조에 맞는 가속화를 구현합니다.
   • 성능: MDLM 경로에서 학습 및 롤아웃 백엔드를 분리 최적화함으로써 전체 RL 파이프라인 속도를 약 4 배 향상시켰으며, BDLM 경로에서는 14 배 이상의 가속화를 달성했습니다.

3. 통합 평가 플랫폼:

   • OpenCompass 를 dLLM 인식 (model-aware) 방식으로 확장하여, 학습 코드와 벤치마크 코드가 독립적으로 발전하는 문제를 해결하고 재현 가능한 평가를 제공합니다.

3. 주요 기여 (Key Contributions)

1. 통합 프레임워크: 마스킹 및 블록 확산 언어 모델 가족을 아우르는 단일 오픈소스 스택에서 dLLM 후학습 (post-training) 과 평가를 통합했습니다.
2. 광범위한 알고리즘 통합: SFT, PEFT, 선호도 최적화, 그리고 VRPO, D1, Coupled-GRPO, CJ-GRPO, SPG, BGPO, EBPO 등 다양한 dLLM 특화 RL 알고리즘을 통합하여 공정한 비교를 가능하게 했습니다.
3. 시스템 최적화 구현: 롤아웃, 학습, 평가 단계에 특화된 dLLM 전용 시스템 최적화 (Fast-dLLM, FlexAttention, fused loss kernels 등) 를 구현했습니다.
4. 재현성 및 비교 가능성: 파편화된 논문별 구현을 통합된 실행 환경으로 전환하여, 다양한 모델과 벤치마크에 대한 재현 가능한 연구 워크플로우를 제시했습니다.

4. 실험 결과 (Results)

저자들은 LLaDA, Dream, SDAR, LLaDA2.x 등 다양한 모델 가족을 대상으로 수학적 추론, 코딩, 계획 (planning) 태스크에서 실험을 수행했습니다.

• 모델 및 알고리즘 성능:
  • LLaDA-8B: 수학 태스크에서는 CJ-GRPO(GSM8K) 와 Coupled-GRPO(MATH) 가 우수했으나, 코딩 태스크에서는 VRPO 가 가장 강력했습니다. 계획 태스크에서는 Coupled-GRPO(Countdown) 와 BGPO(Sudoku) 가 각각 우세했습니다.
  • Dream-7B: CJ-GRPO 가 수학에서 우수했으나, 코딩에서는 Coupled-GRPO 가 가장 좋은 결과를 보였습니다.
  • 결론: 단일 알고리즘이 모든 태스크에서 우세하지 않으며, 백본 모델에 따라 최적의 알고리즘이 달라집니다. 이는 통합 실행 환경 하에서만 명확하게 드러나는 사실입니다.
• 안정성 분석:
  • ELBO 기반 방법론 (SPG, BGPO 등) 은 몬테카를로 샘플 수에 민감하여 수렴이 불안정하거나 reward collapse 가 발생하는 경향이 있었습니다.
  • 반면, d1, Coupled-GRPO, CJ-GRPO는 더 안정적인 보상 곡선을 보이며 효과성과 안정성 사이의 좋은 균형을 이루는 것으로 나타났습니다.
• 성능 향상: DARE 를 적용한 결과, 기존 파편화된 구현 대비 학습 및 롤아웃 속도가 획기적으로 개선되었습니다.

5. 의의 (Significance)

• 연구 인프라의 표준화: dLLM 연구가 모델 정의에서 후학습 및 평가 인프라로 초점이 이동함에 따라, DARE 는 이를 위한 재사용 가능한 연구 기반 (substrate) 을 제공합니다.
• 공정한 비교 체계: 알고리즘의 성능 차이를 실행 환경의 차이와 분리하여, 진정한 알고리즘적 개선 여부를 검증할 수 있는 토대를 마련했습니다.
• 시스템 - 알고리즘 간극 해소: 확산 모델의 고유한 시스템 요구사항 (디노이징, 병렬성, 블록 구조 등) 을 고려한 최적화를 통해, dLLM 의 실용적 배포와 연구를 가속화합니다.
• 미래 지향성: 오픈소스 커뮤니티의 기여를 유도하여, 새로운 확산 모델 (비전, 오미니 등) 과 알고리즘이 지속적으로 통합될 수 있는 확장 가능한 구조를 제시합니다.

요약하자면, DARE 는 확산 기반 LLM 의 후학습과 평가를 위한 첫 번째 통합적이고 최적화된 오픈소스 프레임워크로서, 연구의 재현성을 높이고 알고리즘 비교의 신뢰성을 확보하며 시스템 효율성을 극대화하는 데 핵심적인 역할을 합니다.