arXiv:2410.18514cs.AIcs.CL2024-10ICLR被引 196

提出首个掩码扩散模型缩放定律,11亿参数模型在生成与理解上媲美大模型。

Scaling up Masked Diffusion Models on Text

  • 基于概率框架设计无监督分类器自由引导,有效利用无配对数据提升性能
  • 11亿参数模型在8个零样本任务中4项超越同规模TinyLlama,数学推理媲美70亿参数的Llama-2
  • 采样速度比自回归模型快1.4倍,且能处理双向推理和数据时序变化

掩码扩散模型(MDMs)在语言建模中展现出潜力,但其在文本生成与语言理解等核心任务中的可扩展性与有效性仍待探索。本文首次建立MDM的缩放定律,显示其缩放速率接近自回归模型(ARMs),且计算差距较小。受可扩展性启发,我们训练了最大达11亿参数的MDM家族,系统评估其与同等或更大规模ARMs的表现。充分利用MDMs的概率形式,提出一种简单有效的无监督分类器自由引导方法,有效利用大规模无配对数据,显著提升条件推理性能。在语言理解方面,11亿参数的MDM在8个零样本基准中的4个上优于同规模的TinyLlama;在数学推理上,其表现与70亿参数的Llama-2在GSM8K数据集上相当。在文本生成中,预训练时间增加16倍的MDM在加速采样技术KV-Cache支持下,性能与ARMs相当,且采样速度提升1.4倍。此外,MDMs有效解决自回归模型的挑战,如双向推理与数据时序漂移。值得注意的是,11亿参数的MDM克服了更大规模自回归模型(如130亿参数的Llama-2和1750亿参数的GPT-3)遭遇的反向诅咒问题。代码已开源。

原文摘要 · Abstract (English)

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.

掩码扩散文本生成语言模型大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。