通过任务提示改进 Transformer 的长度泛化能力

近年来观察到 transformer 在某些类型的推理与算术任务上存在长度泛化问题。特别地，在某一长度（例如5位数字）以内任务（如加法）上训练的 transformer 模型，在应用于同一问题的更长实例时性能急剧下降。本文提出一种基于任务提示（task hinting）的方法来应对长度泛化问题。我们的核心思想是：在针对特定任务数据训练模型时，同时训练模型求解一个更简单但相关的辅助任务是有益的。我们以经典排序问题为典型示例来评估该方法。我们设计了一个多任务训练框架，并展示任务提示显著改善了长度泛化。对于排序，我们展示可在长度至多 $20$ 的序列数据上训练模型，并将长度 $100$ 序列的测试准确率从低于1%（标准训练）提升至超过92%（通过任务提示）。我们的研究揭示了长度泛化的若干有趣方面。我们观察到，尽管若干辅助任务先验看来自然，但它们在改善长度泛化上的有效性差异巨大。我们进一步使用探测（probing）与基于可视化的技术来理解模型执行任务的内部机制，并提出与所观察模型学习行为一致的理论构造。基于我们的构造，我们展示在训练过程中引入少量长度相关参数可进一步提升在未见长度上的性能。最后，我们还展示了任务提示方法在排序之外的有效性，令人期待这些技术可应用于更广的语境。