Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Cho, Hanseul, Cha, Jaeyoung, Awasthi, Pranjal, Bhojanapalli, Srinadh, Gupta, Anupam, Yun, Chulhee
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917822929567744
author Cho, Hanseul
Cha, Jaeyoung
Awasthi, Pranjal
Bhojanapalli, Srinadh
Gupta, Anupam
Yun, Chulhee
author_facet Cho, Hanseul
Cha, Jaeyoung
Awasthi, Pranjal
Bhojanapalli, Srinadh
Gupta, Anupam
Yun, Chulhee
contents Even for simple arithmetic tasks like integer addition, it is challenging for Transformers to generalize to longer sequences than those encountered during training. To tackle this problem, we propose position coupling, a simple yet effective method that directly embeds the structure of the tasks into the positional encoding of a (decoder-only) Transformer. Taking a departure from the vanilla absolute position mechanism assigning unique position IDs to each of the tokens, we assign the same position IDs to two or more "relevant" tokens; for integer addition tasks, we regard digits of the same significance as in the same position. On the empirical side, we show that with the proposed position coupling, our models trained on 1 to 30-digit additions can generalize up to 200-digit additions (6.67x of the trained length). On the theoretical side, we prove that a 1-layer Transformer with coupled positions can solve the addition task involving exponentially many digits, whereas any 1-layer Transformer without positional information cannot entirely solve it. We also demonstrate that position coupling can be applied to other algorithmic tasks such as Nx2 multiplication and a two-dimensional task.
format Preprint
id arxiv_https___arxiv_org_abs_2405_20671
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure
Cho, Hanseul
Cha, Jaeyoung
Awasthi, Pranjal
Bhojanapalli, Srinadh
Gupta, Anupam
Yun, Chulhee
Machine Learning
Artificial Intelligence
Computation and Language
Even for simple arithmetic tasks like integer addition, it is challenging for Transformers to generalize to longer sequences than those encountered during training. To tackle this problem, we propose position coupling, a simple yet effective method that directly embeds the structure of the tasks into the positional encoding of a (decoder-only) Transformer. Taking a departure from the vanilla absolute position mechanism assigning unique position IDs to each of the tokens, we assign the same position IDs to two or more "relevant" tokens; for integer addition tasks, we regard digits of the same significance as in the same position. On the empirical side, we show that with the proposed position coupling, our models trained on 1 to 30-digit additions can generalize up to 200-digit additions (6.67x of the trained length). On the theoretical side, we prove that a 1-layer Transformer with coupled positions can solve the addition task involving exponentially many digits, whereas any 1-layer Transformer without positional information cannot entirely solve it. We also demonstrate that position coupling can be applied to other algorithmic tasks such as Nx2 multiplication and a two-dimensional task.
title Position Coupling: Improving Length Generalization of Arithmetic Transformers Using Task Structure
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2405.20671