PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformers are Multi-State RNNs

Matanel Oren, Michael Hassid, Nir Yarden, Yossi Adi, Roy Schwartz

arXiv:2401.06104v2cs.CL

TL;DR

Decoder-only transformers are commonly viewed as conceptually distinct from RNNs, while their key-value cache is a decoding memory bottleneck. This paper recasts them as multi-state RNNs and introduces TOVA, which retains states by attention score; across long-range tasks, it reports TOVA outperforming compression baselines and often approaching full-model performance with a smaller state.

  • Problem

    Transformers are considered conceptually distinct from RNNs, and the key-value cache is a memory bottleneck during LLM decoding.

    Source: §1 · p. 2

  • Method

    The paper models decoder transformers as multi-state RNNs and introduces TOVA, which retains states with the highest attention scores.

    Source: §1 · p. 1, §9 · p. 8

  • Results

    Across four tasks, TOVA outperforms existing compression policies; in 2/4 tasks, 1/8–1/4-size states yield results comparable to the topline, while the other two require half-size states for comparable performance.

    Source: §6.4 · p. 6, §9 · p. 8

  • Takeaways & Limitations

    The findings connect transformers with RNNs and can reduce cache size by up to 88% while increasing throughput by 4.8X.

    Source: §9 · p. 9

    Long-text-generation quality comparisons rely on GPT-4 and may not capture the full breadth of text quality.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Transformers are considered conceptually different from the previous generation of state-of-the-art NLP models - recurrent neural networks (RNNs). In this work, we demonstrate that decoder-only transformers can in fact be conceptualized as unbounded multi-state RNNs - an RNN variant with unlimited hidden state size. We further show that transformers can be converted into $\textit{bounded}$ multi-state RNNs by fixing the size of their hidden state, effectively compressing their key-value cache. We introduce a novel, training-free compression policy - $\textbf{T}$oken $\textbf{O}$mission $\textbf{V}$ia $\textbf{A}$ttention (TOVA). Our experiments with four long range tasks and several LLMs show that TOVA outperforms several baseline compression policies. Particularly, our results are nearly on par with the full model, using in some cases only $\frac{1}{8}$ of the original cache size, which translates to 4.8X higher throughput. Our results shed light on the connection between transformers and RNNs, and help mitigate one of LLMs' most painful computational bottlenecks - the size of their key-value cache. We publicly release our code at https://github.com/schwartz-lab-NLP/TOVA

1 Introduction

The paper reframes decoder-only transformers as unbounded multi-state RNNs and proposes compressing their state with attention-guided token retention. It reports that TOVA can approach full-model performance with a smaller state, while the background sections contrast recurrent state updates with transformer self-attention.

1 Introduction

  • Decoder-only transformers are framed as unbounded multi-state RNNs, with a state for each history token that grows as decoding proceeds.The paper relates autoregressive state preservation to the core principle of RNNs.Source: §1 · p. 1
  • TOVA bounds the state by retaining tokens with the highest attention scores, while the paper also examines which tokens remain in memory.The analysis finds that not all recent tokens are important to retain and that the first token and some possessive endings may remain.Source: §1 · pp. 1–2
  • Using 1/8 of the full multi-state, TOVA can match full-model performance, with throughput up to 4.8X higher.The introduction also reports that TOVA outperforms all baseline compression policies in all setups and supports contexts up to 70K tokens.Source: §1 · p. 2

2.1 RNNs

  • An RNN layer uses the current token representation and previous hidden state to produce an updated representation and a new state.The new state is used in the recurrent computation for the next token.Source: §2.1 · p. 2

2.2 Transformers

  • A transformer layer applies self-attention across the sequence and a feed-forward operation to each token; decoder attention is masked for next-token prediction.During decoding, key and value matrices are commonly cached to avoid recomputing earlier tokens.Source: §2.2 · p. 2

3 Transformers as Multi-State RNNs

