提出PAPL训练方法,让扩散语言模型在规划路径下生成更准。
Planner Aware Path Learning in Diffusion Language Models Training
- 设计新损失函数P-ELBO,让训练对齐规划推理路径
- 蛋白生成提升40%,文本生成MAUVE增4倍,代码生成提23%
- 只需改损失函数,通用性强,适合各类扩散模型应用
扩散语言模型通过灵活并行生成路径实现快速推理,其能力依赖于规划器(planner)迭代选择更优去噪路径。但训练时假设随机去噪路径,而推理使用规划路径,造成训练与推理不一致。本文理论证明标准离散扩散训练目标(ELBO)无法准确描述非均匀规划下的去噪器。为此提出计划式证据下界(P-ELBO),将规划动态直接纳入训练目标。基于此构建规划感知路径学习(PAPL)训练方案,仅需修改标准掩码离散扩散损失,即可实现训练与规划推理对齐。实验表明,PAPL在多个领域持续提升:蛋白序列生成相对提升40%,文本生成MAUVE最高提升4倍,代码生成HumanEval pass@10提升23%。代码已开源。
原文摘要 · Abstract (English)
Diffusion language models have emerged as a powerful alternative to autoregressive models, enabling fast inference through more flexible and parallel generation paths. This flexibility of sampling is unlocked by new engineered sampling strategies, or planners, that select more favorable generation paths by iteratively planning - versus uniformly at random - where to denoise along the sequence. However, by modifying the reverse paths via planning, planners create an irrevocable mismatch between the uniformly random denoising paths assumed during training and planning-based inference. In this paper, we systematically investigate the mismatch of discrete diffusion training and inference under planning and theoretically prove that the standard discrete diffusion training evidence lower bound (ELBO) does not accurately describe a denoiser that uses a non-uniform planner. To address this gap, we derive a new planned evidence lower bound (P-ELBO) that incorporates planner-based reverse dynamics directly into the training objective. Using the P-ELBO, we introduce Planner Aware Path Learning (PAPL), a novel training scheme that aligns training and inference under a planned denoiser. PAPL is implemented as a simple yet effective modification to the standard masked discrete diffusion loss, making it widely applicable and easy to adopt. Empirically, we show PAPL delivers consistent gains across domains, including a 40% relative improvement in protein sequences, improved text generation with up to a 4x relative MAUVE gain, and 23% relative improvement in code generation HumanEval pass@10. Code is available at github.com/pengzhangzhi/PAPL .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。