让语言模型与规划模块联合训练，提升生成质量。
End-to-end Planner Training for Language Modeling
- 用标签概率加权平均嵌入，实现规划与模型的可微衔接。
- 在多个数据集上，困惑度降低1.2~3.8点，效果稳定提升。
- 适合想优化生成模型结构的研究者或工业部署团队。
通过端到端训练预测下一个词元，大语言模型已成为多种任务的重要工具。提升语言建模的核心训练方式，可改善众多下游应用。现有有效方法是引入独立的规划模块，预测未来句子的抽象标签，并以此条件化语言模型。但该方法不可微分，无法联合端到端优化规划模块与语言模型。本文提出一种新方法，实现规划模块与语言模型的联合微调。发现直接使用直通估计器近似标签选择梯度无效；为此，我们建议以预测标签概率作为混合权重，以可微方式将语言模型条件化于标签嵌入的加权平均。这不仅支持联合训练，还使语言模型能利用规划模块输出的完整标签分布，保留更多信息。实验表明，该方法在多个基准数据集上均显著降低困惑度，相对改进达1.2至3.8点。
原文摘要 · Abstract (English)
Through end-to-end training to predict the next token, LLMs have become valuable tools for various tasks. Enhancing their core training in language modeling can improve numerous downstream applications. A successful approach to enhance language modeling uses a separate planning module to predict abstract labels of future sentences and conditions the LM on these predictions. However, this method is non-differentiable, preventing joint end-to-end tuning of the planner with the LM. We propose an effective method to improve this approach by enabling joint fine-tuning of the planner and the LM. We show that a naive way of approximating the gradient of selecting a label via the straight-through estimator is not effective. Instead, we propose to use the predicted label probabilities as mixing weights to condition the LM on a weighted average of label embeddings in a differentiable manner. This not only enables joint fine-tuning of the planner and the LM, but also allows the LM to draw on the full label distribution predicted by the planner, retaining more information. Our experimental results show consistent improvements in perplexity.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。