LongT5: Efficient Text-To-Text Transformer for Long Sequences

In plain words

Transformers get expensive on long documents, because standard attention lets every word look at every other word. LongT5 is a T5 text-to-text model whose encoder is changed so that each word looks mainly at its close neighbours. To still let distant parts of a document influence each other, the model builds a few summary "global" tokens on the fly from fixed-size blocks of the input. These summaries are created for each attention step and thrown away afterwards, so the model needs no extra input slots that a person must design.

The authors also change how the model is pretrained. Instead of only hiding random word spans, they hide whole important sentences and ask the model to write them back, a task borrowed from summarization pretraining.

The authors scale both the input length and the model size, and report state-of-the-art results on several summarization and question-answering benchmarks. These numbers are their own measurements.

TL;DR

LongT5 is a T5.1.1-style encoder-decoder whose encoder uses sparse attention: a local window of radius 127 plus one summary token per 16-token block. It keeps the standard T5 decoder and compatibility with T5 checkpoints §3.1. It is pretrained on C4 with PEGASUS-style principle-sentence generation. Author-reported results beat the compared baselines on arXiv, PubMed, BigPatent, MediaSum and TriviaQA, and are second on Multi-News and roughly comparable on CNN/Daily Mail §4.2.2, §4.3.2.

Problem

Full self-attention costs grow quadratically with input length, which limits T5 to short inputs and makes long-input training slow or out of memory §2. Earlier long-input models (ETC, BigBird, Longformer) need extra global side inputs designed per task, and their pretraining is usually masked-language style rather than generative §1, §6. The paper assumes that scaling input length and model size together is worth doing, and tests this on tasks whose inputs are long (arXiv, PubMed, TriviaQA) §1, §4.2.1.

Method

Architecture §3.1. The encoder of T5 is extended with sparse attention. The decoder is a standard T5 decoder, because the target sequences are short §3.1. Both variants retain T5's relative position encoding, still let several short examples share one input sequence, and load T5 checkpoints §3.1.

Local attention §3.1.1. The encoder self-attention is replaced by a sliding window of radius r. Each token attends to r tokens to its left and r to its right. The paper uses r = 127. This adds no parameters. Cost is O(l × r) for input length l §3.1.1.

Transient Global attention (TGlobal) §3.1.2. The input is cut into fixed blocks of k tokens, with k = 16. For each block, one global token is the sum of the block's token embeddings, followed by normalization. Each input token attends to its local window and to every global token, so the extra work is about l × l/k key-value pairs per layer §3.1.2. Global tokens are rebuilt in every attention layer and are not stored. Two small pieces of state are added: T5-style relative position biases from a token's block to each global token, and T5-style layer-norm parameters for the globals. The paper counts these as about 10k extra parameters for base, 25k for large and 50k for xl §3.1.2, fn. 1. Packed examples are kept apart by masking attention from input tokens to globals of other examples §3.1.2.

The paper gives the complexity as O(l × (r + l/k)) §3.1.2. The pseudocode below restates the TGlobal layer in our own words. It is not a copy of the official implementation long_attention.py.

# One TGlobal self-attention layer of the encoder  [§3.1.1] [§3.1.2]
# Input:  tokens (B, l, d), segment_ids (B, l)
# Output: attended (B, l, d)

