Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Overcoming a Theoretical Limitation of Self-Attention

David Chiang, Peter Cholak

TL;DR

The paper investigates a theoretical limitation of self-attention in transformers for simple regular languages, showing that decisions can become less confident as input length grows. It provides three remedies: exact transformer constructions that achieve perfect accuracy for PARITY and FIRST, a layer-normalization-based method to drive cross-entropy to near zero, and a log-length scaled attention approach to fix learnability and improve length generalization in tasks like machine translation. The results separate expressivity, cross-entropy, and learnability, illustrating that low cross-entropy and successful learning need not coincide. The findings offer practical tools for improving length generalization and addressing limitations in attention-based models.

Abstract

Although transformers are remarkably effective for many tasks, there are some surprisingly easy-looking regular languages that they struggle with. Hahn shows that for languages where acceptance depends on a single input symbol, a transformer's classification decisions become less and less confident (that is, with cross-entropy approaching 1 bit per string) as input strings get longer and longer. We examine this limitation using two languages: PARITY, the language of bit strings with an odd number of 1s, and FIRST, the language of bit strings starting with a 1. We demonstrate three ways of overcoming the limitation suggested by Hahn's lemma. First, we settle an open question by constructing a transformer that recognizes PARITY with perfect accuracy, and similarly for FIRST. Second, we use layer normalization to bring the cross-entropy of both models arbitrarily close to zero. Third, when transformers need to focus on a single position, as for FIRST, we find that they can fail to generalize to longer strings; we offer a simple remedy to this problem that also improves length generalization in machine translation.

Overcoming a Theoretical Limitation of Self-Attention

TL;DR

The paper investigates a theoretical limitation of self-attention in transformers for simple regular languages, showing that decisions can become less confident as input length grows. It provides three remedies: exact transformer constructions that achieve perfect accuracy for PARITY and FIRST, a layer-normalization-based method to drive cross-entropy to near zero, and a log-length scaled attention approach to fix learnability and improve length generalization in tasks like machine translation. The results separate expressivity, cross-entropy, and learnability, illustrating that low cross-entropy and successful learning need not coincide. The findings offer practical tools for improving length generalization and addressing limitations in attention-based models.

Abstract

Although transformers are remarkably effective for many tasks, there are some surprisingly easy-looking regular languages that they struggle with. Hahn shows that for languages where acceptance depends on a single input symbol, a transformer's classification decisions become less and less confident (that is, with cross-entropy approaching 1 bit per string) as input strings get longer and longer. We examine this limitation using two languages: PARITY, the language of bit strings with an odd number of 1s, and FIRST, the language of bit strings starting with a 1. We demonstrate three ways of overcoming the limitation suggested by Hahn's lemma. First, we settle an open question by constructing a transformer that recognizes PARITY with perfect accuracy, and similarly for FIRST. Second, we use layer normalization to bring the cross-entropy of both models arbitrarily close to zero. Third, when transformers need to focus on a single position, as for FIRST, we find that they can fail to generalize to longer strings; we offer a simple remedy to this problem that also improves length generalization in machine translation.
Paper Structure (25 sections, 3 theorems, 48 equations, 4 figures, 1 table)

This paper contains 25 sections, 3 theorems, 48 equations, 4 figures, 1 table.

Table of Contents

  1. Introduction
  2. Background
  3. Notation
  4. Transformers
  5. Input layer
  6. Encoder layers
  7. Output layer
  8. Exact Solutions
  9. FFNN for PARITY
  10. Transformer for PARITY
  11. Transformer for FIRST
  12. Experiments
  13. Layer Normalization
  14. Removing centering
  15. Reducing cross-entropy
  16. ...and 10 more sections

Key Result

Proposition 1

There is a transformer encoder with sigmoid output layer that recognizes (in the above sense) the language $\textsf{PARITY}$ for strings of arbitrary length.

Figures (4)

  • Figure 1: Piecewise linear function equivalent on the integers to $\mathbb{I}[i=k]$.
  • Figure 2: Cross-entropy of exact solutions for PARITY and FIRST computed over 1000 random strings of length $n$. Without layer norm, the cross-entropy quickly approaches its upper bound of one. With layer norm and $\epsilon > 0$, the cross-entropy is better but still grows with $n$. With $\epsilon = 0$, cross-entropy is independent of $n$.
  • Figure 3: The cross-entropy and accuracy of our solution to PARITY are both extremely sensitive to the parameter $\bar{\mathbf{W}}^{\text{V},1,1}_{6,2}$, which is responsible for computing $\frac{k}{n}$. The correct parameter value is 1.
  • Figure 4: Training a transformer on FIRST. Each epoch has 100 training strings of varying length (see legend) and 100 test strings of length 1000. All curves are averaged over 20 runs. Left: Standard transformer with layer normalization ($\epsilon=10^{-5}$). Right: Same, with attention logits scaled by $\log n$.

Theorems & Definitions (5)

  • Proposition 1
  • Proposition 2
  • proof
  • Proposition 3
  • proof