Transformers often struggle with length generalization, meaning they fail to generalize to sequences longer than those encountered during training. While arithmetic tasks are commonly used to study length generalization, certain tasks are considered notoriously difficult, e.g., multi-operand addition (requiring generalization over both the number of operands and their lengths) and multiplication (requiring generalization over both operand lengths). In this work, we achieve approximately 2-3x length generalization on both tasks, which is the first such achievement in arithmetic Transformers. We design task-specific scratchpads enabling the model to focus on a fixed number of tokens per each next-token prediction step, and apply multi-level versions of Position Coupling (Cho et al., 2024; McLeish et al., 2024) to let Transformers know the right position to attend to. On the theory side, we prove that a 1-layer Transformer using our method can solve multi-operand addition, up to operand length and operand count that are exponential in embedding dimension.

本研究解决了变压器在长度泛化方面的不足，特别是在多操作数加法和乘法任务中。通过设计特定任务的临时缓存和多层位置耦合的方法，我们首次在算术变压器上实现了约2-3倍的长度泛化。该工作可能在推进算术理解模型能力方面产生重要影响。

算术变压器可以在操作数长度和数量上实现长度泛化