PaperlayerView exact revision on arXiv

Source-linked AI summary

ETC: Encoding Long and Structured Inputs in Transformers

Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, Li Yang

arXiv:2004.08483v5cs.LGstat.ML

TL;DR

Standard Transformers face quadratic attention costs that limit input length, while few models focus on graph or hierarchical structure among input tokens. ETC combines global-local attention with relative position encodings and CPC pre-training, and sets new state of the art on four datasets requiring long and/or structured inputs.

  • Problem

    Standard Transformer attention scales quadratically and typically limits inputs to around 512 tokens, while few models focus on graph or hierarchical structure among input tokens.

    Source: §1 · p. 1

  • Method

    ETC combines global-local attention over global and long input sequences with relative position encodings and CPC pre-training.

    Source: §1 · pp. 1–2

  • Results

    ETC set new state of the art on four datasets requiring long and/or structured inputs.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The paper reports significant gains from increased input length and further model-quality improvement from representing dataset structure.

    Source: §5 · p. 9

    Experiments were limited to 8192 tokens, although gradient checkpointing allows ETC to scale beyond that.

    Source: §4.1 · p. 7

Abstract

from arXiv · show

Transformer models have advanced the state of the art in many Natural Language Processing (NLP) tasks. In this paper, we present a new Transformer architecture, Extended Transformer Construction (ETC), that addresses two key challenges of standard Transformer architectures, namely scaling input length and encoding structured inputs. To scale attention to longer inputs, we introduce a novel global-local attention mechanism between global tokens and regular input tokens. We also show that combining global-local attention with relative position encodings and a Contrastive Predictive Coding (CPC) pre-training objective allows ETC to encode structured inputs. We achieve state-of-the-art results on four natural language datasets requiring long and/or structured inputs.

1 Introduction

ETC targets two Transformer limitations: quadratic attention constrains input length, and few models represent graph or hierarchical structure among input tokens. Its global-local attention, relative position encodings, and CPC pre-training address these challenges; experiments set new state of the art on four datasets with long and/or structured inputs.

1 Introduction

  • Standard Transformer attention has quadratic computational and memory complexity, typically limiting inputs to around 512 tokens.That limit is problematic for tasks such as question answering across multiple documents that must fit in the model input.Source: §1 · p. 1
  • Few models focus on structured inputs, including graph or hierarchical relations among tokens such as document structure and entity links.ETC aims to represent these structures, including sentence and paragraph organization and web-page DOM trees.Source: §1 · p. 1
  • ETC's global-local attention divides input into global and long sequences to reduce quadratic attention scaling; relative position encodings help handle structure.The architecture can also be initialized from existing pre-trained BERT models.Source: §1 · p. 1
  • CPC pre-training improves performance further on tasks where structure matters, serving as an MLM-like task at sentence-level granularity.The paper focuses on ETC; a related BigBird work explores random sparse attention patterns and extensions to a decoder.Source: §1 · p. 2
  • ETC set new state of the art on Natural Questions, HotpotQA, WikiHop, and OpenKP, datasets requiring long and/or structured inputs.The paper also notes that RoBERTa initialization significantly improves performance and CPC improves it further on tasks where structure is important.Source: §1 · p. 2

2 Background

Prior work scales Transformer attention through sparse attention, recurrence, hierarchical mechanisms, and compressed attention. These approaches restrict, segment, summarize, or compress input representations in different ways.

2 Background

  • Sparse attention restricts each token to a subset of other tokens; examples include predefined patterns, learned spans, hashed neighbors, and online clustering.The cited approaches report costs including O(n√n), O(n log(n)), and O(n^1.5) for particular methods.Source: §2 · p. 2
  • Recurrence extends attention across segments: Transformer-XL processes segments sequentially and attends to the current and previous segments at each layer.At layer k, the representation is influenced by the current segment and the k −1 previous segments.Source: §2 · p. 2
  • Hierarchical mechanisms process input blocks independently into summary embeddings, then apply separate layers to the concatenated summaries.HIBERT applies this idea at the sentence level for extractive summarization.Source: §2 · pp. 2–3
  • Compressed attention selectively summarizes input: BP-Transformer uses tree nodes for distant tokens, while other methods compress token groups or use an auxiliary summary token.The approaches vary in how they replace distant raw-token attention with summarized representations.Source: §2 · p. 3

3 Extended Transformer Construction