The paper defines multi-state RNNs and shows how decoder transformers fit this framework: their key-value states grow with tokens, but can be bounded by a compression policy. It argues that LLMs are unbounded at inference because they accumulate tokens without dropping them, although they may behave as bounded systems in practice.

3.1 Multi-State RNNs

  • An MSRNN uses a matrix rather than a vector as its state, with each row interpreted as a single state; limiting state count bounds its memory capacity.A constant state count yields bounded capacity, whereas an unbounded count permits unbounded capacity.Source: §3.1 · p. 2

3.2 Transformers are Unbounded MSRNNs

  • With the number of states set to the number of input tokens, a transformer is an unbounded MSRNN whose key-value states correspond to tokens.The state grows as tokens arrive.Source: §3.2 · pp. 2–3

3.3 Converting Transformers into Bounded MSRNNs

  • Setting the state count to min(t, k) bounds a transformer’s multi-state, requiring a compression policy once the token count exceeds k.Existing KV-cache methods, including windowed attention and H2O, can be viewed as such policies.Source: §3.3 · p. 3

3.4 LLMs as MSRNNs

  • LLMs are treated as unbounded because inference can process any number of tokens subject to available memory, and their multi-state accumulates tokens without dropping them.The paper notes that they can nevertheless act in practice as bounded MSRNNs.Source: §3.4 · p. 3

4 TOVA: Token Omission Via Attention

TOVA is a training-free policy for compressing an unbounded multi-state RNN into a bounded one. Once capacity is reached, it removes the state with the lowest attention score at each decoding step.

4 TOVA: Token Omission Via Attention

  • TOVA converts an unbounded multi-state into a bounded one by applying a state-compression policy.The policy is training-free.Source: §4 · p. 3
  • After capacity is reached, TOVA drops the token with the lowest attention score at each decoding step.The illustration shows the omitted state becoming transparent in subsequent steps.Source: §4 · p. 3, Figure 2 · p. 3
  • TOVA computes attention scores separately for each head, while preliminary results favor averaging scores across a layer’s heads.Averaging is reported as superior to treating each head individually.Source: §4 · p. 3

5 Experimental Setup

The experiments compare pretrained decoder LLMs under compression policies and evaluate long-range language modeling, understanding, and generation. They use several model families and task-specific datasets, with a fixed 4,096-token training context length.

5 Experimental Setup

  • The setup tests whether bounded MSRNN versions of transformer LLMs can match the full model, using pretrained LLMs and policies requiring no additional training.The authors frame the connection between compression policies and RNNs as novel to their knowledge.Source: §5 · p. 3, §5.1 · p. 3

5.1 Baseline Compression Policies

  • Window discards the oldest state, Window+i also retains early states, and H2O dynamically retains high-attention earlier tokens alongside a recent-token window.H2O typically assigns half of its multi-state capacity to non-window tokens; the full unbounded model serves as the topline.Source: §5.1 · pp. 3–4
  • Unlike baselines that favor recent or early token states, TOVA makes neither a fixed recent-token window nor an early-token preference.The authors characterize the baseline policies as introducing strong inductive biases.Source: §5.1 · p. 4

5.2 Long Range Evaluation

  • Evaluation covers language modeling, long-range understanding, and text generation, using PG-19, SQuALITY, QASPER, and GPT-4 story preferences.The tasks use perplexity, geometric-mean ROUGE-1/2/L, F1, and average GPT-4 win rate, respectively.Source: §5.2 · p. 4

5.3 Models

  • The experiments use roughly 7B LLaMA-2, Mistral, and Yi language models, instruction-tuned models for understanding, and Mytho-Logic for story generation.All models use a 4,096-token training sequence length; long examples are chunked or truncated according to task.Source: §5.3 · p. 4

6 Results: Pretrained Transformers Often Act as Bounded MSRNNs

Across four long-range tasks, TOVA often approaches full-model performance with a smaller multi-state, though retrieval QA and text generation require more capacity. Results vary by task, with some matching the topline at one-eighth to one-quarter capacity and others at about half.

