← 返回信息流
技术博客arXiv cs.CL·3 小时前

全新iLLaDA模型:全双向注意力掩码扩散语言模型

原标题:Improved Large Language Diffusion Models

速览

研究团队提出iLLaDA,这是一个从0开始训练的8B参数掩码扩散语言模型,采用全双向注意力机制。该模型在预训练中使用12T token,并在25B token指令语料上进行微调,显著提升了通用、数学和代码基准测试表现。尽管采用非自回归训练,其性能仍与Qwen2.5 7B相当,证明了全双向扩散训练路径的竞争力。

AI 深度解读

解读 iLLaDA:从头训练的 80 亿参数全双向扩散语言模型

背景

当前,主流的大语言模型(LLMs)几乎都采用自回归(Autoregressive, AR)因子化机制,并结合因果注意力(Causal Attention)进行训练。这种架构虽然成熟,但在并行生成效率和某些特定任务的表现上存在固有局限。与此同时,基于扩散(Diffusion)的生成模型在图像生成领域取得了巨大成功,但在自然语言处理领域的应用仍处于探索阶段。

传统的扩散语言模型往往面临训练不稳定、难以扩展到大规模数据以及推理速度慢等挑战。此外,现有的许多扩散语言模型并非完全从头训练,或者在训练过程中混合使用了自回归损失,未能充分发挥扩散模型在建模全局依赖关系上的潜力。

在此背景下,研究团队提出了 iLLaDA(improved Large Language Diffusion Models),旨在验证一种全新的训练范式:从头开始(From Scratch),使用完全双向注意力(Fully Bidirectional Attention)和纯掩码扩散目标(Masked Diffusion Objective)来训练大规模语言模型。

核心内容

本文介绍了 iLLaDA,这是一个拥有 80 亿(8B)参数的掩码扩散语言模型。与传统的自回归模型不同,iLLaDA 在预训练和监督微调(SFT)阶段始终保留掩码扩散目标,并采用完全双向注意力机制。

1. 训练规模与流程

  • 预训练:模型在高达 12万亿(12T) token 的数据集上进行了预训练。这一规模足以让模型学习到丰富的语言模式和世界知识。
  • 监督微调(SFT):在预训练之后,模型在一个包含 250亿(25B) token 的指令语料库上进行了 12个 epoch 的微调,以增强其遵循指令的能力。
  • 架构特点:iLLaDA 摒弃了因果掩码,使用了完全双向的注意力机制。这意味着在生成过程中,模型可以同时看到上下文的完整信息,从而更准确地预测被掩码的 token。

2. 技术优化

  • 可变长度生成(Variable-length Generation):为了提高生成效率,iLLaDA 引入了可变长度的生成策略。这允许模型根据上下文动态调整生成步骤,避免了固定步长带来的计算浪费或质量损失。
  • 基于置信度的评分(Confidence-based Scoring):针对多项选择评估任务,研究团队引入了一种基于置信度的评分机制。由于扩散模型是概率性的,直接比较生成结果可能不够稳定,通过置信度评分可以更准确地评估模型在选择题等任务上的表现。

3. 性能表现

iLLaDA 在多个基准测试中展现了显著的性能提升,尤其是在通用知识、数学推理和代码生成方面:

  • 相比基线模型 LLaDA 的提升

    • iLLaDA-Base:在 BBH(Big-Bench Hard)基准上提升了 21.6 分,在 ARC-Challenge(科学推理挑战)上提升了 14.9 分。
    • iLLaDA-Instruct:在 MATH(数学推理)基准上提升了 14.5 分,在 HumanEval(代码生成)基准上提升了 16.5 分。
  • 与主流自回归模型的对比: 尽管 iLLaDA 采用的是非自回归的扩散训练方式,但在多个基准测试中,其表现与 Qwen2.5 7B(通义千问 2.5 的 70 亿参数版本)相当,证明了全双向扩散训练路径的竞争力。

关键要点

  • 范式创新:iLLaDA 是首个从头训练、全程使用纯掩码扩散目标且采用完全双向注意力的 8B 级语言模型。它证明了非自回归架构在大规模语言建模中的可行性。
  • 超大规模预训练:模型在 12T token 上预训练,并在 25B token 的指令集上微调,这种数据规模确保了模型具备强大的基础能力。
  • 显著的性能跃升:相比之前的扩散语言模型基线(LLaDA),iLLaDA 在 BBH、ARC-Challenge、MATH 和 HumanEval 等关键基准上取得了两位数的显著提升。
  • 竞争力验证:iLLaDA 的性能足以匹敌当前主流的自回归模型 Qwen2.5 7B,表明扩散模型并非仅仅是自回归模型的替代品,而是一条独立且高效的竞争路径。
  • 效率与评估优化:通过可变长度生成提高了推理效率,并通过基于置信度的评分机制解决了扩散模型在多项选择题评估中的不确定性问题。
  • 开源贡献:研究团队公开了模型权重和代码,为社区探索非自回归语言模型提供了宝贵的资源。

意义与影响

iLLaDA 的研究成果对大语言模型的发展具有重要的理论和实践意义:

  1. 挑战自回归霸权:长期以来,自回归模型主导了 NLP 领域。iLLaDA 的成功表明,完全双向的扩散模型同样可以构建出高性能的语言模型,为打破自回归架构的垄断提供了强有力的实证支持。
  2. 并行生成的潜力:由于扩散模型本质上是并行去噪过程,iLLaDA 展示了非自回归架构在并行生成方面的潜在优势。随着可变长度生成等技术的优化,扩散模型有望在推理速度上超越自回归模型,特别是在长文本生成场景中。
  3. 双向注意力的优势:全双向注意力允许模型在生成时利用完整的上下文信息,这在处理需要全局理解的任务(如数学推理、代码补全)时可能比因果注意力更具优势。iLLaDA 在 MATH 和 HumanEval 上的高分印证了这一点。
  4. 推动扩散模型在 NLP 的落地:通过解决训练稳定性、大规模扩展和评估难题,iLLaDA 为扩散语言模型从实验室走向实际应用铺平了道路。它证明了扩散模型不仅可以用于简单的文本生成,还可以胜任复杂的推理和编码任务。
  5. 开源生态的丰富:模型权重和代码的公开,将激励更多研究人员探索扩散语言模型的架构优化、训练技巧和潜在应用,加速该领域的创新迭代。

总之,iLLaDA 不仅是一个性能强大的 8B 参数模型,更是语言模型架构多元化发展的重要里程碑。它证明了“全双向扩散训练”是一条通往强大语言模型的可行且极具竞争力的路径。

查看原文 →arxiv.org