Scaling up Masked Diffusion Models on Text

Masked diffusion models (MDMs) have shown promise in language modeling, yet their scalability and effectiveness in core language tasks, such as text generation and language understanding, remain underexplored. This paper establishes the first scaling law for MDMs, demonstrating a scaling rate comparable to autoregressive models (ARMs) and a relatively small compute gap. Motivated by their scalability, we train a family of MDMs with up to 1.1 billion (B) parameters to systematically evaluate their performance against ARMs of comparable or larger sizes. Fully leveraging the probabilistic formulation of MDMs, we propose a simple yet effective unsupervised classifier-free guidance that effectively exploits large-scale unpaired data, boosting performance for conditional inference. In language understanding, the 1.1B MDM outperforms the 1.1B TinyLlama model trained on the same data across four of eight zero-shot benchmarks. Notably, it achieves competitive math reasoning ability with the 7B Llama-2 model on the GSM8K dataset. In text generation, MDMs with 16 times more pre-training time offer a flexible trade-off against ARMs with the accelerated sampling technique KV-Cache: MDMs match ARMs in performance while being 1.4 times faster during sampling. Moreover, MDMs address challenging tasks for ARMs by effectively handling bidirectional reasoning and adapting to temporal shifts in data. Notably, a 1.1B MDM breaks the reverse curse encountered by much larger ARMs with significantly more data and computation, such as 13B Llama-2 and 175B GPT-3. Our code is available at https://github.com/ML-GSAI/SMDM.

한 줄 요약

1.1B 매스킹 디퓨전 모델(MDM)이 자동회귀 모델(ARM)과 유사한 성능을 보이며, 수학 추론 및 샘플링 속도에서 우수함.

핵심 기여도

MDM의 첫 번째 스케일링 법칙 제시: ARM과 유사한 스케일링 속도, 16배의 컴퓨트 격차.
1.1B MDM이 1.1B TinyLlama 모델을 8개의 zero-shot 벤치마크 중 4개에서 초과.
1.1B MDM이 7B Llama-2와 유사한 수학 추론 성능을 GSM8K에서 달성.
샘플링 시 1.4배 빠른 MDM, KV-Cache 기반 ARM과 동등한 성능 제공.

핵심 아이디어

기존 자동회귀 모델(ARM)은 언어 생성에서 단방향성과 샘플링 속도 문제를 가지며, 이에 따라 MDM이 대안으로 제시된다. MDM은 시퀀스 내 마스킹된 위치를 확률적으로 복원하는 방식으로, 양방향 맥락 모델링이 가능하다. 본 연구는 MDM의 확장 가능성을 입증하기 위해 최대 1.1B 파라미터를 가진 모델을 훈련하고, 이를 통해 ARM과 유사한 스케일링 속도를 보인다. 또한, **unsupervised classifier-free guidance (CFG)** 기법을 제안하여, 대규모 비결합 데이터를 활용해 조건부 추론 성능을 향상시킨다. 이 기법은 기존의 CFG와 달리 짝이 맞지 않은 데이터를 사용하면서도, 짝이 맞는 데이터가 있을 경우 이를 활용해 성능을 더 높인다.

기술적 접근법

**모델 크기**: 최대 1.1B 파라미터의 MDM 훈련.
**훈련 데이터**: 대규모 언어 데이터셋 사용 (명시되지 않음).
**추론 기법**: **unsupervised classifier-free guidance (CFG)** 도입.
**샘플링 기법**: KV-Cache 기반 ARM과 비교, MDM은 샘플링 시간이 1.4배 빠름.
**스케일링 분석**: IsoFLOP 기반으로 6×10¹⁸ ~ 10²⁰ FLOPs 범위에서 분석.
**성능 평가**: MT-Bench, GSM8K, 8개의 zero-shot 벤치마크 사용.

주요 결과

**언어 이해**: 1.1B MDM이 1.1B TinyLlama 모델을 8개의 zero-shot 벤치마크 중 4개에서 초과.
**수학 추론**: 1.1B MDM이 7B Llama-2 모델과 유사한 성능을 GSM8K에서 달성.
**샘플링 속도**: MDM이 1.4배 빠르게 샘플링하며, 16배 더 많은 훈련 시간이 필요.
**역의 정리 문제 해결**: 1.1B MDM이 13B Llama-2와 175B GPT-3보다 역의 정리(reverse curse) 문제를 더 잘 해결.

의의 및 한계

MDM은 ARM의 단점을 보완할 수 있는 잠재력을 보여주며, 특히 양방향 추론과 시간적 변화 대응에서 우수함. 또한, 수학 추론 및 샘플링 속도에서 ARM과 경쟁 가능하다는 점에서 실용적 가치가 있다. 그러나 MDM은 조건부 생성 및 스케일링 법칙에서 ARM에 비해 여전히 격차가 존재하며, 이는 향후 연구 주제로 제시된다. 또한, MDM의 확장성은 기존의 연속 디퓨전 모델보다 우수하지만, 대규모 실험과 표준 벤치마크 평가가 부족한 점이 한계로 지적된다.

실용적 활용

MDM은 양방향 추론이 필요한 자연어 이해, 수학 문제 해결, 긴 텍스트 생성 등 다양한 언어 모델링 작업에 적용 가능하다. 특히 샘플링 속도가 빠르고, 대규모 비결합 데이터를 효과적으로 활용할 수 있어, 대규모 언어 모델 개발 및 실시간 응용 분야에서 유용할 것으로 기대된다.