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

LongT5 is an extension of the T5 model that handles long sequence inputs more efficiently. We integrated attention ideas from long-input transformers ETC,and adopted pre-training strategies from summarization pre-training PEGASUS into the scalable T5 architecture. The result is a new attention mechanism we call Transient Global(TGlobal), which mimics ETC’s local/globalattention mechanism, but without requiring additional side-inputs. We are able to achieve state-of-the-art results on several summarization and question answering tasks, as well as outperform the original T5 models on these tasks.

Summarization Results

LongT5 achieves state-of-the-art performance on several summarization benchmarks that required longer context or multi-document understanding. The table is showing ROUGE-1 scores. LongT5 base models are all reported with 4k input tokens; large and xl models are trained with 16k tokens for arXiv, PubMed, BigPatent, 8k for MultiNews, and 4k for MediaSum and CNN/Daily News.

QA Results

Natural Questions

For NQ, we compare T5.1.1 and LongT5 with TGlobal attention. We decided to run T5.1.1 (1) with the default 512 input sequence length and (2) with the largest input sequence length that can fit into device memory, and use those as baselines. Since we are comparing against T5.1.1, for LongT5 experiments we report results at 512 input length for base and large, and the largest input length allowed by each model before running out of memory on the same hardware configuration used in our T5.1.1 experiments. For base and large models, we used 4x8 TPUv3 and no model partitioning; for xl model, we used 8x16 TPUv3 and 8 partitions.

Moreover, in our analysis for Input Length vs Speed and Input Length vs Performance sections using NQ, it shows that (1) at shorter sequence length T5.1.1 and LongT5 variants have similar speeds, but as we increase the sequence length, LongT5 becomes significantly faster, (2) T5.1.1 models reach their out-of-memory point much earlier than LongT5 models, and (3) performance increases significantly as input length increases.

TriviaQA

For TriviaQA, we compare LongT5 with various top approaches on the leader board. All LongT5 models are reported with 16k input tokens.

Usage

Data Preprocessing

Most of our tasks are using Tensorflow Datasets which works directly with the SeqIO used in the T5 library. But for Natural Questions and MediaSum we provided our own data preprocessing code. To run the tasks corresponding to these datasets, please specify NQ_DATA_DIR and MEDIASUM_DATA_DIR to the output files produced by the preprocessing code in tasks.py.

Example command for running NQ data preprocessing:

# Data path where the NQ json files are downloaded to.
INPUT_PATH="..."
# Data path where the output files will be generated.
OUTPUT_PATH="..."
LONGT5_DIR="..."  # directory where the LongT5 repo is cloned.

python3 ${LONGT5_DIR}/data/nq_preprocess.py \
  --input_path=${INUT_PATH} \
  --output_path=${OUTPUT_PATH}

Training

The experiments are shown in the tasks.py file. Our architecture, model, and training configuration setups can be found in Flaxformer github repository.

Released Model Checkpoints

We have released the following checkpoints for LongT5 pre-trained models:

• LongT5-Local-Base (250 million parameters): gs://t5-data/pretrained_models/t5x/longt5/local_base
• LongT5-TGlobal-Base (250 million parameters): gs://t5-data/pretrained_models/t5x/longt5/tglobal_base
• LongT5-Local-Large (780 million parameters): gs://t5-data/pretrained_models/t5x/longt5/local_large
• LongT5-TGlobal-Large (780 million parameters): gs://t5-data/pretrained_models/t5x/longt5/tglobal_large
• LongT5-TGlobal-XL (3 billion parameters): gs://t5-data/pretrained_models/t5x/longt5/tglobal_xl

Additionally, we have released the following checkpoints for mLongT5 pre-trained models:

• mLongT5-TGlobal-Base (250 million parameters): gs://t5-data/pretrained_models/t5x/mlongt5/tglobal_base
• mLongT5-TGlobal-Large (780 million parameters): gs://t5-data/pretrained_models/t5x/mlongt5/tglobal_large
• mLongT5-TGlobal-XL (3 billion parameters): gs://t5-data/pretrained_models/t5x/mlongt5/tglobal_xl

Citing LongT5

If you use LongT5 in your research, please cite LongT5: Efficient Text-To-Text Transformer for Long Sequences.

@inproceedings{guo2022longt5,
    title = "{L}ong{T}5: {E}fficient Text-To-Text Transformer for Long Sequences",
    author = "Mandy Guo and Joshua Ainslie and David Uthus and Santiago Onta{\~n}{\'o}n and Jianmo Ni and Yun-Hsuan Sung and Yinfei Yang",
    booktitle = "Findings of the Association for Computational Linguistics: NAACL 2022",
    year = "2022",
    url = "https://aclanthology.org/2022.findings-naacl.55",
    pages = "724--736",
}

For mLongT5, please cite mLongT5: A Multilingual and Efficient Text-To-Text Transformer for Longer Sequences.

@misc{uthus2023mlongt5,
    title = "{mLongT5}: A Multilingual and Efficient Text-To-Text Transformer for Longer Sequences",
    author = "David Uthus and Santiago Onta{\~n}{\'o}n and Joshua Ainslie and Mandy Guo",
    year = "2023",
    eprint = "2305.11129",
    archivePrefix = "arXiv",
    primaryClass = "cs.CL",
    url = "https://arxiv.org/abs/2305.11129"
}