Many natural language processing tasks benefit from long inputs, but processing long documents with Transformers is expensive -- not only due to quadratic attention complexity but also from applying feedforward and projection layers to every token. However, not all tokens are equally important, especially for longer documents. We propose , a long-input Transformer model that builds on this intuition by employing conditional computation, devoting more resources to important tokens in both feedforward and attention layers. We show that achieves stronger performance than with much faster training and inference, achieving SOTA on the long-input SCROLLS benchmark. Moreover, can effectively and tractably make use of extremely long inputs, showing strong gains up to 64k input length.

§ INTRODUCTION

Many natural language processing tasks, such as summarization <cit.> or question answering over long documents <cit.>, require machine learning models to encode long-form text. Processing long documents with a Transformer model is computationally expensive, both because attention cost scales quadratically with input length and because feedforward and projection layers have to be applied to each input token.

Over the past few years, many “efficient Transformer” approaches have been proposed that reduce the cost of the attention mechanism over long inputs <cit.>. However, especially for larger models, the feedforward and projection layers actually make up the majority of the computational burden and can render processing long inputs intractable.

This paper presents (Conditional LongT5), a new family of models that, building on top of  <cit.>, enables fast processing of long inputs by combining architecture improvements for both attention and feedforward layers. is based on the intuition that some tokens are more important than others, and we can achieve better quality for lower cost by devoting more computation to important tokens. Moreover, the fraction of important tokens is likely to diminish with document length, allowing for tractable processing of long documents.

An overview of a Transformer layer with conditional computation. All tokens are processed by light attention and MLP layers, while $q$ routed query tokens perform heavier attention over $v$ routed key-value tokens and $m$ routed tokens are processed by a heavier MLP.

scale only axis,
xlabel=Time per sample (ms),

ylabel=Average perf,

x tick label style=log ticks with fixed point,
grid style=dashed,

legend pos=north west,
ylabel style=align=left, text height=0.2cm,

empty legendInference
[color=longtcolor,mark=,mark size=2pt, line width=2] table

136 43.10
329 45.25
870 46.55
[color=sliccolor,mark=,mark size=2pt, line width=2] table

90 42.36
201 45.31
439 47.42
at (axis cs:132.00,43.20) [anchor= north west, color=darkgrey] B;
at (axis cs:340.0,45.15) [anchor= north, color=darkgrey] L;
at (axis cs:870.0,46.45) [anchor= north, color=darkgrey] XL;

at (axis cs:90.0,42.50) [anchor= north west, color=darkgrey] B;
at (axis cs:215,45.20) [anchor= north, color=darkgrey] L;
at (axis cs:444,47.25) [anchor= north, color=darkgrey] XL;

scale only axis,
xlabel=Time per sample (ms),

x tick label style=log ticks with fixed point,
grid style=dashed,

legend pos=north west,
ylabel style=align=left, text height=0.2cm,

empty legendFine-tuning
[color=longtcolor,mark=,mark size=2pt, line width=2] table

270 43.10
790 45.25
2630 46.55

[color=sliccolor,mark=,mark size=2pt, line width=2] table

154 42.36
496 45.31
1917 47.42
at (axis cs:270.00,43.20) [anchor= north west, color=darkgrey] B;
at (axis cs:825.0,45.15) [anchor= north, color=darkgrey] L;
at (axis cs:2630.0,46.45) [anchor= north, color=darkgrey] XL;

at (axis cs:154.0,42.50) [anchor= north west, color=darkgrey] B;
at (axis cs:525,45.20) [anchor= north, color=darkgrey] L;
at (axis cs:1917,47.25) [anchor= north, color=darkgrey] XL;

legend columns=2,
legend style=
column sep=4ex,
legend entries=,
mark=,mark size=3pt,solid,color=longtcolor

achieves stronger performance than at any speed. Average performance on all datasets as a function of inference and fine-tuning time per sample (ms) for and Base, Large, and XL models. does not use MQA, but we report speed as though it had for a conservative baseline.

In particular, divides each feedforward layer and each attention layer into a light branch which is applied to all tokens and a heavy branch which is applied to a set of important tokens, selected specifically for that input and component. The light feedforward branch has lower hidden dimension than standard while the heavy feedforward branch has higher hidden dimension. The light attention branch has fewer heads and applies only local attention, while the heavy attention branch performs full attention over another separately selected set of important tokens. Figure <ref> provides an overview of the conditional mechanism.

Finally, also includes two other modifications to the architecture. adds multi-query cross-attention <cit.>, significantly speeding up inference. also employs the UL2 <cit.> pre-training objective, which we demonstrate allows for in-context learning over long inputs.

We show that performs much faster fine-tuning and inference with similar or better model quality, improving over on arXiv summarization <cit.> and TriviaQA question answering <cit.> datasets and achieving SOTA on the SCROLLS benchmark <cit.>. Moreover, achieves further gains in quality and speed for tasks with extremely long inputs (64k tokens), with less-than-linear scaling of “focus” tokens.

§ BACKGROUND

Transformer FLOPs follows an extensive line of work in attempting to reduce the computational cost of Transformer models, particularly over long inputs. The computational burden of Transformer models has several distinct elements, and different approaches focus on reducing the cost of different components. For that reason, it is helpful to start by providing a breakdown of the computational cost of Transformer components. Table <ref> shows the FLOPs[Each multiply-add is counted as a single FLOP.] for each component of a Transformer encoder layer <cit.>.

Encoder Layer Component     Flops

Vanilla self-attention computation    $2 n^2 d$

Attention QKV and output projections     $4 n d^2$

Feedforward layer    $8 n d^2$

local attention computation     $2 nwd$

global attention computation     $\frac{n^2 }{8}d$

Computational cost of encoder layer transformer components measured in FLOPs. $n$ is the input length, $d$ is the model dimensionality, and $w$ is the size of the local attention window.

