arXiv:2411.02083cs.CLcs.AI2024-11被引 14

用回归思想优化语言模型生成数字,提升数学推理能力。

Regress, Don't Guess -- A Regression-like Loss on Number Tokens for Language Models

  • 设计基于数值接近度的新型损失函数,直接作用于数字标记
  • 在多个数学数据集上显著提升算术任务表现,媲美专用回归头
  • 可无缝集成至任意大模型,零运行开销,适合需精准数字输出场景

尽管语言模型在文本生成方面表现出色,但在涉及数量推理的任务中(尤其是算术)存在天然缺陷,根源在于交叉熵损失假设标记为名义尺度,无法表达生成数字之间的接近程度。为此,本文提出一种纯标记层面的回归式损失——数字标记损失(NTL),包含两种形式:最小化真实与预测数字标记间 $L_p$ 范数或沃尔德斯坦距离。该损失可轻松添加至任意语言模型训练中,扩展交叉熵目标,且无运行时开销。我们在多个数学数据集上评估,结果表明其能持续提升相关任务性能;在直接对比回归任务中,NTL表现堪比专用回归头。将模型规模扩展至30亿参数后仍见性能提升,证明其可无缝集成至大语言模型。我们希望激励开发者改进预训练目标,并发布轻量级 PyPI 包 ntloss: https://github.com/ai4sd/number-token-loss。完整论文复现代码单独提供。

原文摘要 · Abstract (English)

While language models have exceptional capabilities at text generation, they lack a natural inductive bias for emitting numbers and thus struggle in tasks involving quantitative reasoning, especially arithmetic. One fundamental limitation is the nature of the cross-entropy (CE) loss, which assumes a nominal scale and thus cannot convey proximity between generated number tokens. In response, we here present a regression-like loss that operates purely on token level. Our proposed Number Token Loss (NTL) comes in two flavors and minimizes either the $L_p$ norm or the Wasserstein distance between the numerical values of the real and predicted number tokens. NTL can easily be added to any language model and extend the CE objective during training without runtime overhead. We evaluate the proposed scheme on various mathematical datasets and find that it consistently improves performance in math-related tasks. In a direct comparison on a regression task, we find that NTL can match the performance of a regression head, despite operating on token level. Finally, we scale NTL up to 3B parameter models and observe improved performance, demonstrating its potential for seamless integration into LLMs. We hope to inspire LLM developers to improve their pretraining objectives and distribute NTL as a minimalistic and lightweight PyPI package $ntloss$: https://github.com/ai4sd/number-token-loss. Development code for full paper reproduction is available separately.

语言模型数字生成损失函数数学推理

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