ETC modifies the Transformer encoder with relative position encoding, global-local attention, and CPC pre-training to handle long and structured inputs. Its design also supports hierarchical relations and initialization from existing BERT-family models.

3 Extended Transformer Construction

  • ETC retains the Transformer encoder while adding relative position encoding, global-local attention, and CPC pre-training; the paper leaves decoder work for the future.Source: §3 · p. 3

3.1 Relative Position Encoding

  • Relative position labels encode token offsets, can be extended to arbitrary pairwise relations, and are independent of input length.Offsets are clipped at a maximum distance k.Source: §3.1 · pp. 3–4

3.2 Global-Local Attention

  • Restricting long-to-long attention to radius r reduces attention complexity, while unrestricted global pathways let long-input tokens exchange information through global tokens.With ng = O(2r + 1), attention is linear in the long-input size.Source: §3.2 · p. 4

3.3 Long Inputs and Global-Local Attention

  • For long inputs, ETC can assign a global auxiliary token to each segment and use relative labels to link segment tokens with their word pieces.Hard masking in the global-to-long direction can improve performance on some datasets; global tokens provide an indirect route beyond the local neighborhood.Source: §3.3 · p. 5

3.4 Structured Inputs

  • ETC represents structure through relation labels, hierarchical global summary tokens, and attention masks that exclude selected token pairs.This supports structures such as contexts containing ordered sentences without an order between contexts.Source: §3.4 · p. 5

3.5 Pre-training Tasks

  • ETC combines whole-word masked language modeling with CPC, which trains global sentence-summary tokens to predict masked sentences’ hidden representations.The CPC objective uses a Noise Contrastive Estimation loss.Source: §3.5 · pp. 5–6

3.6 Lifting Weights from Existing Models

  • Initializing ETC from RoBERTa improves performance significantly over pre-training from scratch, although further pre-training is still required.BERT parameters can be copied into ETC’s corresponding parameter matrices.Source: §3.6 · p. 6

4 Empirical Evaluation

ETC was evaluated on four benchmarks with long or structured inputs, with experiments covering model configuration, dev-set comparisons, and official leaderboard performance. The paper reports new state-of-the-art results on several leaderboard measures, including results achieved by single models over leaderboard ensembles.

4 Empirical Evaluation

  • The four benchmarks cover full-article NQ, multi-context HotpotQA, cross-context WikiHop, and OpenKP keyphrase extraction from websites with hierarchical and spatial relations.The evaluation was chosen to include long inputs or interesting input structure.Source: §4 · p. 6

4.1 Training Configuration

  • Pre-training puts wordpieces in the long input and one auxiliary token per sentence in the global input; fine-tuning adds global summaries for paragraphs, sentences, or OpenKP VDOM blocks.Experiments were limited to 8192 tokens.Source: §4.1 · p. 7

4.2 Results on the Dev Set

  • RoBERTa-initialized ETC reached 0.782 / 0.585 on NQ, exceeding RikiNet's 0.754 / 0.593 on long answers but lagging on short answers.Longer inputs improved performance, and removing CPC or sharing attention projections significantly hurt NQ results.Source: §4.2 · pp. 7–8
  • Both ETC base and large models outperform their corresponding Longformer models on HotpotQA and WikiHop.Removing CPC and hard g2l masking hurt HotpotQA; flat structure particularly hurt WikiHop, while sharing Q/K/V seemed to help there.Source: §4.2 · p. 8, Table 3 · p. 8
  • On OpenKP, constrained ETC scored 0.399; scaling to 4096 tokens gave a slight boost, hard g2l helped, and visual features brought the largest benefit.Large models, especially RoBERTa-initialized models, improved results; maximum-logit scoring across occurrences improved on first-occurrence scoring.Source: Table 4 · p. 8, §4.2 · p. 9

4.3 Official Leaderboard Results

  • Single-model submissions set new state of the art on WikiHop, OpenKP, NQ long answer, and HotpotQA Support F1, outperforming leaderboard ensembles.The leaderboard submissions used models selected by their best dev-set results.Source: Table 5 · p. 8, §4.3 · p. 9

5 Conclusions

ETC combines global-local attention, relative position encodings, and CPC pre-training to scale input length and represent structure. The authors report gains from longer inputs and structure, while noting that their experiments use limited structure and single-model comparisons.