Sparse attention The first challenge of applying a Transformer to a long input is that the FLOPs of the self-attention mechanism scales quadratically in the input length, becoming intractable for long inputs. A large body of work focuses on reducing self-attention cost, restricting attention between a subset of inputs <cit.> or to a subset of layers <cit.>. In  <cit.>, the most closely related model to , tokens attend within a local window as well as to a mean-pooled summary representation for each block of 16 tokens in the input. attention leads to sharply reduced (though still non-negligible) FLOPs (Table <ref>).

Conditional computation After applying a sparse attention mechanism, the feedforward and attention projection layers account for the majority of the FLOPs. These costs scale with the length of the input, such that processing long inputs is still prohibitively expensive. A common approach to reduce the remaining cost is to employ some form of conditional computation, avoiding applying all model parameters to the entire input. CALM <cit.> applies a varying number of decoder layers to each decoded token, outputting a token early if the model is confident in its prediction. Mixture-of-Experts models <cit.> route inputs through a small proportion of expert sub-modules, bringing to bear only the parameters most relevant to the input. In the context of retrieval-augmented models, numerous works re-rank retrieved passages by their relevance to the query and process only the highest scoring passages <cit.> and vary the number of processed passages depending on model confidence <cit.>. Concurrent work CoDA <cit.> employs a related conditional computation mechanism, designed for efficient adaptation rather than modeling long documents.

Device utilization FLOPs do not tell the whole story, as modeling choices can influence the effective speed of operations achieved by accelerators. For long text inputs, autoregressive decoder inference is very slow due to memory bandwidth constraints from repeatedly loading the long sequence of keys and values <cit.>. <cit.> introduces multi-query attention (MQA), sharing heads for keys and values to reduce memory bandwidth overhead. <cit.> studies how to shard large models, especially in the context of MQA, to obtain optimal device utilization and therefore speed.

Training objectives T5 introduced the span corruption objective <cit.>, a modification of masked language modeling <cit.>. made use of the PEGASUS <cit.> sentence reconstruction objective for improved summarization performance. <cit.> proposes UL2, a mixture of span corruption, prefix, and causal language modeling, and shows that it leads to strong performance on both short-output and generative tasks.

§.§ Conditional computation

As discussed in the previous section, a large proportion of Transformer FLOPs arise from feedforward and projection layers that scale with the length of the input sequence. Therefore, training and inference on long documents remains expensive.

further reduces the cost of processing long documents through conditional computation, following the intuition that some tokens are more important and therefore benefit more than others from heavy computation. First, some types of tokens may inherently require less computation, such as filler words and punctuation. Second, especially in long documents, large parts of the input may not be relevant to the current question, task, or processing stage.

The conditional computation mechanism consists of three components: routing modules, conditional feedforward layers, and conditional attention layers. All tokens are processed by standard, lightweight attention and feedforward layers. Routing modules additionally select important tokens from an input at each attention or feedforward layer, and a heavy conditional layer applies additional computation to routed tokens. This section describes each component in detail. Figure <ref> provides an overview of the conditional computation mechanism, and Table <ref> compares and FLOPs.

Model     Encoder Layer Flops

T5    $12 n d^2 + 2n^2 d$

$12n d^2 + \frac{n^2}{8} d$

$7 \frac{1}{4} n d^2 + \frac{n^2}{84} d$

uses significantly fewer FLOPs than . Comparison of approximate encoder layer total FLOPs between T5, , and . FLOPs rounded to readable fractions.

In order to separately select important tokens for each component in each layer, we need a learnable and tractable routing function. We follow the simple three-step mechanism from <cit.>: (1) multiply inputs with a learned embedding to obtain routing scores, (2) normalize, and (3) select the top-$k$ highest scoring inputs.

Let $X_i$ be the representation of token $i$, and $u$ a $d$-dimensional learnable embedding. Then the routing score of token $i$ is

\begin{equation*}
s_i = X_i \cdot u
\end{equation*}

We select the top-$k$ highest scoring inputs. In order to provide a learning signal to the scoring embedding, we make sure the contribution of the routed tokens to the layer update is scaled according to the routing score, as will be seen later. To provide a better distributed signal to all tokens, we also globally normalize the routing scores to sum up to the number of desired routed tokens using a generalized softmax, resulting in normalized scores $\tilde{s}_i$.
Each layer has three independent routers, one each for the feedforward layer, attention queries, and attention key-values.

Conditional Feedforward
Intuitively, some token representations may benefit from more processing than others. The conditional feedforward layer applies an additional high-capacity feedforward layer to selected tokens. In particular, let $X_i$ be the model state of the $i$th token and $\tilde{s}_i$ denote the normalized routing score (set to 0 for non-routed tokens). Then the feedforward update for is given by

\begin{equation*}
X_i = X_i + \textrm{FFd}_{\textrm{Light}}(X_i) + \tilde{s}_i \cdot \textrm{FFd}_{\textrm{Heavy}}(X_i)
\end{equation*}

The light and heavy feedforward branches differ only in their hidden dimension, with the light branch having smaller hidden dimension than the standard T5 feedforward layer and the heavy branch larger. Let $n$ denote the number of input tokens, $m$ the number of selected tokens, and $r_L$ and $r_H$ the ratios of light and heavy hidden dimension to standard T5 hidden dimension. Then the FLOPs of the layer are given by

\begin{equation*}
\text{FLOPs}_{\text{FFd}} = \underbrace{8 n r_L d^2}_{\text{Light branch}} + \underbrace{8 m r_H d^2}_{\text{Heavy branch}}
\end{equation*}

We set the light and heavy ratios as $r_L = \frac12$ and $r_H = 4$, half and quadruple the standard T5 hidden dimension respectively. For our main experiments, a fraction $\frac{1}{16}$ of tokens are routed to the heavy branch. As a result the approximate FLOPs from the feedforward layer equals

