[연구] 확산 LLM의 속도를 높이는 ORL 소개 (완전 오픈 소스)

안녕하세요, r/MachineLearning! 저희는 확산 언어 모델(diffusion language models, dLLMs)의 생성 품질을 희생하지 않으면서도 처리 속도를 획기적으로 개선한 새로운 연구인 ORL (On-Policy Residual Learning)을 공유하게 되어 기쁩니다. ORL과 자체 추론 엔진 dlmrt를 사용하여 다음과 같은 결과를 달성했습니다:

• 단일 RTX 5090에서 2,788 tokens/s — LLaDA-2.1-mini + ORL
• 8× B200에서 NVFP4로 7,708 tokens/s — DiffusionGemma 26B MoE + ORL
• QwenDLM-4B 모델의 경우, 동등하거나 더 나은 품질을 유지하면서 DFlash 추측 디코딩(speculative decoding)보다 최대 6.9배 높은 종단 간 처리량(end-to-end throughput)을 달성했으며, 단일 요청 H100 비교에서 정확도가 일치했을 때 1.74배 더 빠릅니다.

*모든 수치는 배치 크기 1의 단일 요청 디코딩에 대한 것이며, 동시 요청 전반의 총 처리량(aggregate throughput)은 아닙니다. 위에 제시된 GPU 수치는 임의의 프롬프트에 대한 평균이 아닌 특정 벤치마크에서의 최고 결과입니다.

저희는 모델, 학습 코드, 평가 파이프라인 및 추론 엔진을 오픈 소스로 공개했습니다.

ORL의 아이디어는 무엇일까요? 한 번에 하나의 토큰을 생성하는 자기회귀적 LLM(autoregressive LLMs)과 달리, 확산 LLM은 여러 토큰을 병렬로 정제하고 최종화할 수 있습니다. 하지만 적은 순방향 패스(forward passes)만으로 이들이 수렴하도록 만드는 것은 어렵습니다. 저희 접근 방식에는 두 가지 핵심 요소가 있습니다:

1. 온-정책 잔차 학습 (On-policy residual learning): 각 디노이징 단계에서 모든 것을 처음부터 예측하는 대신, 모델은 이전 예측을 잔차 컨텍스트(residual context)로 유지하고 점진적으로 이를 정제합니다. 문제는 이러한 잔차가 모델 자체에 의해 생성되며, 훈련 과정 동안 그 분포가 변한다는 것입니다. 오래된 오프라인 잔차(stale, offline residuals)로 훈련하면 추론 시와 불일치(mismatch)가 발생합니다. ORL은 현재 모델이 생성한 신선한 궤적(fresh trajectories)으로 직접 훈련하여, 추론 중에 마주치는 정확한 잔차 상태를 재현(replaying)함으로써 이를 해결합니다.
2. 불확실성을 더 빠르게 줄이는 학습 (Learning to reduce uncertainty faster): 최종 출력만을 단순히 지도하는 것만으로는 모델이 더 빨리 수렴하도록 가르치지 못할 수도 있습니다. 저희는 Adjacent Entropy Contraction을 도입하여, 각 디노이징 단계가 이전 단계에 비해 불확실성(uncertainty)을 줄이도록 장려합니다.

또한 불필요하게 이미 자신감이 높은 예측을 밀어내는 것을 방지하기 위해 손실(loss)에 게이팅(gate)을 적용했습니다. 그 결과: 모델은 정확도를 유지하거나 향상시키면서 포워드 패스(forward pass)당 더 많은 토큰을 확정하는 방법을 학습합니다. 몇 가지 벤치마크 결과를 살펴보겠습니다. 저희는 세 가지 모델 계열에서 ORL을 테스트했습니다. QwenDLM-4B: Qwen3-4B를 변환한 저희의 확산 모델입니다. ORL은 dlmrt를 사용하여 포워드당 토큰 수를 최대 6.2배, 종단 간 처리량(end-to-end throughput)을 약 6.9배 향상시켰습니다. HumanEval pass@1은 82.9%에서 89.6%로 개선되었습니다. LLaDA-2.1-mini: ORL은 자체 생성 학습 토큰 4.2M개와 후처리(post-training)에 H100 GPU-시간 36시간만을 필요로 하며, 비교 가능한 품질에서 최대 3배 높은 처리량을 제공합니다. DiffusionGemma 26B MoE: ORL은 이미 매우 병렬화된 확산 모델의 성능을 더욱 개선하여, 포워드 패스당 최대 1.2배 더 많은 토큰을 달성했습니다. 자가회귀 디코딩(autoregressive decoding)과의 보다 직접적인 비교를 위해, 단일 H100에서 저희의 QwenDLM-4B + ORL은 1,298 tokens/s에 도달하는 반면, DFlash는 746 tokens/s, EAGLE3는 완화된 검증(relaxed verification) 조건 하에서 비교 가능한 측정 정확도로 405 tokens/s를 기록했습니다. 이들은 종단 간 시스템 비교입니다: ORL은 dlmrt를 사용하고, 추측 디코더(speculative decoders)는 vLLM을 사용합니다. 또한 저희는 사용자 지정 추론 엔진인 dlmrt를 구축했습니다. 저희는 낮은 지연 시간의 단일 요청 확산 디코딩에 특화되도록 dlmrt를 개발했습니다. dlmrt는 CPU에서 모든 노이즈 제거(denoising) 단계를 시작하고 동기화하는 대신, 전체 노이즈 제거 루프를 CUDA 그래프로 캡처하여 GPU에서 실행합니다. 또한 소비자용 RTX 5090을 포함한 Blackwell GPU의 NVFP4 추론도 지원합니다. 동일한 모델과 디코딩 설정을 사용했을 때, dlmrt는 단일 H100 비교에서 SGLang/vLLM보다 1.2–1.6배 더 빠르게 작동했습니다. 학습 측면에서도 저희는 효율적인 잔여 상태 캡처(residual-state capture), 특수 어텐션 커널(specialized attention kernels), 그리고 컴퓨팅 균형을 맞춘 패킹(compute-balanced packing) 주변으로 시스템을 공동 설계하여, 학습 단계를 최대 5.2배 가속화했습니다.

모든 것이 오픈 소스입니다. 논문: https://arxiv.org/abs/2610.20687 프로젝트 페이지 및 인터랙티브 데모: https://introspective-diffusion.github.io/orl-dlm/ 학습 코드 및 평가: https://github.com/Introspective-Diffusion/ORL Inference engine (dlmrt): https://github.com/Introspective-Diffusion/DLMRT 모델 체크포인트: https://huggingface.co/yifanyu RTX 5090을 가지고 계시다면, dlmrt를 사용하여 LLaDA-2.1-mini ORL NVFP4 체크포인트를 시도해 볼 수 있습니다. 커뮤니티의 피드백을 환영하며, 특히 결과 재현, 다양한 워크로드 테스트 또는 추론 엔진 기여에 관심 있는 분들의 의견을 듣고 싶습니다. 로컬 추론을 위한 확산 LLM에 대해 어떻게 생각하시는지 궁금합니다. 최적화하고 싶은 주요 병목 현상은 단일 요청 지연 시간(single-request latency)인가요, 아니면 우선순위를 두고 싶은 다른 병목 현상이 있나요? /u/arty_photography가 제출했습니다 [링크] [댓글]