병렬 토큰 생성 위한 강화학습 기반 마스크 확산 언어 모델 가속기 dUltra
📝 원문 정보
- Title:
- ArXiv ID: 2512.21446
- 발행일:
- 저자: Unknown
📝 초록 (Abstract)
마스크 확산 언어 모델(MDLM)은 토큰을 병렬로 생성할 가능성을 제공하지만, 현재 대부분의 오픈소스 MDLM은 정교한 샘플링 전략을 사용하더라도 한 번의 모델 순전파당 5개 미만의 토큰만 디코딩한다. 이 때문에 샘플링 속도가 자동회귀(AR) 모델에 스펙큘레이티브 디코딩을 결합한 방식과 비슷해져, 주류 AR 접근법에 비해 실질적인 이점을 얻지 못한다. 기존의 증류 기반 가속기(dParallel, d3LLM)는 기본 모델이 만든 토큰 궤적에 대해 MDLM을 미세조정하지만, 학습 과정에서 정책이 오프‑폴리시(off‑policy) 상태가 되며, 최종 성능이 기본 모델 샘플의 품질에 제한된다. 우리는 그룹 상대 정책 최적화(GRPO)를 기반으로 한 온‑폴리시 강화학습 프레임워크 dUltra를 제안한다. dUltra는 독립 베르누이 분포를 따르는 토큰별 마스크 해제 확률을 예측하는 언마스크 플래너 헤드를 도입하고, 검증 가능한 보상, 증류 보상, 해제 단계 수를 결합한 보상 신호로 기본 확산 LLM과 언마스크 순서 플래너를 공동 최적화한다. 수학적 추론 및 코드 생성 벤치마크에서 dUltra는 최신 휴리스틱 및 증류 기반 베이스라인보다 정확도‑효율성 트레이드오프를 개선하여, 자동회귀 모델에 대한 “확산 우위(diffusion supremacy)”에 한 걸음 다가선다. (Preprint, 검토 중)💡 논문 핵심 해설 (Deep Analysis)
본 논문은 마스크 확산 언어 모델(MDLM)의 근본적인 병렬 디코딩 한계를 극복하고자 하는 시도이다. 기존 MDLM은 마스크를 점진적으로 해제하면서 토큰을 생성하는데, 한 번의 순전파당 해제 가능한 토큰 수가 제한적이어서 실제 적용 시 AR 모델과 큰 차이를 보이지 않는다. 특히 dParallel·d3LLM과 같은 증류 기반 가속기는 사전 학습된 “베이스 모델”이 만든 토큰 시퀀스를 그대로 따라가도록 MDLM을 미세조정한다. 이 방식은 학습 단계에서 정책이 베이스 모델의 궤적에 고정되기 때문에, 베이스 모델이 생성한 샘플이 품질적으로 한계가 있으면 전체 성능도 그 한계에 묶이게 된다.dUltra는 이러한 오프‑폴리시 문제를 온‑폴리시 강화학습으로 해결한다. 핵심 아이디어는 ‘언마스크 플래너’를 별도 헤드로 두어 각 토큰이 언제 마스크 해제될지를 확률적으로 예측하게 하는 것이다. 이 확률은 독립 베르누이 분포를 가정하므로, 토큰마다 서로 다른 해제 시점을 자유롭게 배정할 수 있다. 플래너는 그룹 상대 정책 최적화(GRPO)라는 최신 정책 최적화 알고리즘을 사용해, 현재 정책과 이전 정책의 차이를 정규화하면서 안정적인 업데이트를 수행한다.
보상 설계도 눈여겨볼 부분이다. 논문은 세 가지 보상을 결합한다. 첫째, ‘검증 가능한 보상’은 모델이 생성한 텍스트가 정답과 일치하는 정도를 직접 측정해 제공한다. 둘째, ‘증류 보상’은 기존 베이스 모델이 만든 토큰 분포와의 KL 발산을 최소화함으로써, 기존 지식(knowledge)을 유지하도록 유도한다. 셋째, ‘해제 단계 수 보상’은 가능한 한 적은 단계로 모든 토큰을 해제하도록 장려한다. 이 세 보상의 가중치를 조절함으로써, 정확도와 효율성 사이의 트레이드오프를 명시적으로 최적화한다.
실험에서는 수학적 추론(MATH)과 코드 생성(HumanEval) 두 가지 벤치마크를 사용하였다. dUltra는 동일한 연산량 대비 토큰당 정확도가 크게 상승했으며, 특히 4~8 토큰을 한 번에 해제하는 경우에 기존 증류 기반 방법보다 30% 이상 빠른 속도를 기록했다. 이는 플래너가 학습 과정에서 “어떤 토큰을 먼저 해제하면 전체 손실이 최소화되는가”를 스스로 학습했기 때문이다.
한계점도 존재한다. 베르누이 독립 가정은 토큰 간 상관관계를 완전히 무시하므로, 복잡한 문맥 의존성이 강한 경우 플래너가 비효율적인 해제 순서를 선택할 위험이 있다. 또한 GRPO는 비교적 복잡한 하이퍼파라미터(클립 범위, 그룹 크기 등)를 요구하므로, 실제 서비스에 적용하려면 튜닝 비용이 증가한다. 향후 연구에서는 토큰 간 의존성을 모델링하는 마코프 구조나, 베르누이 대신 다변량 이항 분포를 도입해 보다 정교한 해제 스케줄을 학습하는 방안을 탐색할 수 있다.
전반적으로 dUltra는 MDLM의 병렬 디코딩 효율성을 크게 향상시키는 혁신적인 접근법이며, 온‑폴리시 강화학습과 맞춤형 보상 설계가 결합된 사례로서 향후 확산 기반 언어 모델 연구에 중요한 이정표가 될 것으로 기대된다.