Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Cho, Hanseul, Cha, Jaeyoung, Bhojanapalli, Srinadh, Yun, Chulhee
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866912332030935040
author Cho, Hanseul
Cha, Jaeyoung
Bhojanapalli, Srinadh
Yun, Chulhee
author_facet Cho, Hanseul
Cha, Jaeyoung
Bhojanapalli, Srinadh
Yun, Chulhee
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2410_15787
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count
Cho, Hanseul
Cha, Jaeyoung
Bhojanapalli, Srinadh
Yun, Chulhee
Machine Learning
Artificial Intelligence
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.
title Arithmetic Transformers Can Length-Generalize in Both Operand Length and Count
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2410.15787