シンプルで効果的なマスク拡散言語モデル

Simple and Effective Masked Diffusion Language Models

June 11, 2024

著者: Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, Volodymyr Kuleshov

cs.AI

要旨

拡散モデルは高品質な画像生成において優れた性能を発揮する一方で、従来の研究では言語モデリングにおいて拡散モデルと自己回帰（AR）手法の間に大きな性能差があると報告されています。本研究では、単純なマスク付き離散拡散が従来考えられていた以上に高性能であることを示します。我々は、マスク付き拡散モデルの性能を向上させる効果的なトレーニング手法を適用し、さらに改善をもたらす簡略化されたRao-Blackwell化目的関数を導出しました。我々の目的関数は単純な形式をしており、古典的なマスク付き言語モデリング損失の混合であり、効率的なサンプラーを許容するエンコーダのみの言語モデルをトレーニングするために使用できます。これには、従来の言語モデルのように半自己回帰的に任意の長さのテキストを生成できるモデルも含まれます。言語モデリングのベンチマークにおいて、現代のエンジニアリング手法でトレーニングされた一連のマスク付き拡散モデルは、拡散モデルの中で新たな最先端を達成し、ARのパープレキシティに近づいています。我々はコードを以下で公開しています：https://github.com/kuleshov-group/mdlm

English

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 release our code at: https://github.com/kuleshov-group/mdlm