[논문 리뷰] Simple and Effective Masked Diffusion Language Models

Subham Sekhar Sahoo, Marianne Arriola외 6명|arXiv (Cornell University)|2024. 06. 11.

Natural Language Processing Techniques인용 수 12

한 줄 요약

MASKED 확산 언어 모델(MDLM)이 SUBS 매개변화와 Rao-Blackwellized ELBO를 통해 언어 벤치마크에서 확산 모델 중 새로운 최첨단 성능을 달성하고, autoregressive perplexity에 다가가며, 효율적인 부분 자기회귀 생성이 가능하다.

ABSTRACT

While diffusion models excel at generating high-quality images, prior work reports a significant performance gap between diffusion and autoregressive (AR) methods in language modeling. In this work, we show that simple masked discrete diffusion is more performant than previously thought. We apply an effective training recipe that improves the performance of masked diffusion models and derive a simplified, Rao-Blackwellized objective that results in additional improvements. Our objective has a simple form -- it is a mixture of classical masked language modeling losses -- and can be used to train encoder-only language models that admit efficient samplers, including ones that can generate arbitrary lengths of text semi-autoregressively like a traditional language model. On language modeling benchmarks, a range of masked diffusion models trained with modern engineering practices achieves a new state-of-the-art among diffusion models, and approaches AR perplexity. We provide the code, along with a blog post and video tutorial on the project page: https://s-sahoo.com/mdlm

연구 동기 및 목표

• 이산 언어 데이터에 대한 확산 모델의 필요성 제시와 autoregressive 모델과의 성능 격차 축소.
• 원리 있는 학습 목표를 갖춘 간단하고 효과적인 MDLM 프레임워크를 개발.
• 향상된 학습을 위한 Rao-Blackwellized 연속시간 변분 하한을 도출.
• 인코더 전용 모델에 대해 부분 자기회귀 생성을 포함한 효율적인 샘플링을 가능하게 한다.
• MDLM 프레임워크를 DNA 서열 같은 비언어 도메인으로 확장하고 생성 능력을 시연한다.

제안 방법

• 데이터와 마스크 토큰 사이를 보간하는 토큰에 대한 이산적 순방향 확산 과정을 정의한다.
• SUBS 도입: 제로 마스킹 확률과 이어지는 비마스킹을 보장하는 치환 기반 역확산 매개변수화.
• Rao-Blackwellized 연속시간 NELBO를 도출하여 MLM 손실의 가중 평균으로 단순화된다.
• 시간 조건부 확산 트랜스포머 아키텍처(DiT)와 분산 감소를 위한 저차분도 샘플러로 학습한다.
• 선조 샘플링과 라운드 간에 접두사를 재사용하는 세미-자기회귀(SAR) 생성 전략으로 빠른 추론을 제공한다.

Figure 1 : (Left) Our proposed masked diffusion language model (MDLM) is trained using a weighted average of masked cross entropy losses. ( Top Right ) In comparison to masked language models (MLM), MDLM’s objective correspond to a principled variational lower bound, and supports generation via ance

실험 결과

연구 질문

• RQ1 효과적인 학습 방식과 함께한 마스킹된 이산 확산이 표준 언어 모델링 벤치마크에서 이전 확산 모델들을 능가할 수 있는가?
• RQ2 간단한 SUBS 매개변화가 MDLM에 대해 더 촘촘하고 분산이 작은 변분 하한을 제공하는가?
• RQ3 효율적인 샘플러를 갖는 인코더 전용 MDLM이 임의 길이의 텍스트를 부분 자기회귀적으로 생성할 수 있는가?
• RQ4 MDLM이 다운스트림 태스크 및 DNA 서열과 같은 교차 도메인 데이터에서 어떻게 성능을 발휘하는가?
• RQ5 학습 선택, 토큰화, 아키텍처가 MDLM의 성능에 미치는 영향이 AR 모델 및 기존 확산 방법과 비교하여 어떠한가?

주요 결과

• MDLM은 LM1B 및 OWT 벤치마크에서 확산 모델 중 새로운 최첨단을 달성한다.
• MDLM은 autoregressive perplexity에 근접하며, 설정에 따라 AR 모델 대비 상대적 격차를 최대 15–25%까지 감소시킨다.
• SUBS 매개변수화와 Rao-Blackwellized 연속시간 ELBO가 가능도 향상과 분산 감소를 달성한다.
• SAR 디코딩은 블록-자기회귀 확산 기준선보다 빠른 생성과 더 나은 생성 perplexity를 제공한다.
• MDLM은 DNA 서열 모델링에서 강력한 생성 및 다운스트림 성능을 보이고, 미세조정 시 GLUE의 다운스트림 지표에서도 경쟁력을 유지한다.
• 아블레이션 연구는 성능 향상에 있어 carry-over unmasking과 제로 마스킹 확률의 결정적 중요성을 보여준다.

Figure 2 : Generative perplexities across wall clock time for generating 64 samples on OWT using a single 32GB A5000 GPU are compared by varying $T\in\{100,500,1000,5000,10000\}$ in the reverse diffusion process. The samples are generated in mini-batches with a batch size of 16 for AR, SEDD, and MDL