Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Single Headed Attention RNN: Stop Thinking With Your Head

Stephen Merity

TL;DR

This work challenges the Transformer-dominated view of language modeling by presenting the Single Headed Attention RNN (SHA-RNN), a memory-augmented LSTM-based model with a single attention head and a Boom feed-forward layer. It demonstrates competitive byte-level language modeling on enwik8 using a single GPU and minimal hyperparameter tuning, arguing that alternative architectures can achieve strong results with efficient compute. The paper explores tokenization attacks, optimization adjustments like minimum-trust LAMB, and the concept of over-parameterizing static learned vectors to provide training history, highlighting practical considerations and community impact. Overall, it advocates architectural diversity and accessible research paths that could spur distillation and broader exploration beyond Transformers.

Abstract

The leading approaches in language modeling are all obsessed with TV shows of my youth - namely Transformers and Sesame Street. Transformers this, Transformers that, and over here a bonfire worth of GPU-TPU-neuromorphic wafer scale silicon. We opt for the lazy path of old and proven techniques with a fancy crypto inspired acronym: the Single Headed Attention RNN (SHA-RNN). The author's lone goal is to show that the entire field might have evolved a different direction if we had instead been obsessed with a slightly different acronym and slightly different result. We take a previously strong language model based only on boring LSTMs and get it to within a stone's throw of a stone's throw of state-of-the-art byte level language model results on enwik8. This work has undergone no intensive hyperparameter optimization and lived entirely on a commodity desktop machine that made the author's small studio apartment far too warm in the midst of a San Franciscan summer. The final results are achievable in plus or minus 24 hours on a single GPU as the author is impatient. The attention mechanism is also readily extended to large contexts with minimal computation. Take that Sesame Street.

Single Headed Attention RNN: Stop Thinking With Your Head

TL;DR

This work challenges the Transformer-dominated view of language modeling by presenting the Single Headed Attention RNN (SHA-RNN), a memory-augmented LSTM-based model with a single attention head and a Boom feed-forward layer. It demonstrates competitive byte-level language modeling on enwik8 using a single GPU and minimal hyperparameter tuning, arguing that alternative architectures can achieve strong results with efficient compute. The paper explores tokenization attacks, optimization adjustments like minimum-trust LAMB, and the concept of over-parameterizing static learned vectors to provide training history, highlighting practical considerations and community impact. Overall, it advocates architectural diversity and accessible research paths that could spur distillation and broader exploration beyond Transformers.

Abstract

The leading approaches in language modeling are all obsessed with TV shows of my youth - namely Transformers and Sesame Street. Transformers this, Transformers that, and over here a bonfire worth of GPU-TPU-neuromorphic wafer scale silicon. We opt for the lazy path of old and proven techniques with a fancy crypto inspired acronym: the Single Headed Attention RNN (SHA-RNN). The author's lone goal is to show that the entire field might have evolved a different direction if we had instead been obsessed with a slightly different acronym and slightly different result. We take a previously strong language model based only on boring LSTMs and get it to within a stone's throw of a stone's throw of state-of-the-art byte level language model results on enwik8. This work has undergone no intensive hyperparameter optimization and lived entirely on a commodity desktop machine that made the author's small studio apartment far too warm in the midst of a San Franciscan summer. The final results are achievable in plus or minus 24 hours on a single GPU as the author is impatient. The attention mechanism is also readily extended to large contexts with minimal computation. Take that Sesame Street.

Paper Structure

This paper contains 28 sections, 2 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Language is humanity's longest running program
  3. Direct motivation
  4. Motivation
  5. Alternate history and the genre of research fiction
  6. Fighting the direction of research
  7. A lone GPU
  8. Model architecture
  9. A simplified attention mechanism
  10. Boom layer
  11. Experiments
  12. Hutter Wikipedia Prize (enwik8)
  13. WikiText
  14. Wordpiece WikiText-103
  15. Results and analysis
  16. ...and 13 more sections

Figures (3)

  • Figure 1: The SHA-RNN is composed of an RNN, pointer based attention, and a "Boom" feed-forward with a sprinkling of layer normalization. The persistent state is the RNN's hidden state $h$ as well as the memory $M$ concatenated from previous memories. Bake at 200$^{\circ}$F for 16 to 20 hours in a desktop sized oven.
  • Figure 2: The attention mechanism within the SHA-RNN is highly computationally efficient. The only matrix multiplication acts on the query. The $A$ block represents scaled dot product attention, a vector-vector operation. The operators $\{qs, ks, vs\}$ are vector-vector multiplications and thus have minimal overhead. We use a sigmoid to produce $\{qs, ks\}$. For $vs$ see Section \ref{['sec:opslv']}.
  • Figure 3: Bits Per Character (BPC) reported on the enwik8 validation set. Each epoch took approximately 4020 seconds (67 minutes). The learning rate was cut in half at epoch 16.