5 Conclusions

  • ETC scales input length linearly and combines global-local attention with relative position encodings and CPC to encode structured inputs.The paper reports that lifting weights from existing BERT models significantly improves results.Source: §5 · p. 9
  • ETC gains from increased input length, and representing dataset structure further improves model quality.The experiments also show ETC exploiting existing structure despite the datasets containing less structure than graph datasets.Source: §5 · p. 9
  • Longformer-large reported better results using a two-stage approach; the paper's table presents single-model results for comparison.The authors identify the Longformer-large results as 81.0 / 85.8.Source: §5 · p. 9
  • Future work includes complementary attention mechanisms, greater scalability, and evaluation on datasets with richer structure.Source: §5 · p. 9

Global-Local Attention Implementation

The implementation projects queries, keys, and values, then uses padded blocks to compute sliding-window local attention efficiently. Blocking saves memory at the cost of extra compute, with masking removing out-of-radius attention.

Global-Local Attention Implementation

  • The implementation extends sliding-window local attention with flexible masking, relative position encoding, and global tokens as side keys and values.It projects each embedding into queries, keys, and values for each attention head.Source: Global-Local Attention Implementation · p. 12
  • For local radius r, inputs are divided into blocks of length r + 1, with padding blocks added at both ends.For r = 2, the example uses blocks of length 3.Source: Global-Local Attention Implementation · p. 12
  • Each query block attends to its own key block and the adjacent blocks, avoiding materialization of a separate key window for every query.Keys and values use matching blocks; scaled dot-product attention is applied independently for each row.Source: Global-Local Attention Implementation · p. 12
  • Blocking can expose queries to tokens beyond radius r, which are masked out; this saves memory at the expense of extra compute.Source: Global-Local Attention Implementation · p. 12

Efficient Relative Attention Implementation

The appendix describes a relative-attention implementation that avoids materializing a pairwise key-vector tensor, then details ETC's CPC objective. CPC uses two shared-weight encoders to align representations of masked segments, with within-batch negatives and an NCE loss.

Efficient Relative Attention Implementation

  • The relative-attention implementation dots each query with unique relative-key vectors, then gathers scalar scores for query-key pairs instead of materializing pairwise vectors.This avoids O(n^2dz) memory overhead and permits different relative-key vectors per head without additional activation memory.Source: Efficient Relative Attention… · p. 12

Appendix B: CPC Loss in ETC

  • ETC adapts CPC as a dual-encoder task using two instances of the same model with shared weights.One encoder is the main model; its long input is divided into segments, each paired with a global token.Source: Appendix B · p. 12
  • Masked segments are encoded independently, and the encoders train to align their corresponding global-token representations using within-batch negatives and an NCE loss.Source: Appendix B · p. 13

Appendix C: Training Details

Training uses Wikipedia and Books pre-training, with CPC and MLM masking when both objectives are enabled. Fine-tuning and prediction procedures vary by task; the appendix details NQ data preparation, model selection, and answer aggregation.

Appendix C: Training Details

  • Default pre-training uses BERT's Wikipedia and Books corpora, filters documents under seven sentences, and runs 33 epochs with batches of 512 sequences of 4096 tokens.ETC-large runs for 66 epochs with batch size 1024; concatenating shorter documents with masking yields roughly 3x speedup for 4096-token pre-training.Source: Appendix C · p. 13
  • When pre-training with CPC, 10% of sentences are selected for CPC masking, then 15% of remaining tokens are masked for MLM.When both objectives are used, the combined loss weights MLM at 0.8 and CPC at 0.2.Source: Appendix C · p. 13
  • For NQ, the long input starts with CLS and the question, followed by the document; global tokens mark the question and each paragraph.Sentences are word-tokenized before BERT/RoBERTa tokenization to align word-indexed ground truth.Source: Appendix C · p. 13
  • NQ fine-tuning sweeps learning rates and epochs, predicts answer boundaries and type, and uses sliding windows for inputs longer than the configured length.The sweep tests learning rates 3 × 10^-5 and 5 × 10^-5, with three or five epochs, or two or three for large models.Source: Appendix C · p. 13
  • Model selection tests one random seed per hyperparameter configuration and chooses the model with the highest average dev F1 at fine-tuning's end.The best large model initialized from RoBERTa was trained for two epochs at learning rate 3 × 10^-5.Source: Appendix C · p. 13
  • NQ inference selects the best start position before choosing a later end position, limits short answers to 38 words, and replaces short answers with yes/no when those logits are higher.Source: Appendix C · p. 13

