变换器通过合适的嵌入实现算术运算
机器学习
2024-12-24 v2 人工智能
摘要
变换器在算术任务上的表现较差,主要原因是它们难以跟踪大范围数字中每个数字的确切位置。我们通过为每个数字添加一个嵌入向量来解决这个问题,该嵌入向量编码了该数字相对于数开头位置的相对位置。除了这些嵌入本身提供的提升外,我们还展示,这一修复方法使我们能够实现架构修改,如输入注入和循环层,从而进一步提高性能。当位置问题得到解决后,我们可以研究变换器的逻辑外推能力。它们能否解决比训练数据中更大更复杂的算术问题?我们发现,只需在单个 GPU 上训练 20 位数字的训练数据一天,就能达到最新的性能水平,在 100 位加法问题中实现最高 99% 的准确率。最后,我们表明,这些数值能力的提升也 unlocks 对其他多步骤推理任务的改进,包括排序和乘法。
引用
@article{arxiv.2405.17399,
title = {Transformers Can Do Arithmetic with the Right Embeddings},
author = {Sean McLeish and Arpit Bansal and Alex Stein and Neel Jain and John Kirchenbauer and Brian R. Bartoldson and Bhavya Kailkhura and Abhinav Bhatele and Jonas Geiping and Avi Schwarzschild and Tom Goldstein},
journal= {arXiv preprint arXiv:2405.17399},
year = {2024}
}