The Role of Sparsity for Length Generalization in Transformers

TL;DR

This paper develops a theoretical framework showing that length generalization in transformers depends on sparse token dependencies, and introduces a new training method to enhance this ability in language models.

Contribution

It formalizes the concept of sparse dependencies for length generalization and proposes Predictive Position Coupling to improve transformer performance on longer sequences.

Findings

Length generalization occurs with fixed small token dependencies.

Sparse dependency structure is key to successful length generalization.

Predictive Position Coupling broadens tasks for effective length generalization.

Abstract

Training large language models to predict beyond their training context lengths has drawn much attention in recent years, yet the principles driving such behavior of length generalization remain underexplored. We propose a new theoretical framework to study length generalization for the next-token prediction task, as performed by decoder-only transformers. Conceptually, we show that length generalization occurs as long as each predicted token depends on a small (fixed) number of previous tokens. We formalize such tasks via a notion we call $k$-sparse planted correlation distributions, and show that an idealized model of transformers which generalize attention heads successfully length-generalize on such tasks. As a bonus, our theoretical model justifies certain techniques to modify positional embeddings which have been introduced to improve length generalization, such as position coupling. We support our theoretical results with experiments on synthetic tasks and natural language, which confirm that a key factor driving length generalization is a ``sparse'' dependency structure of each token on the previous ones. Inspired by our theory, we introduce Predictive Position Coupling, which trains the transformer to predict the position IDs used in a positional coupling approach. Predictive Position Coupling thereby allows us to broaden the array of tasks to which position coupling can successfully be applied to achieve length generalization.