HotpotQA

HotpotQA and WikiHop use ETC’s long and global inputs to represent context structure and task-specific information. The sections describe dataset preparation, model selection, and prediction procedures.

HotpotQA

  • HotpotQA’s long input concatenates the question and contexts, while global tokens represent the question, each context, and its sentences; inputs are cropped to 4096 tokens without windowing.The global input length is 256, and 16 instances without extractable answers are filtered from training.Source: HotpotQA · p. 14
  • HotpotQA selects models by dev joint F1 after a hyperparameter sweep and three seeds; the best large RoBERTa-initialized model used five epochs and a 3 × 10−5 learning rate.The selected model used batch size 32 and supporting-fact threshold 0.4.Source: HotpotQA · p. 14
  • HotpotQA predicts supporting facts and answer type from global embeddings, then selects a constrained answer span from long-input embeddings.Eligible spans must lie within a sentence or title, belong to a supporting fact, and be no longer than 30 tokens.Source: HotpotQA · p. 14

WikiHop

  • WikiHop represents candidate answers as global tokens and links them to text mentions by string matching; long inputs are cropped to 4096 tokens without windowing.When cropping contexts, a MaxHeap prioritizes retaining contexts with fewer sentences.Source: WikiHop · p. 14
  • WikiHop selects the checkpoint with highest dev accuracy and predicts by choosing the candidate with the highest logit from global embeddings.The selected large RoBERTa-initialized model used ten epochs and a 5 × 10−5 learning rate.Source: WikiHop · p. 14

OpenKP

OpenKP combines document text, VDOM-node global tokens, and visual features for keyphrase prediction. The appendix also describes how ETC reuses compatible BERT/RoBERTa weights and illustrates structure encoding.

OpenKP

  • OpenKP adds one global token per VDOM node and embeds visual features alongside text, using 4096-token long and 512-token global inputs by default.Instances are truncated according to whichever input limit is more restrictive; the processed set contains 133,374 training examples.Source: OpenKP · p. 15
  • OpenKP predicts keyphrases from word and n-gram embeddings; its improved max loss takes the highest logit across occurrences of each keyphrase rather than using only the first occurrence.The selected large model used max loss and RoBERTa initialization; inference returns up to five deduplicated keyphrases.Source: OpenKP · p. 15

Appendix D: Lifting Weights from BERT/RoBERTa

  • ETC can reuse compatible BERT/RoBERTa feed-forward, attention, layer-normalization, and token-embedding weights, while initializing CPC and relative-position components randomly.Lifting requires matching layer count, hidden size, attention-head count, and feed-forward intermediate size.Source: Appendix D · p. 15
  • Figure 5 illustrates relative-position labels on attention arrows and where the corresponding tokens appear in ETC’s input.Figure 4 plots wall time per step against sequence length; ETC’s global length is set to one-sixteenth of long length until capped at 512 at long length 8192.Source: Figure 5 · p. 16, Figure 4 · p. 16

Appendix E: Model Computational Requirements

The appendix reports ETC’s computational scaling, memory headroom, and parameter-count properties, then illustrates how global tokens represent hierarchy in structured text.

Appendix E: Model Computational Requirements

  • With global input fixed at 512 tokens on a TPU v3 core, ETC fit long inputs of 22,656 tokens for base models and 8,448 for large models.These experiments used gradient checkpointing and removed optimizer gradient moments.Source: Appendix E · p. 16
  • Beyond about 1500 input tokens, ETC becomes faster per step than BERT, whose plotted run ends earlier because of memory constraints.ETC is initially slower, and its wall time is not linear in the figure because global input length also increases.Source: Appendix E · p. 16
  • ETC attention complexity is linear in long-input size when global length scales with relative-position range; parameter count does not depend on input length.Separate global and long projections yield about 50% more parameters than BERT, while shared projections yield a similar count.Source: Appendix E · p. 16

Appendix F: Structured Input Example

  • For WikiHop-style inputs, ETC uses global context and sentence tokens to encode hierarchy, with relative-position labels marking token, sentence, and context relations.The example keeps sentence order within contexts while allowing contexts themselves to have no order.Source: Appendix F · p. 17
Loading 2004.08483v5…