def tglobal_attention(tokens, segment_ids):
    # Queries, keys and values use the T5 projections  [§3.1]
    queries = project_q(tokens)                     # (B, l, h, d_head)
    keys = project_k(tokens)                        # (B, l, h, d_head)
    values = project_v(tokens)                      # (B, l, h, d_head)

    # Cut the sequence into fixed blocks of k = 16 tokens  [§3.1.2]
    blocks = reshape_blocks(tokens, block_size=16)  # (B, l/k, k, d)

    # One global token per block: sum the block, then normalise  [§3.1.2]
    global_tokens = layer_norm(sum_over(blocks))    # (B, l/k, d)
    global_keys = project_k(global_tokens)          # (B, l/k, h, d_head)
    global_values = project_v(global_tokens)        # (B, l/k, h, d_head)

    # Local scores over the window of radius r = 127  [§3.1.1]
    local_scores = window_scores(queries, keys, radius=127)

    # Scores against every global token  [§3.1.2]
    global_scores = queries @ transpose(global_keys)  # (B, l, h, l/k)

    # Packing: a token sees only globals of its own example  [§3.1.2]
    global_segment_ids = block_segment_ids(segment_ids)
    global_scores = mask_other_examples(
        global_scores, segment_ids, global_segment_ids)

    # T5 relative position biases, local and token-to-global  [§3.1.2]
    scores = concat(local_scores, global_scores) + relative_bias

    weights = softmax(scores, axis=-1)
    return combine(weights, values, global_values)

PEGASUS-style pretraining §3.2. Pretraining uses gap-sentence generation with the "Principle, Ind, Uniq" settings of PEGASUS §3.2. Each sentence is scored by the ROUGE-F1 of its content against the rest of the document. Scoring is independent per sentence (Ind), and each n-gram counts once (Uniq). The top-scoring sentences are masked and the model generates them as a single target string §3.2. The masked sentence ratio is 0.2 in these experiments, not the 0.45 of the original PEGASUS §4.1.1. The pseudocode below is our own summary of this selection.

# Principle-sentence generation example  [§3.2] [§4.1.1]
# Input:  document, a list of sentences
# Output: masked_document, target_text

def make_psg_example(document):
    # Score each sentence by ROUGE-F1 against the rest of the document,
    # each sentence scored on its own  [§3.2]
    scores = [rouge1_f1(s, others_of(document, s)) for s in document]
    # Keep the top-scoring sentences; the ratio is 0.2  [§4.1.1]
    chosen = top_fraction(document, scores, ratio=0.2)
    # Mask the chosen sentences; the model must write them as one string
    # [§3.2]
    masked_document = replace_with_mask(document, chosen)
    target_text = join(chosen)
    return masked_document, target_text

Data and tasks §4.2.1, §4.3.1. Summarization: CNN/Daily Mail, PubMed, arXiv, BigPatent, MediaSum, Multi-News. QA: Natural Questions (short answer only, with the long answer ignored) and TriviaQA. Input-length statistics are in tab. 1 and tab. 3.

Hyperparameters

Results

Author-reported. Setup: summarization fine-tuned with per-dataset input lengths; QA on NQ and TriviaQA, with NQ results against T5.1.1 baselines run by the authors §4.2.2, §4.3.2. Each configuration was run once §4.3.2, §5.2.

Summarization ROUGE tab. 2 (main-text rows; the full table is in app. A, tab. 8).

The caption of tab. 2 notes that each dataset uses a different input length, so rows are not controlled for input length.

Question answering tab. 4 (EM / F1; full list in app. B, tab. 7).

Speed and memory §5.1, fig. 3. At input length 2048, base T5.1.1 trains at 479 sequences per second, LongT5 base TGlobal at 765 and LongT5 base Local at 860 §5.1. The longest inputs before out-of-memory on a 4x8 TPUv3 slice are 6k for T5.1.1 base, 12k for TGlobal base and 36k for Local base. For large models: 3k for T5.1.1, 6k for TGlobal and 10k for Local §5.1.

Input length vs NQ F1 §5.2, fig. 4. Large models: F1 rises with input length. Local attention is clearly below TGlobal and T5.1.1 on NQ. The paper attributes single-point dips (such as 6k for TGlobal) to noise from one run §5.2.

Pretraining ablation, NQ and arXiv, TGlobal base, dev set tab. 5.

PSG = principle-sentence generation; SC = span corruption. The paper states that using both objectives together performs worse than PSG alone §5.3.

