TESS 2: A Large-Scale Generalist Diffusion Language Model

📄 arXiv: 2502.13917v2 📥 PDF

作者: Jaesung Tae, Hamish Ivison, Sachin Kumar, Arman Cohan

分类: cs.CL

发布日期: 2025-02-19 (更新: 2025-06-01)

备注: ACL 2025 camera-ready

🔗 代码/项目: GITHUB

💡 一句话要点

TESS 2：一种大规模通用扩散语言模型，性能媲美自回归模型。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 扩散模型 指令跟随 自回归模型 预训练 指令调优 奖励引导 自然语言处理

📋 核心要点

  1. 现有指令调优扩散模型性能不足，难以匹敌自回归模型，限制了其应用。
  2. TESS 2通过持续预训练和指令调优，有效提升了扩散模型的指令跟随能力。
  3. 实验表明，TESS 2在指令跟随任务上超越了现有扩散模型，并能与自回归模型竞争。

📝 摘要（中文）

我们介绍了TESS 2，一种通用的指令跟随扩散语言模型，它优于现有的指令调优扩散模型，并且能够匹配甚至超越强大的自回归（AR）模型。我们通过以下方式训练TESS 2：首先，通过使用通常的交叉熵作为扩散损失进行持续预训练来调整一个强大的AR模型，然后进行进一步的指令调优。我们发现，适应性训练以及基础模型的选择对于训练良好的指令跟随扩散模型至关重要。我们进一步提出了奖励引导，这是一种新颖且模块化的推理时引导程序，用于对齐模型输出，而无需训练底层模型。最后，我们表明TESS 2随着推理时计算量的增加而进一步改进，突出了扩散LM在对推理时使用的计算量进行细粒度控制方面的效用。代码和模型可在https://github.com/hamishivi/tess-2获取。

🔬 方法详解

问题定义：论文旨在解决现有指令调优扩散语言模型在性能上不如自回归模型的问题。现有方法在指令跟随能力上存在不足，限制了扩散模型在自然语言处理任务中的应用。

核心思路：论文的核心思路是首先利用一个强大的自回归模型作为基础，通过持续预训练使其适应扩散模型的训练方式，然后再进行指令调优。这种方法结合了自回归模型的优势和扩散模型的灵活性。

技术框架：TESS 2的训练框架包含两个主要阶段：1) 适应性训练：使用交叉熵损失，对预训练的自回归模型进行持续预训练，使其适应扩散模型的训练方式。2) 指令调优：在适应性训练的基础上，使用指令数据对模型进行进一步的调优，提升其指令跟随能力。在推理阶段，论文提出了奖励引导方法，用于在不改变模型参数的情况下，对模型输出进行调整。

关键创新：论文的关键创新在于：1) 提出了一种有效的训练扩散语言模型的方法，即先进行适应性训练，再进行指令调优。2) 提出了奖励引导方法，可以在推理时对模型输出进行调整，而无需重新训练模型。3) 强调了基础模型的选择对于训练良好的指令跟随扩散模型的重要性。

关键设计：论文的关键设计包括：1) 选择一个强大的自回归模型作为基础模型。2) 使用交叉熵损失进行适应性训练。3) 使用指令数据进行指令调优。4) 设计奖励函数用于奖励引导。具体的参数设置和网络结构细节在论文中未详细说明，属于未知信息。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TESS 2在指令跟随任务上取得了显著的性能提升，超越了现有的指令调优扩散模型，并且能够匹配甚至超越强大的自回归模型。此外，论文还表明，TESS 2的性能随着推理时计算量的增加而进一步提高，突出了扩散模型在计算资源利用方面的优势。

🎯 应用场景

TESS 2具有广泛的应用前景，例如对话生成、文本摘要、机器翻译等。其强大的指令跟随能力使其能够更好地理解用户意图，生成更符合用户需求的文本。此外，奖励引导方法为模型输出的控制提供了更大的灵活性，可以应用于各种需要对模型输出进行精细调整的场景。

📄 摘要（原文）

We introduce TESS 2, a general instruction-following diffusion language model that outperforms contemporary instruction-tuned diffusion models, as well as matches and sometimes exceeds strong autoregressive (AR) models. We train TESS 2 by first adapting a strong AR model via continued pretraining with the usual cross-entropy as diffusion loss, and then performing further instruction tuning. We find that adaptation training as well as the choice of the base model is crucial for training good instruction-following diffusion models. We further propose reward guidance, a novel and modular inference-time guidance procedure to align model outputs without needing to train the underlying model. Finally, we show that TESS 2 further improves with increased inference-time compute, highlighting the utility of diffusion LMs in having fine-grained controllability over the amount of compute used at inference time. Code and models are available at https://github.com/hamishivi/tess-2.