6.1 Language Modeling

  • On PG-19, TOVA is within 0.4 perplexity points of the topline at one-eighth context and consistently outperforms the baselines.The baselines require at least half of the full context length to reach full-model results.Source: §6.1 · p. 5, Figure 3 · p. 5

6.2 Long Range Understanding

  • On SQuALITY, TOVA beats all baselines and reaches within one point of the topline using one-quarter context for Mistral and Yi or one-eighth for LLaMA-2.The reported metric is the geometric mean of ROUGE-1/2/L.Source: §6.2 · p. 5, Figure 4 · p. 5
  • On QASPER, TOVA outperforms baselines across multi-state sizes and comes within one F1 point of the topline at half context.The gap over baselines exceeds 5 F1 points in some cases.Source: §6.2 · pp. 5–6, Figure 5 · p. 5

6.3 Text Generation

  • Average story length is 1,566 tokens for the full model and falls to 1,503 at 512-token capacity and 1,361 at 256.Length remains 1,566 tokens at a multi-state size of 1,024.Source: §6.3 · p. 6
  • At 256 tokens, TOVA stories lose to the topline in 47% of GPT-4 comparisons; the loss rate falls to 19% at 512 and 6% at 1,024 tokens.TOVA is preferred over the topline in 5–10% of cases at each tested size.Source: §6.3 · p. 6, Figure 6 · p. 6

6.4 Discussion

  • In 2/4 tasks, TOVA matches the topline with 1/8–1/4 of the multi-state size; text generation and retrieval QA require larger sizes but remain comparable at half-size.The authors suggest that conversion to an RNN reintroduces difficulties in retrieving distant information (Arjovsky et al., 2016; Jelassi et al., 2024).Source: §6.4 · p. 6

7 Analysis

The analysis finds that TOVA extrapolates to 70K tokens while outperforming Window+4, and examines memory efficiency and which tokens it retains. Recent tokens matter, but TOVA also keeps selected older tokens, including the first token and some content-specific tokens.

7.1 TOVA is Time- and Memory-Efficient

  • Table 1 compares cache memory at batch size 1, maximum batch size, and relative throughput, showing reductions in memory alongside increases in batch size and throughput.Its maximum-batch and throughput rows use decoding of the same number of tokens on a single V100 machine.Source: Table 1 · p. 6

7.2 Extrapolation with TOVA

  • For extrapolation, TOVA compresses gaps between cached token positions with ln(ln(g)), leaving gaps of 10 or less unchanged.Preliminary tests of ln(g) and sqrt(g) performed worse.Source: §7.2 · pp. 6–7
  • Up to 70K tokens, TOVA maintains performance within 0.5 PPL points of shorter contexts and outperforms Window+4.The comparison averages perplexity over the first 70K tokens of 52 PG-19 books, using a multi-state size of 512.Source: Figure 7 · p. 7, §7.2 · p. 7

7.3 Which Tokens Matter?

  • Only 73–76% of kept tokens are recent, so TOVA retains older tokens as well as a clear window of recent ones.The recent-token proportion is averaged across examples, layers, and positions; TOVA identifies recent tokens automatically rather than using a handcrafted window.Source: §7.3 · p. 7
  • Across multi-state sizes, TOVA keeps the first token until the sequence ends, while other early tokens are dropped sooner.Figure 9 averages the number of decoding steps kept across layers and examples.Source: §7.3 · p. 7, Figure 9 · p. 7
  • Alongside punctuation and special symbols, possessive endings and proper nouns tend to remain in memory longer.These token categories are among those kept longest in the reported part-of-speech analysis.Source: §7.3 · p. 8, Table 2 · p. 8

8 Related Work

Related work connects transformers with recurrence, limits KV-cache size, develops new RNN variants, and simplifies transformer attention. The paper distinguishes its approach by mapping tokens to explicit states and applying compression to existing LLMs without dedicated training.