Pretraining ablation, T5.1.1 at 512 input tab. 6. arXiv R-1 / R-2 / R-3: SC 44.59 / 18.34 / 40.65; PSG 45.78 / 18.94 / 41.53; LongT5 4k 45.66 / 19.22 / 41.49; LongT5 16k 48.21 / 21.7 / 44.03. PubMed R-1 / R-2 / R-3: SC 47.86 / 22.14 / 44.39; PSG 48.74 / 23.42 / 45.24; LongT5 4k 48.47 / 23.38 / 45.01; LongT5 16k 50.12 / 24.78 / 46.56 §5.3.

Figures. fig. 1 plots average ROUGE of LongT5 and four baselines on arXiv and PubMed against input length. fig. 2 illustrates Local and TGlobal attention patterns. fig. 3 and fig. 4 are described in the Results section above. Figures are not reproduced here.

Paper vs code

The official code at code_commit is google-research/longt5, Apache-2.0. It contains the tasks, the PEGASUS data preprocessing, and the NQ and MediaSum preprocessing scripts. It does not contain the model architecture or training configurations. The paper's footnote points the architecture to Flaxformer, and the README says the architecture and training setups are in Flaxformer README.md. The Flaxformer longt5 directory on main holds the architecture (longt5_architecture.py, long_attention.py, relative_position_biases_general.py) and tests. It holds no training configs. The code cards below cite Flaxformer by file, and not by line, because the Flaxformer files were read through WebFetch, not from a local clone, and I did not verify their line numbers. Flaxformer is not pinned to a commit.