\begin{equation*}
\text{FLOPs}_{\text{FFd}} = \underbrace{4 n d^2}_{\text{Light branch}} + \underbrace{2 n d^2}_{\text{Heavy branch}}
\end{equation*}

consuming 75% of the FLOPs of a standard T5 feedforward layer.

Model     Avg     2|c|Speed    TQA     NQA      QAS     QuAL     CNLI     arXiv      SumS     QMS      GovR

inf     fn     F1     F1      F1     EM     EM     Rgm      Rgm     Rgm      Rgm

-B     43.1     0.6 / 7.4   3.7       82.2     23.0     46.6     37.9     85.6      35.4     19.2     20.4      37.7

-B     42.4     11.2   6.5     82.4        23.3     42.1     36.5       86.5     35.3     18.7     18.4     37.9

-L      45.3     0.3 / 3.0   1.3      84.2     27.2     52.3     40.6     87.3      35.7     19.1     21.4      39.5

-L      45.3     5.0   2.0      84.5     27.7     49.8    39.9     88.7      35.9     20.5     21.0      39.7

-XL      46.6     0.2 / 1.2    0.4      85.3     29.3     53.1     46.0     88.2      35.9     19.4     21.3      40.5

-XL     47.4     2.3     0.5       86.1     31.1     53.9     48.1     88.4     36.1     20.0     22.5      40.5

Performance comparison of and Base, Large and XL models on question-answering datasets TriviaQA (TQA), NarrativeQA (NQA), QASPER (QAS), and QuALITY (QuAL), NLI dataset ContractNLI (CNLI), and summarization datasets arXiv, SummScreenFD (SumS), QMSum (QMS), and GovReport (GovR). SCROLLS results are on leaderboard test set where -XL achieves SOTA. Average speed is reported in samples per second for inference (inf) and fine-tuning (fn). does not use MQA but inference speed is reported without/with MQA for conservative baseline. Rgm stands for the geometric mean of ROUGE-1,2,L.

Conditional Attention

An overview of the attention pattern. The light branch performs local attention for each token. In the higher capacity heavy branch $q$ selected query tokens (2 in the figure) attend to $v$ separately selected key and value tokens (4 in the figure).

conditional attention operates on the intuition that most tokens have simple, local interactions, but some tokens benefit from heavier processing and long-range interactions. The conditional attention layer applies an additional high-capacity attention layer that attends from selected query tokens to selected key-value tokens. Let $\tilde{s}_i^{q}$ denote the normalized routing query score for token $i$, and $\tilde{s}^{kv}$ the key-value scores for all tokens (set to $0$ if not routed). Then the attention update for is given by

\begin{equation*}
X_i = X_i + \textrm{A}_{\textrm{Light}}(X_i, X) + \tilde{s}_i^{q} \cdot \textrm{A}_{\textrm{Heavy}}(X_i, \tilde{s}^{kv}X)
\end{equation*}

The light and heavy branches differ in the number of heads and tokens attended to: the light branch has fewer heads and attends to a local context window, while the heavy branch has more heads and attends to all routed key-value tokens. Separately selecting query and key-value tokens also allows the model to differentiate between tokens that require additional information and those that possess such information. Figure <ref> shows the attention pattern. Let $q, v$ be the number of selected query and key-value tokens, $w$ the size of the local attention window and $r_L, r_H$ the proportion of light and heavy heads relative to standard T5. Then the FLOPs of the attention layer are given by

\begin{align*}
\text{FLOPs}_{\text{Att}} &= \underbrace{4 n \cdot r_L d^2}_{\text{Local projection}} + \underbrace{2 n  w \cdot r_L d}_{\text{Local attention}} \\ &+ \underbrace{2 q \cdot r_H d^2 + 2 v \cdot r_H d^2}_{\text{Global projection}} + \underbrace{2 q v \cdot r_{H} d}_{\text{Global attention}}
\end{align*}

We set the light and heavy head ratios as $r_L = \frac14$ and $r_H = \frac34$, keeping the total number of heads across the light and heavy branches equal to standard T5 heads. For our main experiments a fraction $\frac{1}{16}$ query tokens and $\frac{1}{8}$ key-value tokens are routed to the heavy branch. Ignoring local attention computation, we approximate attention FLOPS by[Global projection and attention FLOPs rounded to readable fractions, exact values are $\frac{9}{32}$ and $\frac{3}{256}$. Complexity assumes constant fraction of routed tokens; we show we can do better in practice for extremely long inputs.]

\begin{equation*}
\text{FLOPs}_{\text{Att}} \approx \underbrace{n d^2}_{\text{Local proj.}} + \underbrace{\frac{1}{4}n d^2}_{\text{Global proj.}}  + \underbrace{\frac{1}{84}n^2 d}_{\text{Global att.}}
\end{equation*}

with less than half projection FLOPs and order-of-magnitude smaller quadratic length scaling compared to . Table <ref> shows total FLOPs for the layer. In general, we set $q = m$ and $v = 2m$, and use $m$ to summarize the number of routed tokens going forward.

§.§ Multi-query Attention

Conditional computation effectively reduces the computational cost of the encoder. However, for encoder-decoder models with long inputs the majority of inference time is spent in the decoder due to memory bandwidth constraints <cit.>. We apply multi-query attention <cit.> (MQA) in cross-attention layers for much faster inference.

§.§ UL2

The UL2 pre-training objective <cit.> combines different denoising objectives and has been shown to lead to improved in-context learning. We train on UL2 instead of PEGASUS <cit.>, endowing with in-context learning capabilities.

§ EXPERIMENTS

In order to evaluate , we perform the following experiments: (1) our main results compare and on a collection of long input datasets using input length of 16k tokens; (2) we evaluate on extremely long inputs up to 64k tokens and compare scaling against ; (3) demonstrate 's few-shot capability, investigating how performance changes as input length and number of shots increase, (4) perform a series of ablations to understand the effect of individual components, and (5) investigate empirical routing patterns. The remainder of the section outlines our experimental setup, and then describes each of the experiments above.