Transformers and RNNs

  • Earlier transformer–RNN approaches use hybrids, convolution, inference-time recurrence, or repeated layers, while related bounded-memory work treats memory as a single state rather than mapping tokens to states.The paper also contrasts its approach with works that require dedicated training and cannot operate on existing LLMs.Source: §8 · p. 8

Limited KV cache

New RNN variants and Simplifying transformers

  • Other work revives RNNs through state-space and trainable RNN variants, while transformer simplification prunes heads or replaces attention with efficient alternatives.The paper positions bounded MSRNNs as another way to reduce transformer decoders.Source: §8 · p. 8

9 Conclusion

The conclusion frames decoder transformers as unbounded multi-state RNNs and presents TOVA as a way to bound their state by selecting tokens with attention scores. TOVA performs comparably to the unbounded model at smaller state sizes and supports inputs up to 70K tokens.

9 Conclusion

  • Decoder transformers can be viewed as unbounded multi-state RNNs and compressed to bounded versions by limiting tokens handled at each decoding step.The conclusion presents this as a connection between transformer operation and RNNs.Source: §9 · p. 8
  • TOVA selects tokens by attention scores and is reported as superior to existing compression policies, with performance comparable to the full model in many cases.Those cases require 1/8–1/4 of the multi-state size.Source: §9 · pp. 8–9
  • TOVA supports processing inputs up to 70K tokens and is reported to reduce cache size by up to 88% and increase throughput by 4.8X.Source: §9 · p. 9

Limitations

The paper notes that long-text evaluation is costly and difficult to reproduce, and that GPT-4 comparisons may not capture the full breadth of text quality. Its evaluation focuses on English; languages with more flexible word order may require larger multi-state sizes.

Limitations

  • Long-text generation evaluation is computationally expensive, and GPT-4 comparison may miss aspects of text quality.The authors say this cost might limit reproduction of their results.Source: Limitations · p. 9
  • The evaluation framework focuses on English, and languages with more flexible word order may require a larger multi-state size.Source: Limitations · p. 9

Ethics Statement

  • The authors say reducing transformer LLM memory could potentially increase adoption among users with limited hardware access.Source: Ethics Statement · p. 9
  • The work uses no newly collected data, relying instead on open-source models and public data collected by other sources.Source: Ethics Statement · p. 9

A Policy Ablation

The ablation compares cache-compression policies and TOVA variants, and finds that layer-wise TOVA performs best. The section also describes its implementation, prompts, evaluation protocol, and experimental setup.

A Policy Ablation

  • The ablation tests Window, Window+1/+4, head- and layer-wise H2O and TOVA, and TOVA variants preserving initial tokens, using LLaMA-2-7B.The initial-token variants preserve one or four tokens.Source: A Policy Ablation · p. 12

B Formal Description of Method

C Prompts

D Details of Generation Evaluation

  • GPT-4 evaluates generated-text pairs in both presentation orders, and a win counts only when the same approach is preferred in both.The evaluation uses gpt-4-0613 and drops cases where generation stops before the memory limit because the stories are identical.Source: D Details of Generation Evaluation · p. 12

E Experimental Details

  • Experiments use bfloat16 on Nvidia V100 GPUs, with static masks for Window policies and attention-weight-adjusted masks for H2O and TOVA.The mask modifications parallelize language modeling over sequence tokens.Source: E Experimental Details · pp. 12–13

A Policy Ablation

  • TOVA dominates the PG-19 perplexity ablation across all tested multi-state sizes, and layer-wise TOVA outperforms the other policies while head-wise TOVA performs worse.The authors attribute the layer-wise advantage to requiring agreement among heads when selecting important tokens; explicitly preserving initial tokens leaves results relatively unchanged.Source: A Policy Ablation · p. 12, Table 3 · p. 13

F Long Range Understanding with Base Models

The section reports base-model results on SQuALITY and QASPER, illustrates TOVA's retained tokens across layers and decoding steps, and provides fuller part-of-speech analysis.

F Long Range Understanding with Base Models

G Illustration of the Tokens Retained by

H Full Part-of-Speech Tag Analysis

Loading 2401.06104v2…