• Model size. Paper: base ~220M, large ~770M §4.1. Released checkpoints are 250M, 780M and 3B parameters README.md.
• Pretraining objective. Paper: "Principle Ind-Uniq" from PEGASUS §3.2. Code: parser_strategy='dynamic_rouge', ROUGE-1 F-score, parser_rouge_ngrams_size=1 and an empty stopword file preprocessors.py. The parser itself is imported from an external pegasus package [code: longt5/preprocessors.py#L31], which is not in this repo, so Ind and Uniq cannot be confirmed here.
• Undocumented sampling knobs. The code also sets parser_mask_sentence_option_prob=[.9, 0, .1, 0], parser_rouge_noise_ratio=.2 and parser_dynamic_mask_min_ratio=.33 preprocessors.py. The paper states none of these. Their meaning depends on the external parser.
• Which pretraining task produced the released checkpoints. The repo defines a PEGASUS parser task and a mixture of PEGASUS with span corruption (mean noise 3) at rate 1.0 tasks.py. The paper says the PSG-only setup is the main one and that SC+PSG is worse §5.3. The released pretraining mixture is not stated.
• Length and ratio. Input 4096, target 910, masked-sentence ratio 0.2 match the paper preprocessors.py.
• Pretraining data. The repo uses c4/en:2.2.0 for both pretraining tasks tasks.py, consistent with §4.1.1.
• NQ split. Paper: 90/10 of the official training set §4.3.1. Code: shards 0–8 train and shard 9 validates, with the official dev set as test [code: longt5/tasks.py#L114-L123]. This is about 90/10 by shard count, not an exact split. The paper does not give exact counts.
• TriviaQA split. Paper: 90/10 §4.3.1. Code: train[:78785] and train[78785:], with a separate task that trains on train+validation for the leaderboard tasks.py.
• Table values in the README. The README gives LongT5 large MediaSum R-1 as 35.53, where the paper says 35.54 README.md. It gives LED arXiv R-1 as 46.64, where the paper says 46.63 [code: README.md#L14; tab. 2]. These are rounding-level differences.
• TGlobal global key/value projections. The Flaxformer long_attention.py comment says global keys and values reuse the local key and value projections. The paper does not say this §3.1.2. This came from a WebFetch read, so verify before relying on it.
• TGlobal normalization. The Flaxformer code applies a T5 layer norm to the summed global tokens, consistent with "summing (and then normalizing)" in §3.1.2.
• Fine-tuning and decoding. The learning rate, dropout, batch size, input lengths, and greedy decoding §4.1, §4.1.2 are not in the repo. Flaxformer's longt5 directory has no configs. The T5X or other config locations were not checked, so these remain as stated in the paper.

No discrepancies were found for the TGlobal block size (k = 16 in the paper and "typical 16 or 32" in Flaxformer) or the local radius (127 in both).

Pitfalls

• Single runs. Each configuration was run once, and the paper flags noise at single points such as 6k for TGlobal large §4.3.2, §5.2. Compare trends, not individual points.
• Suspicious table entry. T5.1.1 base at 6k on NQ has EM = F1 = 56.73 in tab. 4 and in app. B, tab. 7. Treat it as possibly a transcription error in the paper, and do not use it as a baseline without checking the source.
• Unequal input lengths. tab. 2 compares models at different input lengths for each dataset (4k, 8k, 16k). The caption says this explicitly. Gains may come from input length as much as from the architecture.
• Greedy decoding. All results use greedy decoding. Beam search may improve them, and the paper says so §4.1.
• NQ is not a leaderboard comparison. The NQ setup predicts answer text rather than spans, uses a 90/10 split of the training data, and reports on the official dev set. The authors say the numbers are not comparable to the leaderboard §4.3.1.
• Asymmetric baselines. The T5.1.1 baselines are run at 512 tokens and at the longest length that fits on the same hardware §4.3.2. The memory-limited length depends on the TPU topology (4x8 vs 8x16 with 8 partitions) §4.3.2, fn..
• Cost of TGlobal. Each token attends to l/k extra global tokens in addition to its local window. At 16k input and k = 16, that is 1,024 globals per token, per head, per layer. This is derived from the formula in §3.1.2, not stated in the paper. Local attention has no such term, which is why it reaches 36k on base but underperforms on NQ §5.1, §5.2.
• Packing. Globals must be masked by example. Without that mask, examples packed into one sequence attend to each other §3.1.2.
• Relative bias to globals. The Flaxformer side bias from tokens to global blocks is required, and the code raises if it is missing (WebFetch read; not line-verified).
• Local window size. The paper defines the window as r tokens on each side §3.1.1, which gives 2r + 1 = 255 positions including the token itself. A Flaxformer docstring, as read through WebFetch, warns that a naive block implementation can reach up to 2r + 1 tokens away on each side. Check the mask before relying on the exact window.
• Pretraining comparisons. The PSG versus SC ablation uses T5.1.1 pretrained at 512 input, while LongT5 is pretrained at 4096 §5.3, tab. 6. Differences mix objective and input length.
• Scoring details. The PEGASUS sentence scorer uses ROUGE-1 F-score with an empty stopword file preprocessors.py. The paper does not state stopword handling.
• Output length. The decoder is full attention, and only the encoder is sparse. Long outputs are not addressed §3.1, §6.

Lineage

• Builds on: T5 (arXiv:1910.10683) for the backbone and text-to-text format §2; ETC (EMNLP 2020; arXiv:2004.08483 from the repo README) for the local-global attention idea, with TGlobal removing its side inputs §3.1.2; PEGASUS (arXiv:1912.08777) for principle-sentence pretraining §3.2; BigBird (arXiv:2007.14062) for sparse attention patterns §3.1.
• Compared against: BigBird-PEGASUS (arXiv:2007.14062), HAT-BART (arXiv:2104.07545), PRIMER (arXiv:2110.08499), LED / Longformer (arXiv:2004.05150), TG-MultiSum (arXiv:2110.11207), DANCER PEGASUS, BART, Fusion-in-Decoder (arXiv:2007.01282), ReadTwice (arXiv:2105.04241) [§4.2.2, §4.3.2].
• Later work by the same authors: mLongT5 (arXiv:2305.11129), named in the repo README README.md.