§.§ Experimental setup

is based on the T5.1.1 architecture  <cit.>, implemented with JAX <cit.>, Flax <cit.>, and Flaxformer[https://github.com/google/flaxformer]. Following , we experiment with Base, Large, and XL model sizes. models use the same embedding dimension, number of layers, and total attention heads as corresponding models of the same size, with more overall parameters (but less compute) due to the conditional branch. See Appendix <ref> for additional details on model configuration.

We pre-train for 1M steps on a variant of the UL2 objective <cit.> with batch size 256, input length 4096, and output length 910. In particular, our mixture contains four objectives in equal proportion: prefix-LM with noise rate 0.5, and span corruption <cit.> with noise rate 0.15 and average span lengths 3, 8, and 64. We use the Adafactor optimizer <cit.> with the T5.1.1 inverse square root learning rate schedule and no dropout. is trained with the T5X <cit.> framework. For pre-training, we route $m = 512$ tokens, $\frac{1}{8}$th of the input length.

For fine-tuning we use a constant learning rate of 0.001, batch size 128, and dropout rate 0.1 for all tasks. Main results use input length of 16384 for all datasets other than ContractNLI, which uses 8192. Question answering datasets use output length 128 and summarization datasets use output length 512, except for GovRep which uses output length 1024. We route $m = 1024$ tokens, $\frac{1}{16}$th of the input length. We train until convergence and select the checkpoint with the highest dev performance. We use greedy decoding for inference.

We evaluate on TriviaQA <cit.>, arXiv <cit.>, and the SCROLLS benchmark <cit.>. SCROLLS contains question-answering datasets  NarrativeQA <cit.>, QASPER <cit.>, QuALITY <cit.>, NLI dataset ContractNLI <cit.>, and summarization datasets SummScreenFD <cit.>, QMSum <cit.>, and GovReport <cit.>. Table <ref> provides an overview of the size and input length for each dataset.

Dataset     Type     Samples     Median     90%

TriviaQA     QA      157,053     8,858     28,956

arXiv     Sum     215,913     8,519     20,170

NarrativeQA     QA     71,187     57,829     176,862

QASPER     QA     5,692     5,472     8,657

QuALITY     QA     6,737     7,171     8,276

ContractNLI     NLI     10,319     2,148     4,485

SummScreen     Sum     4,348     9,046     15,172

QMSum     Sum     1,810     14,197     27,761

GovRep     Sum      19,402     8,841     18,835

Median and 90th percentile input length by dataset measured in SentencePiece tokens.

We report time per sample per TPUv4 chip, as measured by xprof <cit.>. For inference we use a single TPUv4 with batch size 16 or the largest that fits in memory. For fine-tuning we profile with 8 TPUv4 chips, sharded separately for each model to maximize throughput.

§.§ Main results

Figure <ref> compares the quality-speed trade-off for [Note that does not use MQA, but for profiling we add MQA to for a conservative baseline.] and , showing that is better at any speed. For 16k input length, matches or exceeds quality for Large and XL with 35-75% training speedup and 50-100% inference speedup on top of the order-of-magnitude inference speedup from MQA. Encoder speedups are even greater (Appendix <ref>). -XL also achieves SOTA performance on the SCROLLS benchmark. Table <ref> contains all main results.

§.§ Scaling to extremely long inputs

scale only axis,
xlabel=Time per sample (ms),

grid style=dashed,
legend columns=1,
legend cell align=left,
legend pos=south east,
[color=longtcolor,mark=,mark size=2pt, line width=2] table

100 23.04

222 25.57

613 28.43
[color=sliccolor,mark=,mark size=2pt, line width=2] table

69 23.62

103 26.69

262 29.87

473 32.20

at (axis cs:100,23.04) [anchor= west, color=darkgrey] 8k;
at (axis cs:200,25.57) [anchor= north west, color=darkgrey] 16k;
at (axis cs:613,28.35) [anchor= north, color=darkgrey] 32k;

at (axis cs:69,23.62) [anchor= west, color=darkgrey] 8k;
at (axis cs:103,26.69) [anchor= west, color=darkgrey] 16k;
at (axis cs:248,29.74) [anchor= north west, color=darkgrey] 32k;
at (axis cs:435,32.0) [anchor= north west, color=darkgrey] 64k;

effectively scales to extremely long inputs, achieving stronger performance and faster speed than . F1 on NarrativeQA as a function of inference time per sample for and Large models using varying input lengths.

We hypothesize that the advantage of over strengthens with input length, as the fraction of important tokens decreases and can route a greater proportion of important tokens to the heavy branch. Figure <ref> compares the quality-speed trade-off for and on NarrativeQA, sweeping over input length rather than model size. The number of routed tokens is $\frac{1}{16}$th of the input length, except that we do not increase routed tokens going from 32k to 64k. achieves both stronger performance and faster inference speed at all input lengths and is able to effectively make use of extremely long inputs. We note that achieves large quality gains by going from 32k to 64k tokens even while keeping the number of routed tokens constant, providing more evidence for our hypothesis.

§.§ In-context learning

[baseline=(current bounding box.south),scale=1.0]

enlarge y limits=0.2,
major y tick style = transparent,
xmajorgrids = true,
symbolic y coords=1k, 2k, 4k, 8k, 16k,
ytick = data,
y dir=reverse,

axis y line*=none,
axis x line*=bottom,
empty legendNaturalQ
coordinates (0.0361,1k) (0.1001,2k) (0.1836,4k) (0.2651,8k) (0.3062,16k);

[baseline=(current bounding box.south),scale=1.0]

enlarge y limits=0.2,
major y tick style = transparent,
xmajorgrids = true,

symbolic y coords=1k, 2k, 4k, 8k, 16k,
ytick = data,
y dir=reverse,

axis y line*=none,
axis x line*=bottom,
/pgf/number format/fixed,
empty legendTriviaQA
coordinates (0.0137,1k) (0.0303,2k) (0.0591,4k) (0.0840,8k) (0.1055,16k);

can use its long-input capability to benefit from more shots for in-context learning. Few-shot exact match for -Large on Natural Questions and TriviaQA dev sets as a function of input tokens, fitting as many examples as possible. Each example contains question, context, and answer. Inputs length used are 1024, 2048, 4096, 8192, 16384.

Models trained on the UL2 objective have shown strong few-shot in-context learning (ICL) capabilities even at smaller sizes <cit.>. enables tractable inference with long inputs. Here, we leverage this for scaling the number of examples used for in-context learning.

We test the above hypothesis by evaluating few-shot learning performance on Natural Questions <cit.> and TriviaQA as a function of input length, using as many examples as fit in the context. We consider the open book setting, such that each example consists of question, context document, and answer. Table <ref> shows the number of examples by input length. We evaluate on the full dev set, randomly sampling examples from the training set for each dev sample until no further examples fit in the input length. We found that can perform in-context learning only up to the input length it was trained on, so for these experiments we continued pre-training a -Large model on input length 16384 for another 100k steps. For the same reason we route $m = 512$ tokens as in pre-training.

Figure <ref> displays few-shot performance as a function of input length, showing that is able to apply its long-input capabilities to extract information from increasing numbers of examples.

Dataset     1024     2048     4096      8192     16384

NQ     0.1     0.7     1.7     3.4     5.6

TriviaQA     1.6     2.3     3.8     7.0     9.8

Number of Natural Questions and TriviaQA examples that fit in input length.

§.§ Ablations

2*Ablation     2*Model     Avg     Inf     TQA     NQA      QAS     QuAL     CNLI     arX      SumS     QMS      GovR

S/s     F1     F1      F1     EM     EM     Rgm      Rgm     Rgm      Rgm

Baseline     -B        42.5     11.2     82.4        23.1     38.3     36.6       87.8     35.3     19.3     20.5     39.4

2*Routing     Static        40.5     11.6     79.7        19.2     34.2     34.5       86.4     34.9     18.1     18.9     38.8

Share QKV        42.0     11.8     82.1        21.9     37.5     36.2       87.0     35.2     18.2     20.4     39.7

2*Attention     v=all        42.5     9.4     82.4     22.3        38.6     37.2     87.8       35.3     19.1     20.3     39.8

v=q        42.3     11.5     82.5        22.5     37.3     37.0       85.9     35.2     19.0     20.5     39.7

22ptRouted Tokens     m=512        41.6     12.2     81.9        22.1     37.3     35.4       84.6     35.2     18.9     19.5     39.6

m=1536        42.9     10.4     82.6        23.5     39.8     37.5       87.5     35.4     19.4     20.8     40.0

Encoder     -B     42.1     7.4     82.0        21.4     38.4     35.8       88.0     35.5     18.7     20.4     38.5

Decoder     Multi-head     42.9     0.7     82.7        22.9     40.2     35.8       87.7     35.5     19.7     21.2     40.3

Objective     PEGASUS         42.8     11.2     82.6        22.6     40.5     37.3       87.3     35.3     19.6     20.8     39.6

ablations. Each experiment modifies a component of the recipe for -Base. Static routing divides the input into equal-length blocks and selects the first token in each block to be routed. Shared QKV routing shares routing decisions for queries and keys/values. In v=all the routed queries attend to the entire input, while v=q selects the same number of key and value tokens as query tokens. m=512 and m=1536 use different numbers of routed tokens. -B uses a encoder while retaining other parts of the training recipe such as MQA and the UL2 objective. Multi-head refers to using multi-head cross-attention. The final ablation replaces the UL2 objective with PEGASUS as in .

This section studies the effect of different choices in the recipe. Table <ref> contains results of a series of experiments that change a single component for Base.

Routing First, we note that static routing -- evenly distributing routed tokens over the input -- leads to massive drop in performance. The importance of routing provides evidence that the model learns to devote capacity to important tokens and the advantage of is not merely a result of additional parameters. Sharing routing decisions for query and KV tokens should be compared with v=q, and leads to a modest reduction in quality and increase in speed.

The optimal number of routed tokens represents a trade-off between improved performance and computational cost of applying heavier layers. Table <ref> shows strong gains going from 512 to 1024 (baseline) routed tokens and diminishing returns for further increases.

relies on routing to identify not only tokens that can benefit from important information elsewhere in the input, but also which tokens contain such important information. We study whether is successful in this task by comparing performance with two different attention settings -- v=all, in which routed tokens attend to the entire input, and v=q, which uses equal number of routed keys and values as queries, rather than twice as many. appears to occupy a sweet spot, as using fewer routed key-values modestly decreases performance at similar speed but attending to all inputs barely helps at sharply increased cost.

We compare to with multi-query cross-attention, confirming that indeed does not achieve an unexpected quality gain from MQA, and our conservative assumptions in Figures <ref>, <ref> are valid. Next, we evaluate multi-head cross-attention for , finding that it leads to modestly improved performance. However, as MHA exhibits order-of-magnitude slower inference, MQA is clearly favored. Finally, PEGASUS appears to fine-tune slightly better than UL2, though the difference is small and UL2 enables few-shot learning.

§.§ Routing analysis

bar width=.4cm,

ylabel=Routed proportion,
symbolic x coords=MLP, Query, KV,

enlarge x limits=0.25,

legend style=at=(0.25,0.95),
anchor=north,legend columns=1
[fill=othercolor] coordinates (MLP,0.212) (Query,0.211) (KV,0.343);
[fill=answercolor] coordinates (MLP,0.236) (Query,0.262) (KV,0.432);
[fill=questioncolor] coordinates (MLP,0.396) (Query,0.365) (KV,0.608);
Other, Answer, Question

Proportion of tokens routed for answer (string match), question, and other tokens by routing component for Large model, averaged over examples in TriviaQA dev set and all layers of model.

It is interesting to ask whether routed tokens line up with what we consider intuitively important tokens in each document. We investigate this question by studying routing patterns of a Large model fine-tuned on TriviaQA. We divide tokens into three categories: (1) question tokens, (2) answer tokens, and (3) other tokens. Figure <ref> shows the average fraction of each type of token that is routed through the heavy path for MLP and attention layers on TriviaQA. We note that question and answer tokens are significantly more likely to be routed than other tokens, for feedforward as well as attention queries and keys/values. Appendix <ref> presents more detailed routing analysis; e.g., semantically important tokens are much more likely to be selected in later layers.

§ CONCLUSION

We propose , a new model for long-range inputs that employs conditional computation for higher quality and faster speed. has light feedforward and attention layers that apply to the entire input, as well as heavy branches that are applied only to a subset of important tokens selected by a learned router. We show that achieves stronger performance at any speed compared to on a variety of long-input datasets, and can effectively and efficiently make use of extremely long inputs up to 64k tokens.

Model     Layers     Model dim     MLPlight dim     MLPheavy dim     Headslight     Headsheavy     Params

-B     12     768     2048     N/A      12     N/A     248m

-B     12     768     1024     8096      4     8     433m

-L     24     1024     2816     N/A      16     N/A     783m

-L     24     1024     1408     11264      4     12     1462m

-XL     24     2048     5120     N/A      32     N/A     2850m

-XL     24     2048     2560     20480      8     24     5297m

Hyperparameters for and models. T5.1.1 hyperparameters match . parameters are sparsely accessed as a result of conditional computation, so parameter counts do not reflect compute.

Model     2c|Average     2c|16k in, 128 out     2c|16k in, 512 out
2c|16k in, 1024 out     2c8k in, 128 out

Enc     Tot    Enc     Tot    Enc     Tot    Enc     Tot    Enc     Tot

-B    77     136     84     98     84     165     84      296     27     39

-B    29     90     30     45     30     113     30     256     18     30

-L    164     329     173     222     179     392     179     799      66     100

-L     70     201     73     103     73     250     73     578     45     69

-XL     390     870     412     557     423     1081     423     2065     166     290

-XL     177     439     185     239     185     525     185     1253     115     163

Comparison of total and encoder inference time per sample (ms) for and Base, Large, and XL models at different input and output lengths. Average time per sample is computed as a weighted average over input and output lengths, weighted by the number of tasks in our evaluation that use the corresponding setting (4 for 16k/128, 3 for 16k/512, and one each for 16k/1024 and 8k/128).

§ CONTRIBUTIONS

Joshua led the project, developed the initial conditional attention mechanisms, and conducted most experimental ablations.  Tao developed the heavy/light formulation for heterogeneous conditional computation, comprising the routing and conditional feedforward mechanisms, and iterated with Joshua on initial experiments demonstrating feasibility.  Michiel helped to scope the paper, performed most of the writing, and oversaw speed benchmarking.  Santiago designed and conducted all the few-shot experiments, initiated the routing analysis visualization, and integrated UL2 into the codebase.  Siddhartha developed the separate routing for query and key/value tokens in the conditional attention component and demonstrated the resulting quality improvements.  Yury designed and conducted all experiments for inputs larger than 16k tokens, demonstrating favorable scaling up to 64k.  David integrated all SCROLLS tasks into the codebase and ran early experiments, especially comparing UL2 with PEGASUS.  Mandy developed the leaderboard comparisons with LongT5 and helped run several experiments.  James advised on and ran early comparisons with MoE conditional computation.  Yi advised on the adaptation of UL2 to 4k input length pre-training.  Finally, Yun-Hsuan and Sumit provided guidance and support for the project overall.

§ MODEL HYPERPARAMETERS

Table <ref> shows and hyperparameters, including parameter counts. For , we report numbers for the TGlobal configuration, which match T5.1.1. Notice that 's parameter counts are larger due to using conditional compute. Similar to other conditional compute architectures such as mixture-of-experts, computational cost does not necessarily increase with parameter count.

We use the same 127-token local radius for as .  This results in a local attention window $w$ of 255 since 127 tokens are attended to the left and 127 to the right.

§ ROUTING NORMALIZATION HYPERPARAMETERS

To normalize the routing scores for differentiable top-$k$ token selection, we use the iterative soft top-$k$ algorithm from <cit.> and <cit.> with $\epsilon = 1.0$ and 50 iterations.  During training we allow the top $\frac{9}{8} k$ tokens to have nonzero weight instead of just the top $k$ in order to provide a slightly improved training signal.

§ ADDITIONAL EXPERIMENTAL RESULTS

Table <ref> compares and inference speed in more detail, splitting off encoder and total time per sample. Since applies conditional computation only in the encoder, encoder speed gains are larger than overall speed gain, and total speed gains are largest for shorter output length. Trade-offs are even more in the favor of when paired with other decoder optimizations.

Table <ref> shows full (Rouge-1, Rouge-2, Rouge-L) results for summarization datasets.

scale only axis,

ylabel=Routing proportion,

x tick label style=log ticks with fixed point,
grid style=dashed,

legend pos=north west,
ylabel style=align=left, text height=0.2cm,

empty legendMLP
[color=othercolor,mark=,mark size=1pt, line width=1] table
0 0.21161997946429376
1 0.2125872223817629
2 0.21139287510891192
3 0.2120075226278527
4 0.21143393511373654
5 0.21185861990537036
6 0.2113462075224012
7 0.21165094318084943
8 0.21246731457213297
9 0.21207230297236887
10 0.21240816743978747
11 0.21173454768636082
12 0.21308165538967852
13 0.21261318091433634
14 0.20987843682723828
15 0.2119161537179608
16 0.21203253878894812
17 0.21229383346521144
18 0.21107189040051932
19 0.21156451254256683
20 0.20975093698681574
21 0.21084240583540229
22 0.21221006045407584
23 0.2080541786840686
[color=answercolor,mark=,mark size=1pt, line width=1] table
0 0.23404124981875588
1 0.12133974985640475
2 0.2773626670220826
3 0.1990351298117969
4 0.26216721309024876
5 0.24163464687057307
6 0.2982911783714974
7 0.2850873355229606
8 0.21702661150057342
9 0.27718464853413505
10 0.18387851142672684
11 0.2737455132017648
12 0.11462717450160506
13 0.14369442677229646
14 0.5881649708680756
15 0.10881932430568753
16 0.12051417298406941
17 0.09281284504910445
18 0.1376113019123887
19 0.18240325725650058
20 0.4718293490780316
21 0.10874523946867246
22 0.12584635703993513
23 0.6010882598061337
[color=questioncolor,mark=,mark size=1pt,line width=1] table
0 0.20856889202407008
1 0.21226976756629298
2 0.2629011370564126
3 0.2335030021419846
4 0.2742826105589716
5 0.21929366584665527
6 0.325394583890978
7 0.28125829559056403
8 0.23776303913582714
9 0.22415379958167925
10 0.2285775261957093
11 0.2940282392779054
12 0.1701509477917461
13 0.25734053143717994
14 0.3895200056455131
15 0.5509839725342021
16 0.5038410777931918
17 0.43938062321848653
18 0.6880240312117615
19 0.5884735326060652
20 0.7361686221028755
21 0.7852749420317643
22 0.43978855486954743
23 0.9620789504443227

scale only axis,

x tick label style=log ticks with fixed point,
grid style=dashed,

legend pos=north west,
ylabel style=align=left, text height=0.2cm,

empty legendQuery
[color=othercolor,mark=,mark size=1pt, line width=1] table
0 0.19549139457384002
1 0.21209221538348624
2 0.2125791879007595
3 0.2105169686202275
4 0.2112944328572769
5 0.21240783221189677
6 0.21176481164882838
7 0.2130972302762573
8 0.20935626333610893
9 0.2130512618262021
10 0.20952047660539894
11 0.21263229572024034
12 0.21287331516477234
13 0.2125744576897517
14 0.20913046305089916
15 0.21240088471819554
16 0.21199353612183974
17 0.21297168167351868
18 0.20954572724295997
19 0.20935358805721618
20 0.21268238211396587
21 0.20873566927429005
22 0.21232251657333964
23 0.20697159425681427
[color=answercolor,mark=,mark size=1pt, line width=1] table
0 0.09856966095334675
1 0.14060753632977327
2 0.13186202080930925
3 0.40304072073742486
4 0.24739154421037465
5 0.1672355452586664
6 0.23150961204544812
7 0.11160585670330919
8 0.46910581322547396
9 0.14235919997308963
10 0.3146520498904203
11 0.15289618580039627
12 0.1479256285183321
13 0.11205143872629633
14 0.5839836673398893
15 0.1519279027796636
16 0.09239308751278266
17 0.11036272263682104
18 0.4037651750456416
19 0.37856251427068704
20 0.0901048019000607
21 0.6743836892905413
22 0.08463163214863184
23 0.8586354041149576
[color=questioncolor,mark=,mark size=1pt, line width=1] table
0 0.2099036197909046
1 0.22926699697975958
2 0.211451565719695
3 0.26243562420203304
4 0.2800879162310573
5 0.2684250004017433
6 0.3169072893202554
7 0.13571953945826448
8 0.463885009902794
9 0.06026191930443222
10 0.796024658192565
11 0.12316350899062913
12 0.0691505313549957
13 0.22311366730813595
14 0.5122468448194271
15 0.22110591251264508
16 0.49468819575532663
17 0.14179612527010294
18 0.7373704741411858
19 0.8387805202703116
20 0.2971193860421291
21 0.6228262860671179
22 0.412059198614116
23 0.8251401341248926

scale only axis,

x tick label style=log ticks with fixed point,
grid style=dashed,

legend pos=north west,
ylabel style=align=left, text height=0.2cm,

empty legendKV
[color=othercolor,mark=,mark size=1pt, line width=1] table
0 0.35019249035137523
1 0.3525438069145887
2 0.35293969478952314
3 0.3512627081448169
4 0.35180610639939575
5 0.3516696604571174
6 0.3516573057818594
7 0.35250147529909165
8 0.3502812027267641
9 0.35110776222753687
10 0.3519654084233534
11 0.35073425660566815
12 0.3526173688518903
13 0.3527463777811796
14 0.2456468071859964
15 0.35132045050321287
16 0.3525094426405284
17 0.3510333293545492
18 0.2792762573520494
19 0.35100672693816265
20 0.35176224045963533
21 0.3321378796247362
22 0.3503188617750819
23 0.34622569438400463
[color=answercolor,mark=,mark size=1pt, line width=1] table
0 0.37919229692764433
1 0.32862221726090235
2 0.2663435888759715
3 0.5158641264872711
4 0.4348650637532726
5 0.42079635832020107
6 0.44208448995530036
7 0.31435037054317644
8 0.6003177818950836
9 0.3867071926342268
10 0.27786296098789925
11 0.464732708828246
12 0.29523005085614473
13 0.2391429766550423
14 0.35034875120751774
15 0.4691989400084632
16 0.287680735209191
17 0.3805258248305361
18 0.4406435230754643
19 0.5181929256800606
20 0.38277884263242634
21 0.6655974345329255
22 0.6420381393763785
23 0.862727037011765
[color=questioncolor,mark=,mark size=1pt, line width=1] table
0 0.3574036293389754
1 0.3714281853222321
2 0.40176561417062046
3 0.45976470574889095
4 0.41888616643100607
5 0.5517346042962912
6 0.5092366638237913
7 0.4357807900764302
8 0.6431377355400922
9 0.8062974080445435
10 0.7008200519289048
11 0.9459240178802629
12 0.2784017532337479
13 0.5438503167829559
14 0.32421232699312585
15 0.7310340349223263
16 0.59942834172575
17 0.7376625469667542
18 0.5351831860687264
19 0.8712681993442024
20 0.7532590154318346
21 0.7746499807418749
22 0.8897478233991253
23 0.962236743816725

legend columns=3,
legend style=
column sep=4ex,

legend entries=Other, Answer, Question
mark=,mark size=3pt,solid,color=answercolor

Proportion of tokens routed for answer (string match), question, and other tokens by routing component and layer for Large model, averaged over examples in TriviaQA dev set.

Visualization of token routing weights for some fragments of an example on TriviaQA.

§ COMPUTATIONAL RESOURCES

For pre-training we generally used 128 TPUv4 chips for Base and 256 TPUv4 chips for Large and XL. Pre-training took approximately 2.5 days for Base, 3.7 days for Large, and 12.8 days for XL. For fine-tuning we generally used 64, 128, and 256 TPUv4 chips for Base, Large, and XL, respectively, with training time varying with dataset size.

§ ROUTING ANALYSIS

In this section we take a closer look at the routing mechanisms in . There are three routing processes in each layer of : (1) Routing of attention keys and values (“KV-routing”), (2) routing of attention queries (“Q-routing”) and (3) routing of MLP tokens (“MLP-routing”). For simplicity, we will say that a token is selected, when it is routed to the heavy alternative (of either MLP or attention). We are interested in understanding what tokens are selected and whether these mechanisms select similar or different tokens in each layer.

Which tokens are selected

We divide input tokens into three categories: (1) question tokens, (2) answer tokens (found via simple normalized string match of the ground truth answer), and (3) other tokens. Figure <ref> shows the proportion of each token type that is routed by a fine-tuned -Large model on the TriviaQA dev set, by layer and routing component.

Earlier we showed that question and answer tokens are more likely to be selected, but separating routing decisions by layer reveals interesting patterns. At early layers question and answer tokens are only modestly more likely to be selected, with routing probability sharply increasing at later layers and peaking in the last layer. This makes intuitive sense: in early layers the model has not yet had the opportunity to identify which tokens and parts of the document are important. However, the increase is not monotonic and there is strong variation between layers. This variation may imply that different layers focus on different types of tokens, or that some routing components do not successfully learn to identify important tokens.

To gain a better insight into this, Figure <ref> visualizes routing on two sample fragments from a TriviaQA example (notice that, given the large input length used in , we do not show the complete example in the figure). The two fragments shown correspond to the beginning of the example (where the question is located), and the part of the context surrounding the correct answer. We have added a colored background to the figure, where each of the three CMY channels are mapped to the KV-routing weights in different layers of the model. Cyan corresponds to layer 1, Magenta to layer 12, and Yellow to layer 24. As we can see, question and answer are heavily yellow colored, showing those tokens are selected in the last layer.

Correlation between routing processes. Table <ref> shows the Pearson correlation coefficient between the routing weights of the different routing mechanisms in each layer in a  Large model (MLP-routing correlation with KV-routing, MLP-routing with Q-routing, and KV-routing with Q-routing). We show numbers for both the pre-trained checkpoint, as well as a fine-tuned model on TriviaQA. As we can see, the routing of keys/values and routing of queries is highly correlated at all layers except the first two, while the routing of tokens in the MLP has lower correlation to the other two processes. Interestingly correlation between MLP and attention routing increases in the last layers of the model.

3c||Pre-trained     3c|Fine-tuned
MLP-KV     MLP-Q     KV-Q     MLP-KV     MLP-Q     KV-Q
1     -0.06    	-0.06    	-0.09    	-0.06    	-0.09    	-0.26

2     0.27    	0.52    	0.04    	0.27    	0.39    	0.02

3     -0.05    	-0.03    	0.75    	0.05    	-0.01    	0.69

4     0.05    	0.09    	0.76    	0.18    	0.14    	0.72

5     0.02    	-0.01    	0.75    	0.22    	0.26    	0.68

6     0.02    	-0.01    	0.78    	0.31    	0.33    	0.70

7     0.02    	0.00    	0.73    	0.26    	0.27    	0.70

8     0.00    	-0.02    	0.44    	0.11    	-0.07    	0.29

9     0.13    	0.11    	0.74    	0.36    	0.40    	0.70

10     -0.06    	-0.08    	0.08    	-0.15    	-0.15    	0.12

11     -0.05    	-0.07    	0.31    	-0.08    	-0.03    	0.18

12     -0.04    	-0.08    	0.27    	0.03    	0.00    	0.28

13     -0.10    	-0.09    	0.87    	-0.13    	-0.03    	0.72

14     -0.04    	-0.05    	0.76    	-0.06    	-0.12    	0.67

15     0.53    	0.64    	0.69    	0.51    	0.55    	0.67

16     0.08    	0.12    	0.63    	0.06    	0.57    	0.24

17     0.28    	0.30    	0.65    	0.27    	0.32    	0.69

18     0.28    	0.02    	0.84    	0.31    	0.20    	0.76

19     0.45    	0.77    	0.59    	0.19    	0.38    	0.64

20     0.30    	0.39    	0.64    	0.38    	0.47    	0.62

21    	0.05    	-0.04    	0.49    	0.18    	0.11    	0.47

22     0.05    	0.00    	0.69    	0.21    	0.16    	0.68

23     0.39    	0.33    	0.68    	0.60    	0.79    	0.69

24     0.43    	0.39    	0.59    	0.57    	0.63    	0.65

Pearson correlation coefficient between the routing weights of the different routing mechanisms in each layer in a  Large model. We show numbers for both the pre-trained checkpoint, as well as a fine-tuned model on TriviaQA. Blue bars visualize positive correlation, whereas red bars visualize negative correlation.