CoLT5: Faster Long-Range Transformers with Conditional Computation
Authors: Joshua Ainslie
Tao Lei
Michiel de Jong
Santiago Ontanon
Siddhartha Brahma
Yury Zemlyanskiy
David Uthus
Mandy Guo
James Lee-Thorp
Yi Tay
Yun-Hsuan Sung
Sumit Sanghai
Venue: EMNLP 2023 Main (2023)
Keywords: long context
conditional computation
efficient nlp
TL;DR: We present a long-input conditional computation model that scales to 64k tokens and sets a new SCROLLS state of the art.

Abstract: 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 CoLT5, 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 CoLT5 achieves stronger performance than LongT5 with much faster training and inference, achieving SOTA on the long-input SCROLLS benchmark. Moreover, CoLT5 can effectively and tractably make use of extremely long inputs, showing strong gains up to 64k input length.

CoLT5: Faster Long-Range Transformers with Conditional Computation
Joshua Ainslie, Tao Lei, Michiel de Jong, Santiago Ontañón Siddhartha Brahma, Yury Zemlyanskiy, David Uthus, Mandy Guo James Lee-Thorp. Yi Tay, Yun-Hsuan Sung, Sumit Sanghai
Google Research
Abstract
Many natural language processing tasks bene- fit from long inputs, but processing long doc- uments with Transformers is expensive not only due to quadratic attention complexity but also from applying feedforward and projec- tion layers to every token. However, not all tokens are equally important, especially longer documents. We propose CoLT5, long-input Transformer model that builds on this intuition by employing conditional compu- tation, devoting more resources to important tokens both feedforward and attention lay-
Light MLP
Heavy MLP
m
Routi ng
Light Attention
Heavy Attention
We show that CoLTS achieves stronger performance than LONGTS with much faster training and inference, achieving SOTA on long-input SCROLLS benchmark. Moreover, CoLTS can effectively and tractably make use of extremely long inputs, showing strong gains up to 64k input length.
Routing
Light branch
Heavy branch
Figure 1: An overview a COLT5 Transformer layer with conditional computation. All tokens are processed by light attention and MLP layers, while routed query tokens perform heavier attention over e routed key- value tokens and routed tokens processed by heavier MLP.
1 Introduction
Many natural language processing tasks, such as summarization (Cohan et al., 2018) or question an- swering over long documents (Joshi et al., 2017). require machine learning models to encode long- form text. Processing long documents with a Trans- former model is computationally expensive, both because attention cost scales quadratically with in- put length and because feedforward and attention projection layers have to be applied to each input token.
the computational burden and can render process- ing long inputs intractable.
This paper presents CoLT5 (Conditional LongT5), new family of models that, building top of LONGT5 (Guo al., 2022), enables fast pro- cessing of long inputs by combining architecture improvements for both attention and feedforward layers. CoLTS 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 dimin- ish with document length, allowing for tractable processing of long documents.
Over the past few years, many "efficient Trans- former" approaches have been proposed that re- duce the cost of the attention mechanism over long inputs (Child et al., 2019; Ainslie al., 2020; Belt- agy et al., 2020; Zaheer al., 2020; Wang al., 2020; Tay et al., 2021; Guo et al., 2022). However, especially for larger models, the feedforward and projection layers actually make up the majority of
In particular, CoLTS divides each feedforward layer and each attention layer into a light branch
"Author contrihutions are outlined in Appendix A. Corre- spondence author: jainslic @ @googe..com

---

Inference
Fine-tuning
XL
Average perf
46
XL
46
XL.
44
44
B
42
42
200
100
600
800
500
1,000
1,500 2,000 2,500
Time per sample (ms)
Time per sample (ms)
LONGT5
CoLT5
Figure 2: COLTS achieves stronger performance than LONGTS at any speed. Average performance on all datasets as function of inference fine-tuning time per sample (ms) for LONGT5 and CoLT5 Base, Large, and XL models. LONGT5 does not MQA. but report speed though it had for conservative baseline.
which is applied to all tokens and heavy branch which is applied to a set of important tokens, se- lected specifically for that input and component. The light feedforward branch has lower hidden di- mension than standard LONGT5 while the heavy feedforward branch has higher hidden dimension.
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 help- ful to start by providing a breakdown of the compu- tational cost of Transformer components. Table shows the FLOPs' for each component of Trans- former encoder layer (Kaplan et al., 2020).
The light attention branch has fewer heads and ap- plies only local attention, while the heavy attention branch performs full attention over another sepa- rately selected set of important tokens. Figure provides an overview of the CoLT5 conditional mechanism.
Encoder Layer Component
Flops
Vanilla self-attention computation
2n
Attention QKV and output projections Feedforward layer LONGT5 local attention computation LONGT5 global attention computation #d
4nd2
Finally, CoLT5 also includes two other mod- ifications to the LONGTS architecture. CoLTS adds multi-query cross-attention (Shazeer, 2019), significantly speeding up inference. CoLT5 also employs the UL.2 (Tay et al., 2022) pre-training ob- jective, which we demonstrate allows for in-context learning over long inputs.
8nd2
2nadd n²
Table 1:
Computational cost of encoder layer trans- former components measured in FLOPs. is the input length, is the model dimensionality. and ws is the of the local attention window.
We show that CoLT5 performs much faster fine- tuning and inference with similar or better model quality, improving over LONGT5 5 on arXiv summa- rization (Cohan et al., 2018) and TriviaQA question answering (Joshi et al., 2017) datasets and achiev- ing SOTA on the SCROLLS benchmark (Shaham et al., 2022). Moreover, CoLT5 achieves further gains in quality and speed for tasks with extremely long inputs (64k tokens), with less-than- linear seal- ing of "focus" tokens.
Sparse attention The first challenge applying a Transformer to a long input is that the FLOPs of the self-attention mechanism scales quadrati- cally in the input length, becoming intractable for long inputs. A large body of work focuses on ducing self-attention cost, restricting attention be- tween subset of inputs (Child al., 2019; Ainslie et al., 2020; Beltagy et al., 2020; Zaheer al., 2020; Wang et al., 2020: Guo et al., 2022) or a subset of layers (Zemlyanskiy et al., 2021). ). LONGT5 (Guo et al., 2022), the most closely re- lated model to CoLT5, tokens attend within a lo-
2 Background
Transformer FLOPS CoLTS follows an exten- sive line work attempting to reduce the com- putational cost of Transformer models, particularly
Each multiply-add is counted as a single FLOPP.

---

cal window as well as mean--oolee summary representation for each block of 16 tokens in the input. LONGT5 attention leads to sharply reduced (though still non-negligible) FLOPs (Table I).
tion objective for improved summarization perfor- mance. Tay et al. (2022) proposes UL2, a mixture of span corruption, prefix, and causal language modeling, and shows that it leads to strong perfor- mance on both short-output and generative tasks.
Conditional computation After applying sparse attention mechanism, the feedforward and attention projection layers account for the major- ity of the FLOPs. These costs scale with the length of the input, such that processing long in- puts 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 the entire input. CALM (Schuster et al., 2022) applies 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 (Shazeer et al., 2017; Fedus et al., 2021; Zoph et al., 2022) route inputs through small proportion of expert sub-modules, bringing to bear only the parame- ters most relevant to the input. In the context of retrieval-uugmented models, numerous works re- rank retrieved passages by their relevance to the query and process only the highest scoring pas- sages (Mao et al., 2021; Wang et al., 2018; Yu et al., 2022) and vary the number of processed pas- sages depending on model confidence (Kratzwald and Feuerriegel, 2018; Varshney et al., 2022). Con- current work CoDA (Lei et al., 2023) employs related conditional computation mechanism, de- signed for efficient adaptation rather than modeling long documents.
CoLT5
3.1 Conditional computation
As discussed in the previous section, a large propor- tion of Transformer FLOPs arise from feedforward and projection layers that scale with the length of the input sequence. Therefore, LONGT5 training and inference long documents remains expen- sive.
CoLT5 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 to- kens 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 process- ing stage.
The CoLT5 conditional computation mecha- nism consists of three components: routing mod- ules, conditional feedforward layers, and condi- tional attention layers. All tokens are processed by standard, lightweight attention and feedforward lay- ers. Routing modules additionally select important tokens from an input at each attention or feedfor- ward layer, and a heavy conditional layer applies additional computation to routed tokens. This sec- tion describes each component in detail. Figure provides overview of the CoLT5 conditional computation mechanism, and Table 2 compares CoLTS and LONGTS FLOPs.
Device utilization FLOPs do not tell the whole story, modeling choices can influence the effec- tive speed of operations achieved by accelerators. For long text inputs, autoregressive decoder infer- ence is very slow due to memory bandwidth con- straints from repeatedly loading the long sequence of keys and values (Shazeer, 2019; de Jong et al., 2022). Shazeer (2019) introduces multi- query at- tention (MQA), sharing heads for keys and values reduce memory bandwidth overhead. Pope et al. (2022) studies how to shard large models, espe- cially in the context of MQA, to obtain optimal device utilization and therefore speed.
Model
Encoder Layer Flops
T5
12nd2 2n
LONGT5
12nd + 5 d
COLT5
nd + odd
Table 2: COLTS uses significantly fewer FLOPs than LONGT5. Comparison of approximate encoder layer total FL.PPs between T5, LONGTS, and COLT5. CoLT5 FLOPs rounded to readable fractions.
Training objectives T5 introduced the span cor- ruption objective (Raffel et al., 2020), modi- fication of masked language modeling (Devlin et al., 2019). LONGTS made use of the PEGA- SUS (Zhang et al., 2020) sentence reconstruc-
Routing In order to separately select important tokens for each component in each layer, we need

---

learnable and tractable routing function. We follow the simple three-step mechanism from Lei et al. (2023): (1) multiply inputs with learned embedding to obtain routing scores, (2) normalize, and (3) select the top-k highest scoring inputs.
Attention keys
Attention queries
Let X; be the representation of token i, and a d-dimensional learnable embedding. Then the routing score of token is
8
X; 76
We select the top-k highest scoring inputs. In order to provide learning signal to the scoring embed- ding, 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 bet- ter distributed signal to all tokens, we also globally normalize the routing scores to sum up to the num- ber of desired routed tokens using a generalized softmax, resulting in normalized scores $,. Each CoLT5 layer has three independent routers, one each for the feedforward layer, attention queries, and attention key-values.
Figure 3: An overview of the CoLT5 attention pattern. The light branch performs local attention for each to- ken. In the higher capacity heavy branch q selected query tokens (2 in the figure) attend to separately se- lected key and value tokens (4 in the figure).
heavy branch. As result the approximate FLOPs from the CoLT5 feedforward layer equals
Conditional Feedforward Intuitively, some to- ken representations may benefit from more pro- cessing than others. The CoLT5 conditional feed- forward layer applies an additional high-capacity feedforward layer to selected tokens. In particular, let X; be the model state of the ith token and 8; denote the normalized routing score (set to 0 for non-routed tokens). Then the feedforward update for CoLT5 is given by
FLOPPprr
And
+
2nd
Light branch
Henvv branch
consuming 75% of the FLOPs of standard T5
feedforward layer.
Conditional Attention CoLT5 conditional at- tention operates the intuition that most tokens have simple, local interactions, but some tokens benefit from heavier processing and long-range in- teractions. The CoLT5 conditional attention layer applies an additional high-capacity attention layer that attends from selected query tokens to selected key-value tokens. Let sy denote the normalized routing query score for token i, and g/tr the key- value scores for all tokens (set to 0 if not routed). Then the attention update for CoLT5 is given by
X; = X; + Fighiig((X) + 8; FFd (X;)
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 to- kens, m the number of selected tokens, and ry and Th the ratios of light and heavy hidden dimension to standard T5 hidden dimension. Then the FLOPs of the CoLT5 layer are given by
X; = X; + ALightl X;; X) +)+33 Atteay (X,, ghu =
X)
The light and heavy branches differ in the number of heads and tokens attended to: the light branch has fewer heads and attends to local context win- dow, while the heavy branch has more heads and attends all routed key-value tokens. Separately selecting query and key-value tokens also allows the model to differentiate between tokens that re- quire additional information and those that possess
FLOPSFFa = d + dd
Light branch
Heavy branch
We the light and heavy ratios rz and Th 4, half and quadruple the standard T5 hid- den dimension respectively. For our main exper- iments, a fraction i of tokens are routed to the
16

---

Model
Avg
Speed
TQA
NQA
QAS
QuAL
CNLI
arXiv
SumS
QMS
GovR
inf
fn
F1
F1
F1
EM
EM
Rgm
Rgm
Rgm
Rg-
LONGTS-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
CoLT5-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
LONGTS-L
45.3
0.3/3.0
84.2
27.2
52.3
40.6
87.3
35.7
19.1
COLTS-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
LONGTS-XL
46.6
0.27 1.2
0.4
85.3
29.3
53.1
46.0
88.2
35.9
19.4
21.3
40.5
COLTS-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
Table 3: Performance comparison of CoLTS and LONGTS Base, Large and models on question-answering datasets TriviaQA (TQA), NarrativeQA (NQA), QASPER (QAS), and QuALITY (QuAL), NLI dataset Con- tractNLI (CNLI), and summarization datasets arXiv, SummScreenFD (SumS), QMSum (QMS), and GovReport (GovR). SCROLLS results are on leaderboard test where CoLT5-XL achieves SOTA. Average speed is re- ported in samples per second for inference (inf) and fine-tuning (fn). LONGT5 does not use MQA but inference speed is reported without/with MQA for conservative baseline. Rgm stands for the geometric mean of ROUGE- I,2,L., Similar to SCROLLS, we take simple average across all datasets even though the datasets different performance metrics.
such information. Figure 3 shows the CoLTS at- tention pattern. Let q, be the number of selected query and key-value tokens, w the size of the lo- cal attention window and T1, TH the proportion of light and heavy heads relative to standard T5. Then the FLOPs of the CoLT5 attention layer are given by
3.2 Multi-query Attention
Conditional computation effectively reduces the computational cost of the encoder. However, for encoder-decoder models with long inputs the ma- jority of inference time is spent in the decoder due to memory bandwidth constraints (Shazeer, 2019; de Jong et al., 2022). Most of the overhead is caused by repeatedly reading all the input token keys and values from memory for every output to- ken that is autoregressively decoded during cross attention. Multi-query attention (Shazeer, 2019) (MQA) allows all query heads to share single key and value head, alleviating this bottleneck. Accord- ingly, apply MQA in cross-attention layers for much faster inference. Note however that MQA does improve training speed since target tokens are processed in parallel during training, avoiding this memory bandwidth bottleneck.
FLOPSAR
4n r + 2nw rid
Local projection
Local attention
+ 2q r + 2v d
2qu Thd
Global projection
Global attention
We set the light and heavy head ratios as rz and FH 4. keeping the total number of heads across the light and heavy branches equal to stan- dard T5 heads. For our main experiments a fraction 1 16 query tokens and key-value tokens are routed to the heavy branch, so q = 16 ing local attention computation, we approximate attention FLOPS by?
n and e = 8. Ignor-
3.3 UL2
The UL2 pre-training objective (Tay et al., 2022) combines different denoising objectives, extending the span corruption pre-training used in T5 to a variety of noise rates / average span lengths and adding a prefix language modeling objective more similar to typical decoder-only model pre-training. UL.2 has been shown lead to improved in-context learning. We train CoLT5 on UL2 instead of PE- GASUS (Zhang et al., 2020), endowing CoLT5 with in-context learning capabilities.
1 n'd
FLOPSAn
nd
+
nd²
84
Local proj.
Global proj.
Global att.
with less than half projection FLOPs and order-of- magnitude smaller quadratic length scaling com- pared to LONGT5. Table 2 shows total FLOPs for the COLTS layer. In general, we set g m and 2m, and use to summarize the number of routed tokens going forward.
4 Experiments
'Global projection and attention FLOPss rounded to read- 9 able fractions. exact values are 37 and a Complexity as- 256 sumes constant fraction of routed tnkens; we show we can do hetter in practice for extremely long inputs.
In order to evaluate COLTS, we perform the fol- lowing experiments: (1) our main results com-

---

pare COLTS and LONGT3 on collection of long input datasets using input length of 16k tokens; (2) we evaluate CoLT5 on extremely long in- puts up to 64k tokens and compare scaling against LoNGT5; (3) demonstrate CoLT5's few-shot ca- pability, investigating how performance changes as input length and number of shots increase, (4) per- form series of ablations to understand the effect of individual CoLT5 components, and (5) inves- tigate empirical routing patterns. The remainder of the section outlines our experimental setup, and then describes each of the experiments above.
and select the checkpoint with the highest dev per- formance.
We use greedy decoding for inference.
Data We evaluate CoLT5 on TriviaQA (Joshi et al., 2017), arXiv (Cohan et al., 2018), and the SCROLLS benchmark (Shaham et al., 2022). SCROLLS contains question-answering datasets: NarrativeQA (Kočiský et al.. 2018), QASPER (Dasigi et al., 2021), and QuAL. ITY (Pang et al., 2021), NLI dataset: Con- tractNLI (Koreeda and Manning. 2021), and sum- marization datasets: SummSereenFD (Chen et al., 2022), QMSum (Zhong et al., 2021), and Gov- Report (Huang et al.. 2021). Table 4 provides an overview of the size and input length for each dataset.
4.1 Experimental setup
Configurations CoLT5 based on the T5.1.1 architecture (Raffel et al., 2020), implemented with JAX (Bradbury et al., 2018). Flax (Heek et al., 2020), and Flaxformer'. Following LONGTS, we experiment with Base, Large, and XL model sizes. CoLT5 models use the same embedding dimen- sion, number of layers, and total attention heads as corresponding LONGT5 models of the same size, with more overall parameters (but less compute) due to the conditional branch. See Appendix for additional details on model configuration.
Dataset
Type
Samples
Median
90%
TriviaQA
QA
157,053
8,858
28,956
arXiv
Sum
215.913
8,519
20,170
Narrative
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
Pre-training We pre-train CoLT5 for IM steps on the C4 dataset (Raffel et al., 2020) using vari- ant of the UL2 objective (Tay et al., 2022) with batch size 256, input length 4096, and output length 910. In particular, our mixture contains four objec- tives in equal proportion: prefix-LM with noise rate 0.5, and span corruption (Raffel et al., 2020) with noise rate 0.1.5 and average span lengths 3, 8, and 64. We use the Adafactor optimizer (Shazeer and Stern, 2018) with the 15.1.1 inverse square root learning rate schedule and no dropout. CoLT5 is trained with the T5X (Roberts al., 2022) frame- work. For pre-training, we route m 512 tokens,
Table 4: Median and 90th percentile input length by dataset measured in SentencePiece tokens.
Timing We report time per sample per TPUv4 chip, as measured by xprof (Google, 2020). For inference we use single TPUv4 with batch size 16 or the largest that fits in memory. For fine-tuning profile with 8 TPUv4 chips, sharded separately for each model maximize throughput.
g th of the input length.
4.2 Main results
Fine-tuning For fine-tuning we constant learning rate of 0.001, batch size 128, and dropout rate 0.1 for all tasks. Main results use input length 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, To th of the input length. We train until convergence
Figure 2 compares the quality-speed trade-off for L and CoLT5, showing that CoLT5 is better at any speed. For 16k input length, CoLT5 matches exceeds LONGT5 quality for Large and XL with 35-75% training speedup and 50-100% in- ference speedup on top of the order-of-magnitucc inference speedup from MQA. Encoder speedups are even greater (Appendix D). CoLT5-XL. also achieves SOTA performance on the SCROLLS benchmark. Table 3 contains all main results.

---

Natural
TriviaQA
32
1k
64k
2k
30
32k
4k
F
28
32k
8k
16k
26
16k
16k
24
LONGTS
0.1
0.2 0.3
0.05
0.1
8k
CoLT5
Figure 5: COLTS can use its long-input capability to benefit from more shots for in-context learning. Few-shot exact match for CoLT5-Large on Natural Questions and TriviaQA dev sets as function of in- put tokens, fitting as many examples as possible. Each example contains question, context, and answer. Inputs length used are 1024, 2048, 4096, 8192, 16384.
100
200
300
400
500
600
Time per sample (ms)
Figure 4: COLTS effectively scales to extremely long inputs, achieving stronger performance and faster speed than LONGT5. Fl on NarrativeQA as func- tion of inference time per sample for LONGTS and CoLTS Large models using varying input lengths.
We test the above hypothesis by evaluating few-shot learning performance on Natural Ques- tions (Kwiatkowski et al., 2019) and TriviaQA as function of input length, using as many exam- ples fit in the context. We consider the open book setting, such that each example consists of question, context document, and answer. Table 5 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 the input length. We found that CoLT5 can perform in-context learning only to the input length it was trained on, so for these experiments we continued pre-training COLT5-Large model on input length 16384 for another 100k steps. For the same reason route m 512 tokens in pre-training.
4.3 Scaling to extremely long inputs
We hypothesize that the advantage of CoLTS over LONGT5 strengthens with input length, as the frac- tion important tokens decreases and CoLT5 can route a greater proportion of important tokens to the heavy branch. Figure 4 compares the quality- speed rrf for LONGT5 and CoLT5 on Nar- rativeQA, sweeping over input length rather than model size. The number of routed tokens is 16 th of the input length, except that do not increase routed tokens going from 32k to 64k, so at 64k we route only nd of the input length. COLT5 achieves both stronger performance and faster in- ference speed at all input lengths and is able to effectively make use of extremely long inputs. We note that CoLTS 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.
Figure 5 displays CoLT5 few-shot performance as a function input length, showing that CoLT5 able to apply its long-inpun capabilities to extract information from increasing numbers of examples.
4.4 In-context learning
Dataset
1024
2048
4096
8192
16384
Models trained on the UL2 objective have shown strong few-shot in-context learning (ICL. capa- bilities even smaller sizes (Tay al., 2022). COLT5 enables tractable inference with long in- puts. Here, we leverage this for scaling the number of examples used for in-context learning.
NQ
0.1
0.7
1.7
3.4
5.6
TriviaQA
1.6
2.3
3.8
7.0
9.8
Table 5: Average number of Natural Questions and TriviaQA few-shot examples that fit in input length.
4.5 Ablations
*Note that LONGT5 does not use MQA. we add MQA LONGT5 for conservative baseline.
profiling
This section studies the effect different choices in the CoLT5 recipe. Table 6 contains results of a series experiments that change a single compo-
"We initially evaluated ICL for models pre-trained with PEGASUS but found performance to he nearly O.

---

Avg
Inf
TQA
NQA
QAS
QuAL
CNLI
SumS
QMS
GovR
Ablation
Model
S/s
F1
F1
F1
EM
EM
Rgm
Rgm
Rgn
Rgm
Baseline
CoLT5-B
42.5
11.2
82.4
23.1
38.3
36.6
87.8
35.3
19.3
20.5
39.4
Static
40.5
11.6
79.7
19.2
34.2
34.5
86.4
34,9
18.1
18.9
38.8
Routing
Share QKV
42.0
11.8
82.1
21.9
37.5
36.2
87.0
35,2
18.2
20.4
39.7
v=all
42.5
9.4
82.4
22.3
38.6
37.2
87.8
35.3
19.1
20.3
39.8
Attention
v=q
42.3
11.5
82.5
22.5
37.3
37.0
85.9
35.2
19.0
20.5
39.7
Routed Tokens
m=512
41.6
12.2
81.9
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
LONGT5-B 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
11.2
82.6
22.6
40.5
37.3
87.3
35.3
19.6
20.8
39.6
Table 6: CoLT5 ablations evaluated on validation sets. Each experiment modifies a component of the COLT5 recipe for CoLT5-Base Static routing divides the imput 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 entire imput, while v=q selects the same number of key and value tokens as query tokens. m=512 and m=1536 different numbers of routed tokens. LONGT5-B uses LONGT5 encoder while retaining other parts of the CoLT5 training recipe such MQA and the UL2 objective. Multi-head refers to using multi-head cross-attention. The final ablation replaces the UL.2 objective with PEGASUS as LONGT5.
nent for COLT5 Base.
0.6
Other
Routed proportion
Routing First, note that static routing evenly distributing routed tokens over the input leads to massive drop in performance. The impor- tance routing provides evidence that the model learns to devote capacity to important tokens and the advantage of COLTS 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 modest reduction in quality and increase in speed.
Answer
Question
0.4
0.2
MLP
Query
KV
Figure 6: Proportion of tokens routed for answer (string match), question, and other tokens by routing compo- nent for CoLTS Large model, averaged over examples in TriviaQA dev set and all layers of model.
The optimal number routed tokens represents a trade-off between improved performance and computational cost of applying heavier layers. Ta- ble 6 shows strong gains going from 512 to 1024 (baseline) routed tokens and diminishing returns for further increases.
Attention CoLT5 relies on routing identify not only tokens that can benefit from important in- formation elsewhere in the input, but also which tokens contain such important information. We study whether CoLT5 is successful in this task by comparing performance with two different atten- tion settings v=all, in which routed tokens attend to the entire input, and v=g. which equal num- ber of routed keys and values queries, rather than twice many. CoLT5 appears to occupy sweet spot, as using fewer routed key-values modestly de- creases performance at similar speed but attending to all inputs barely helps at sharply increased cost.
Other We compare CoLTS LONGTS with multi-query cross-attention. confirming that LoNGT5 indeed does not achieve an unexpected quality gain from MQA, and our conservative assumptions in Figures 2, 4 are valid. Next, we evaluate multi-head cross-attention for CoLT5, finding that it lcads to modestly improved CoLTS performance. However, as MHA exhibits order- of-magnitude slower inference, MQA clearly favored. Finally, PEGASUS appears fine-tune slightly better than UL2, though the difference is small and UL2 enables few-shot learning.

---

4.6 Routing analysis
Acknowledgements
It is interesting to ask whether COLTS routed to- kens line up with what consider intuitively important tokens in each document. We investi- gate this question by studying routing patterns of a Large CoLT5 model fine-tuned on TriviaQA. We divide tokens into three categories: (1) question tokens, (2) answer tokens, and (3) other tokens. Figure 6 shows the average fraction of each type of token that is routed through the heavy path for MLP and attention layers TriviaQA. We note that question and answer tokens are significantly more likely to be routed than other tokens, for feedfor- ward as well as attention queries and keys/values. Appendix F presents more detailed routing analy- sis; e.g.. semantically important tokens are much more likely to be selected in later layers.
We would like to thank Srinadh Bhojanapalli, Luke Vilnis, Zachary Fisher, Jianmo Ni, Tal Schuster, Vaclav Cvicek, Sudeep Gandhe, Bhargav Kanagal, Kenton Lee, Ming-Wei Chang, Afroz Mohiuddin, Raphael Hoffmann, and others Google Research for helpful advice and discussion.
References
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Va- clav Cvicek, Zachary Fisher, Philip Pham, Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. 2020. ETC: Encoding long and structured inputs in transformers. arXiv preprint
Ttt::200.08488.
Iz Beltagy, Mathew E Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. arXiv preprint 2rXiv:2004.05/33
5 Conclusion
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson. Chris Leary. Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. 2018 JAX: composable transformations of Python+NumPy programs.
We propose CoLT5, new model for long-range inputs that employs conditional computation for higher quality and faster speed. CoLT5 has light feedforward and attention layers that apply to the entire input, as well as heavy branches that are ap- plied only to a subset of important tokens selected by a learned router. We show that CoLTS achieves stronger performance at any speed compared to LONGTS on variety of long-input datasets, and effectively and efficiently make use of ex- tremely long inputs up to 64k tokens.
Mingda Chen, Zewei Chu, Sam Wiseman, and Kevin Gimpel. 2022. SummScreen: A dataset for ab- stractive screenplay summarization. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 8602-80 1610. Dublin, Ireland. Association for Computational Linguistics.
Rewon Child, Scott Gray, Alec Radford, and Generating long
Ilya Sutskever. 2019. quences with sparse transformers. arXiv preprint arXiv: 10509.
Limitations
CoLT5 applies conditional computation only the encoder. Applying conditional computation in the decoder is more complicated; the routing method in CoLTS is not causal, so it isn't 't applicable when generating token by token. Since decoder-only models and applications with long outputs have become more popular recently, this is strong limi- tation of the current approach. Although the rout- ing method in CoLT5 could potentially be applied to the input context in a decoder-only model, we didn't investigate this setup.
Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Scokhwan Kim, Walter Chang, and Na- zli Goharian. 2018 A discourse-awer awaree attention model for abstractive summarization of long docu- ments. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Lunguage Tech- nologies, Volume (Short Papers), pages 615-621, New Orleans, Louisiana. Association for Computa- tional Linguistics.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A. Smith, and Matt Gardner. 2021. dataset of information-ssekknn questions and answers an- chored in research papers. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Hu- man Language Technologies, pages 4599-4610, On- line. Association for Computational Linguisties.
CoLT5 is specialized towards long sequences and has to be trained from scratch. For large-scale training and deployment, it is desirable either train single model that can handle both short and long sequences, or develop long-input architec- ture that can be adapted from an existing large model.
Michiel de Jong, Yury Zemlyanskiy, Joshua Ainslie, Nicholas FizzGerald, Sumit Sanghai, Fei Sha, and

---

William Cohen. 2022. FIDO: indio-ii-deeoder op- timized for stronger performance and faster infer- ence. arXiv preprint arXiv:2212.08155.
Yuta Koreeda and Christopher Manning. 2021. Con- tractNLI: A dataset for document-level natural lan- guage inference for contracts. In Findings of the Association for Computational Linguistics: EMNLP 2021, pages 1907-1919, Punta Cana, Dominican Re- public. Association for Computational Linguistics.
Jacob Devlin, Ming-Wor King-Wei Chang, Kenton Lee, and Kristina Toutanova. 2019. BERT: pre-training of deep bidirectional transformers for language under- standing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, NAACL-HLT 2019. Minneapolis, MN, USA, June 2-7, 2019, Volume (Long and Short Pa- pers). pages 4171 Association for Computa- tional Linguisties.
Bernhard Kratzwald and Stefan Feuerriegel. 2018. Adaptive document retrieval for deep question an- swering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, Brussels, Belgium, October 31 November 4, 2018, pages 576-581. Association for Computational Lin- guistics.
William Fedus, Barret Zoph, and Noam Shazeer. 2021. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.03961.
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Red- field, Michael Collins, Ankur P. Parikh, Chris AL- berti, Danielle Epstein, Illia Polosukhin, Jacob De- vlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wes -ei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019. Natural questions: benchmark for question answer- ing research. Trans. Assoc. Comput. Linguisties. 7:452-466
Google. 2020. Profile your model with cloud tpu tools. https://cloud. google. cwwc/o//////t comm/pp/docs// cloud-tpu-tools, Accessed: 2022-11-11.
Mandy Guo, Joshua Ainslie, David Uthus, Santiago Ontañón, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. 2022. LongT5: Efficient text-to-text trans- former for long sequences. In Findings of the Associ- ation for Computational Linguistics: NAACL 2022, pages 724-736, Seattle, United States. Association for Computational Linguistics.
Tao Lei, Junwen Bai, Siddhartha Brahma, Joshua Ainslie, Kenton Lee, Yanqi Zhou, Nan Du, Vin- cent Y. Zhao, Yuexin Wu, Bo Li. Yu Zhang, and Ming- Wei Chang. 2023. Conditional adapters: Parameter-efficient transfer learning with fast infer- ence. In Advances in Neural Information Processing Systems.
Jonathan Heek, Anselm Levskaya, Avital Oliver, Mar- vin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. 2020. Flax: A neural network library and ecosystem for JAX.
Yuning Mao, Pengeheng Hc. Xiaodong Liu, Ye- long Shen, Jianfeng Gao, Jiawei Han, and Weizhu Chen. 2021. Reader-guided passage reranking for open-domain question answering. In Findings of the Association for Computational Linguistics: ACL/IJCNP 2021, Online Event, August 1-6, 2021, volume ACL/UJCNLP 2021 of Findings of ACL, pages 344-350. Association for Computational Lin- guistics.
Luyang Huang, Shuyang Cao. Nikolaus Parulian, Heng Ji, and Lu Wang. 2021. Efficient attentions for long document summarization. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Hu- man Language Technologies, pages 1419-1436, On- line. Association for Computational I Linguistics.
Richard Yuanzhe Pang, Alicia Parrish, Nitish Joshi, Nikita Nangia, Jason Phang, Angelica Chen, Vishakh Padmakumar, Johnny Ma, Jana Thompson. He He, and Samuel R. Bowman. 2021. QuAL- ITY: Question answering with long input texts, yes! arXiv preprint arXh::21
Mandar Joshi, Funsol Choi, Daniel S. Weld, and Luke Zettlemoyer. 2017. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehen- sion. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics. Van- couver, Canada. Association for Computational Lin- guistics.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya. Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. 2022. Efficiently scaling transformer in- ference. arXiv preprint arxiv:2211.05102.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess. Rewon Child, Scott Gray. Alec Radford, Jeffrey Wu, and Dario Amodei. 2020. Scaling laws for neural language models. CoRR, abs/2001.0880000 108661.
Yujic Qian. Jinhyuk Lee, Sai Meher Karthik Duddu, Zhuyun Dai, Siddhartha Brahma, Iftekhar Naim, Tao Lei, and Vincent Y Zhao. 2022. Multi- vector retrieval as sparse alignment. arXiv preprint arXiv:2211.01267.
Tomáš Kočiský. Jonathan Schwarz, Phil Blunsom, Chris Dyer. Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. 2018. The NarrativeQA read- ing comprehension challenge. Transactions of the Association for Computational Linguisties, 6:317- 328.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang. Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2020. Exploring the limits

---

of transfer learning with a unified text-to-text trans- former. J. Mach. Learn. Res., 00:140:1-140:60
Neeraj Varshney, Man Luo, and Chitta Baral, 2022. Can open-domain QA reader utilize exter- nal knowledge efficiently like humans? CoRR, abs/2211.12700
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor, Sha- ran Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowycz, Alex Salcianu, Mare van Zee, Jacob Austin, Sebas- tian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jas- mijn Bastings, Jannis Bulian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee. Dan Garrette, James Lee- Thorp, Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Maitin- Shepard, Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov. Joshua Newlan, and Andrea Gesmundo. 2022. Scaling up models and data with t5x and seqio. arXiv preprint arXiv:2203.11 17199.
Shuohang Wang, Mo Yu, Xiaoxiao Guo, Zhiguo Wang, Tim Klinger, Wei Zhang, Shiyu Chang. Genry Tesauro, Bowen Zhou, and Jing Jiang. 2018. R²: Reinforeed ranker-reader for open-domain question answering. In Proceedings of the Thirty-Second AAAI Conference on Artificial Intelligence, (AAAI- 18), the 30th innovative Applications of Artificial In- telligence (IAAI-18), and the 8th AAAI Symposium Educational Advances in Artificial Intelligence (EAAI-18), New Orleans, Louisiana, USA, February 2-7, 2018, pages 1988-5988. AAAI Press.
Sinong Wang. Belinda Z Li. Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self- attention with linear complexity. arXiv preprint 047668.
Tal Schuster, Adam Fisch, Jai Gupta, Mostafa De- hghani, Dara Bahri, Vinh Q Tran, Yi Tay, and Don- ald Metzler. 2022. Confident adaptive language modeling, arXiv preprint arXiv:2207.0777 .77661.
Donghan Yu, Chenguang Zhu, Yuwei Fang. Wenhao Yu, Shuohang Wang. Yichong Xu, Xiang Ren, Yim- ing Yang, and Michael Zeng. 2022. Kg-fid: Infus- ing knowledge graph in fusion-in-decoder for open- domain question answering. In Proceedings of the 60th Annual Meeting of the Association for Compu- tational Linguistics ( (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, pages
Uri Shaham, Elad Segal, Maor Ivgi, Avia Efrat, Ori Yoran, Adi Haviv, Ankit Gupta, Wenhan Xiong, Mor Geva, Jonathan Berant, and Omer Levy. 2022. Scrolls: Standardized comparison over long lan- guage sequences. ArXiv, abs/2201.03555 .330533..
4961-4974, Association for Computational Linguis-
Noam Shazeer. 2019.
Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150.
tics.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago On- tağón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang, et al. 2020. Big bird: Transformers for longer sequences. Advances in Neural Information Processing Systems, 33: 13:17383-17297.
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis Quoe V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously large neural net- layer.
works: The sparsely-gater
mixture-of-exxpee In 5th International Conference on Learning Rep- resentations, ICLR 2017, Toulon, France, April 24- 26, 2017, Conference Track Proceedings. OpenRe- view.net.
Yury Zemlyanskiy, Joshua Ainslie, Michiel de Jong. Philip Pham, Ilya Eckstein, and Fei Sha. 2021. Readtwice: Reading very large documents with memories. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, pages 5189-5195.
Noam Shazeer and Mitchell Stern. 2018 Adafac- tor: Adaptive learning rates with sublinear memory cost. In Proceedings of the 35th International Con- ference on Machine Learning, ICML 2018 Slock- holmsmässan, Stockhoim, Sweden, July 10-15, 2018 volume 80 of Proceedings of Machine Learning Re- search, pages 4603 4611. PMLR.
Jingging Zhang, Yao Zhao, Mohammad Saleh, and Pe- ter Liu. 2020. Pegasus: Pre-training with extracted gap-sentences for abstractive summarization. In In- ternational Conference on Machine Learning, pages 11328-11339, PMLR.
Yi Tay, Mostafa Dehghani, Samira Abnar. Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. 2021. Long range arena : A benchmark for efficient trans- formers. In International Conference on Learning Representations.
Ming Zhong, Da Yin, Tao Yu, Ahmad Zaidi, Mutethia Mutuma, Rahul Jha, Ahmed Hassan Awadallah, Asli Celikyilmaz, Yang Liu, Xipeng Qiu, and Dragomir Radev. 2021. QMSum: A new benchmark for query- based multi-domain meeting summarization. In Pro- ceedings of the 2021 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 5905 Online, Association for Computational Linguistics.
Yi Tay, Mostafa Dehghani, Vinh Q Tran, Xavier Gar- cia, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby, and Donald Metzler. 2022. Unify- ing language learning paradigms. arXiv preprint orXivv2205.05131.

---

Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du. Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. 2022. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08906.

---

Model
Layers
Model dim
MLP light dim
MePy dim
HeadSight
HeadSpewyy
Params
LONGTS-B
768
2048
N/A
N/A
248m
CoLTS-B
12
1024
8096
433m
LONGTS-L.
24
1024
2816
N/A
16
N/A
783m
CoLTS-L
24
1024
1408
11264
4
12
1462m
LONGTS-XL
24
2048
5120
N/A
32
N/A
2850m
COLTS-XL
24
2048
2560
20480
8
24
5297m
Table 7: Hyperparameters for LONGT5 and CoLT5 models. T5.1.1 hyperparameters match LoNGT5, CoLT5
parameters are sparsely accessed as a result of conditional computation, so parameter counts do not reflect compute, and for given model size CoLT5 is in fact faster than LONGT5 despite having more parameters.
A Contributions
computation. Yi advised on the adaptation of UL2 to 4k input length pre- pre-training Finally, Yun-Hsuan and Sumit provided guidance and support for the project overall.
Joshua led the project, developed the initial con- ditional attention mechanisms, and conducted most experimental ablations. Tao developed the heavy/light formulation for heterogencous condi- tional computation, comprising the routing and con- ditional feedforward mechanisms, and iterated with Joshua on initial experiments demonstrating fea- sibility. Michiel helped scope the paper, per- formed 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 rout- ing for query and key/value tokens in the condi- tional 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 devel- oped the leaderboard comparisons with LongT5 and helped run several experiments. James advised and early comparisons with MoE conditional
B Model Hyperparameters
Table shows LONGTS and CoLT5 hyperparam- eters, including parameter counts. For LONGTS5, we report numbers for the TGlobal configuration, which match T5.1.1. Notice that CoLT5's parame- ter counts larger due to using conditional com- pute. Similar to other conditional compute archi- tectures such as mixture-of-experts, computational cost does not necessarily increase with parameter count.
We use the same 127-token local radius for CoLT5 as LONGT5. This results in local atten- tion window of 255 since 127 tokens are attended to the left and 127 to the right.
C Routing Normalization Hyperparameters
To normalize the routing scores for differentiable top-k k token selection, we use the iterative soft top- algorithm from Lei al. (2023) and Qian et al.
Model
Average
16k in, 128 out
16k in, 512 out
16k in, 1024 out
8k in, 128 out
Ene
Tot
Enc
Tot
Enc
Tot
Enc
Tot
Enc
Tot
LONGT5-B
77
136
84
98
84
165
84
296
27
39
CoLTS-B
29
90
30
45
30
113
30
256
18
30
LONGT5-L
164
329
173
222
179
392
179
799
66
100
COLTS-L
70
201
73
103
73
250
73
578
45
69
LONGT5-XL
412
557
1081
2065
166
COLT5-XL.
177
439
185
239
525
1253
115
Table 8: Comparison of total and encoder inference time per sample (ms) for LONGT5 and CoLT5 Base, Large, and XL models different input and output lengths. Average time per sample is computed 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, for 16k/512, and one each for 16k/1024 and 8k/128).

---

Model
arXiv
SummScreenFD
QMSum
GovRep
R-1
R-2
R-L
R-1
R-2
R-L
R-1
R-2
R-L
R-1
R-2
R-L
LONGTS-B
47.4
21.4
43.5
34.8
9.3
20.7
35.1
11.1
23.4
59.3
30.1
33.0
CoLT5-B
47.5
21.3
43.6
35.6
9.7
21.0
34.6
10.9
23.0
60.2
31.0
32.8
LONGTS-L
47.9
21.7
43.8
35.3
20.8
12.0
24.1
61.4
32.5
COI.S5-L
48.4
21.7
44.3
35.7
10.1
21.4
36.8
12.6
24.7
61.8
32.7
34.4
LONGTS-XL
48.2
21.8
44.1
36.6
10.3
21.5
37.0
12.5
24.7
61.8
33.2
34.8
COLTS-XL
48.4
22.0
44.3
36.3
10.0
21.5
37.4
13.0
25.1
62.2
33.3
34.9
Table 9: Full performance comparison with Rouge-1, Rouge-2, and Rouge-L. metries CoLTS and LONGT5 Base, Large, and XL models summarization dev sets. Results based on checkpoint that maximizes Rgm in Table
(2022) with 1.0 and 50 iterations. During training we allow the top k tokens have nonzero weight instead of just the top & in order provide a slightly improved training signal.
TPUv4 chips for Base, Large, and XL, respectively, with training time varying with dataset size.
Routing Analysis
Additional Experimental Results
In this section we take a closer look at the routing mechanisms in COLT5. There are three routing
Table 8 compares LONGT5 and CoLT5 inference speed in more detail, splitting off encoder and total time per sample. Since CoLT5 applies conditional computation only the encoder, encoder speed gains are larger than overall speed gain, and total speed gains are largest for shorter output length. Trade-offt are even more in the favor CoLT5 when paired with other decoder optimizations. Table 9 shows full (Rouge-1, Rouge-2, Rouge-L) results for summarization datasets.
processes in each layer of CoLT5: (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 sim- plicity, we will say that token is selected, when it routed to the heavy alternative (of either MLP 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 to- kens into three categories: (1) question tokens, (2) answer tokens (found via simple normalized string match of the ground truth answer), and (3) other to- kens. Figure 7 shows the proportion each token type that is routed by a fine-tuned CoLT5-Large model the TriviaQA dev set, by layer and rout- ing component.
E 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
Routing proportion
MLP
0.8
Query
KV
0.8
0.8
0.6
M
0.6
0.6
0.4
0.4
0.2
0.4
0.2
0.2
10
20
10
20
10
20
Layer
Layer
Layer
Other
Answer
Question
Figure 7: Proportion of tokens routed for answer (string match), question, and other tokens by routing component and layer for CoLT5 Large model, averaged over examples in TriviaQA dev set.

---

Saaviig leeng barrone
portunrral ae pirtannt open mage
Question
heavily routed to the expensive alternative last layers of the model.
madden pi? hovinatt bunkow OUght tone bantone bavolrrbrrtte S00d bortoon dnamint bartonn, bar noble bontons and the hast barkono butboy fiee use tore The howr smergett certonaae tonan,
secondi
frmetto driided
kiaeiin opendi: real
codendied
contayy soadly sacred porputoo musi, this slage was liccusity usad bogining encompaneing used deectbe soeage chor sitllmmeed
owesd voiass dnolucting
cantury taly
vivoce
seaglig rangge
boguus turtone
and yau canker jroeed many opentt
cwnkay hhrw radtodd thet
pOAdE your Ior wed bankbes carchccnne aod Chool sig botter ddu Da upper tagge biinpe nola End roice sharee prano pacsappio and necondd paodieggo with dommetc beer and inboeiiite fit and iespective
hance cocld taaned os linor berytr mod sport "polits pa (lincisande iSoude Scuuday bogo cor Biso locket oties sort EgE ylouice
The same tree for tokens around the correct answer ("pupageno" in this example).
Furho cieeaicc mandererd
serien Taery byre
sounding baribre
drainlic
(ycilipnng mozal) segggessntta
Jain dos "yrappers mompesi thomes adso maree horiia cplccitoo puott Migees, this tatter seaile gpooe (soii)) "oontiee caimee 1prorgss hore "op the ander land parss oogland handd cuatton the knndle burnor baatonn medal voice Ear3ang both bec and andiatattii thr nodie with whoooodooooooo lopoo comems uang beice oe gabows midde jA0 not pownrt
Figure 8: Visualization of token routing weights for some fragments of an example on TriviaQA.
Earlier we showed that question and answer to- kens are more likely to be selected, but separat- ing routing decisions by layer reveals interesting patterns. At early layers question and answer to- kens 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 im- ply that different layers focus on different types of tokens, or that some routing components do not successfully learn to identify important tokens.
routing with Q-routing, and K KV-routing with Q- routing). We show numbers for both the pre-trained checkpoint, as well as a fine-tuned model Trivi- aQA. As 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.
see, the routing of keys/values and
To gain better insight into this, Figure 8 vi- sualizes routing on two sample fragments from TriviaQA example (notice that, given the large in- put length used in COLT5, we do not show the complete example in the figure). The two frag- ments 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 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 10 shows the Pearson correlation coefficient be- tween the routing weights of the different routing mechanisms in each layer in a CoLT5 Large model MMLP-routing correlation with KV-routing, MLP.

---

Pre-trained
Fine-tuned
MLP-KV
KV-Q
MLP-KV
M KV-Q
-0,06
-0.06
-009
-0,06
-0,09
-0,26
0.27
0.52
0.04
0.27
0.39
0.02
-0,05
-0,03
0.75
0.05
-0.01
0.69
0.05
0.09
0.76
0.18
0.14
0.72
0.02
-0.01
0.75
0.22
0.26
0.68
0.02
-0.01
0.78
0.31
0.33
0.70
0.02
0.00
0.73
0.26
0.27
0.70
0.00
-0,02
0.44
0.11
-0,07
0.29
0.13
0.11
0.74
0.36
0.40
0.70
10
-0,06
-0108
0.08
-0,15
-0,15
0.12
11 0,05
-0.07
0.31
-0,08
-003
0.18
12
-0
0.04
-0,08
0.27
0.03
0.00
0.28
13
O IO
0,09
0.87
-0,13
-0033
0.72
14
-004
-0,05
0. 76
-0,66
-0.12
0.67
15
0.53
0.64
0.69
0.51
0.55
0.67
16 0.08
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
Table 10: Pearson correlation coefficient between the routing weights of the different routing mechanisms in each layer in CoLT5 Large model. We show num- bers for both the pre-trained checkpoint, as well as a fine-tuned model TriviaQA. Blue bars visualize pos- itive correlation, whereas red bars visualize negative correlation.


Reviews:
[{"review_id": "DdP7Ozq5lF", "reviewer": "Reviewer_tPfd", "summary": "", "questions": "A. Apologies if I’ve missed that part but after going through the pre-training section (page 6) I realized the dataset on which CoLT5 is pre-trained isn’t mentioned.\n\nB. Reproducibility - Numbers of CoLT5 -B in Table 6 and Table 3 are different. I wonder why that is happening. QAS F1 score is 42.1 v/s 38.3 in Table 3 and Table 6 respectively. Same for QMS and GovR Rgm scores have a difference of almost 5% which is a bit high. Let me know if I’m missing something that’s different in both these tables. \n\nC. Are all these experiments done on multiple seeds because I couldn’t see that anywhere.\n\nD. In Table 5, when you say the number of Natural Questions or TriviaQA, the number stated is in decimals like 0.1. Is it % or in hundreds, thousands? It will be great if you can clear that.\n\nE. Why didn’t you pre-train CoLT5 on input length equivalent to that of the dataset with the highest median input length in Table 4? Was it because there were no significant gains after a certain input length or did it take too long?\n\nF. I can’t see the number of parameters for different versions of CoLT5. The paper only mentions that it has more overall parameters as compared to LongT5 but uses less compute (Update: The concerned numbers are mentioned in Appendix, I’d recommend moving them in the main paper. Not the whole table but the parameter count)\n", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thorough review and engaging questions!\n\nWe agree that CoLT5 evaluation on multilingual tasks would be very interesting too but unfortunately could not find an extensive suite of multilingual long input tasks like SCROLLS.\n\nQuestion A:  \nGreat question.  We pre-trained on the C4 dataset used by T5 (and LongT5).  We’ll be sure to mention this in the pre-training section.\n\nQuestion B:  \nThank you for the careful comparisons.  The explanation for the difference is that Table 3 uses official SCROLLS leaderboard test set results, while all ablations in Table 6 use validation sets instead.  Most of the SCROLLS datasets have a relatively modest number of training examples (10k or less), so we saw that the CoLT5 Base model “overfit” slightly with higher metrics on the validation set than the leaderboard test set for QMS and GovR as you noticed.  Interestingly we noticed less such “overfitting” for CoLT5 Large and even less for CoLT5 XL.  On the other hand, for the QAS dataset we found the validation metrics to be noticeably lower than leaderboard metrics in general, and we suspect this is due to the QAS test set having more examples with multiple references (98% in test vs. 74% in validation as reported in the original Qasper paper).\n\nQuestion C:  \nWe unfortunately lacked the resources to repeat fine-tuning for each dataset using multiple seeds and relied more on averaging _across_ the 9 datasets to compare overall performance with less noise.\n\nQuestion D:  \nWe’ll clarify the Table 5 caption to highlight that it’s showing the average number of few-shot examples that fit in context along with the example we’re answering the question for.  The 0.1 value for NQ at 1024 length for example indicates that very few examples (e.g. 1 in 10) have space to fit even 1 shot, so effectively 1024 length degenerates to zero-shot for most NQ examples.\n\nQuestion E:  \nFor comparison with LongT5, we wanted to make the pre-training recipe as similar as possible, so we stuck with LongT5’s 4096 input length for pre-training.  (Unlike LongT5, we did adopt the UL2 task, but this was necessary to yield any interesting few-shot metrics.)\n\nQuestion F:  \nIn case of acceptance we’ll have some extra space we can use to see if there’s a natural spot to add this in the main paper.\n\nThank you also for all the helpful suggestions.  We’ll incorporate the grammatical clarifications and take everything into consideration in a final draft in case of acceptance.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "* Motivation and intuition behind the paper are explained really well.\n* Changes to the existing architecture of LongT5 are quite simple and novel.\n* CoLT5 is not just doing well when it comes to numbers but also efficient as compared to the baseline mentioned.\n* Conditional computing section is explained thoroughly except for Conditional Attention sub-section\n* Utilizes SCROLLS benchmark precisely made for the task at hand covering summarization, question answering, and natural language inference tasks.\n* Promising performance with extremely long inputs (~64k tokens). This is one of the main highlights of the paper for me.\n* Ablation study is done quite well serving as a foundation for upcoming models dealing with large input length\n", "reasons_to_reject": "* Results look great, however, SCROLLS is limited to English only. Testing on benchmarks focusing on other languages as well could have been helpful.\n* Selection of light and heavy hidden dimensions (rL, rH) looks a bit hacky especially for calculating FLOPS for the CoLT5 attention layer  (rL=¼, rH=¾). An ideal scenario would have been to show trade off with different hidden dimensions, corresponding parameter count and show its effect on downstream tasks (apologies if it’s already mentioned anywhere)\n", "questions_to_authors": "A. Apologies if I’ve missed that part but after going through the pre-training section (page 6) I realized the dataset on which CoLT5 is pre-trained isn’t mentioned.\n\nB. Reproducibility - Numbers of CoLT5 -B in Table 6 and Table 3 are different. I wonder why that is happening. QAS F1 score is 42.1 v/s 38.3 in Table 3 and Table 6 respectively. Same for QMS and GovR Rgm scores have a difference of almost 5% which is a bit high. Let me know if I’m missing something that’s different in both these tables. \n\nC. Are all these experiments done on multiple seeds because I couldn’t see that anywhere.\n\nD. In Table 5, when you say the number of Natural Questions or TriviaQA, the number stated is in decimals like 0.1. Is it % or in hundreds, thousands? It will be great if you can clear that.\n\nE. Why didn’t you pre-train CoLT5 on input length equivalent to that of the dataset with the highest median input length in Table 4? Was it because there were no significant gains after a certain input length or did it take too long?\n\nF. I can’t see the number of parameters for different versions of CoLT5. The paper only mentions that it has more overall parameters as compared to LongT5 but uses less compute (Update: The concerned numbers are mentioned in Appendix, I’d recommend moving them in the main paper. Not the whole table but the parameter count)\n", "paper_topic_and_main_contributions": "* This paper deals with the problem of handling longer inputs with Transformers. \n* Problem arises not only due to quadratic attention complexity but also because of  feedforward and attention layer projection to every token\n* Authors build on top of LongT5 and suggest to divide each feedforward and attention layer into light branch and heavy branch and perform routing of tokens.\n* Consists of Light branch and heavy branch\n    * Light branch : caters to all tokens, has lower hidden dimension, performs local attention, has fewer heads than LongT5\n    * Heavy branch : caters to important tokens, has higher hidden dimension, performs global attention, has more heads than LongT5\n* Authors suggest learning routing function using learned embeddings.\n* Apart from architectural change, authors also use a different objective for pre-training and add multi-query cross-attention to speed up inference\n* CoLT5 is faster when it comes to inference or fine-tuning as compared to LongT5. Does better on almost all the datasets as compared to LongT5 and achieves SOTA results on SCROLLS benchmark\n* As input length scales up, the model performance improves and it achieves faster speed than LongT5 (model scales well)\n", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "* One minor correction: Page 6, “SCROLLS contains question-answering datasets: Narrative QA …, NLI dataset: ContractNLI.., and summarization datasets: SummScreenFD.. ”. Minor punctuation adjustments that can improve clarity and readability.\n* “Feedforward and projection layer” and “feedforward and attention layer” have been used throughout the paper. I’d recommend sticking to one of those preferably feedforward and attention projection layers.\n* Figure 2 doesn’t mention the metric of average perf which is F1 score. Also, the caption should be: CoLT5 achieves faster performance than LongT5 for the same number of parameters instead of saying CoLT5 achieves stronger performance than LongT5 at any speed.\n* In Table 3 you have computed the average score under the Avg column, however, these metrics may not always be directly comparable, so averaging them might not accurately represent the overall performance. You should have averaged similar metrics like TQA, NQA, QAS under avg F1, QuAL, CNLI under avg EM, and arXiv, SumS, QMS, GovR under avg Rgm.\n* I feel the equation for calculating FLOPs (attn), the global projection part, can be explained a bit. It took me some time to understand that (this one is totally up to you, if you think readers can figure this out easily then don’t)\n"}}, {"review_id": "PY6rAyg2TB", "reviewer": "Reviewer_gfnU", "summary": "", "questions": "A. In the computation of FLOPs for the FFds, it's mentioned in lines 231 and 231 that r_L and r_H represent the dimension ratios relative to the standard T5. Considering this, it appears that the formula on line 235 should include squares for both r_H and r_L since they are applied to \"d.\" If that adjustment is made, the entire FLOPs formulation and final FLOPs formulas would likely change, although this change should not affect the experimentation results.\n\nB: Is it possible for this architecture to accommodate a more customized setting where each layer has its own ratio? For instance, the model could have larger local attention windows \"w\" and lower ratios \"m\" for the early layers, while the higher layers could have smaller local attention windows and higher important token ratios, considering their increased contextualization.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful review and questions!\n\nTo address your concern about the selected important token ratio, it is indeed possible to change this ratio during fine-tuning.  For example, for pre-training we used m=512 important tokens out of 4096 sequence length, yielding a ratio of 1/8.  On the other hand, for fine-tuning we used m=1024 important tokens by default out of 16384 sequence length for most tasks, yielding a smaller 1/16 ratio.  For our length scaling experiment for NarrativeQA, we actually found we only needed m=2048 tokens for 64k (more precisely 65,536) input length, yielding a ratio of 1/32.  We’ll clarify this in the text.  We agree it’s quite nice to be able to change this ratio depending on the task and language as you mentioned!\n\nQuestion A:  \nGood question about dimensions.  The reason we do not square r_L or r_H is because we only apply this ratio to the feedforward “hidden dimensions”, not the original “model dimension”.  For example, for T5 or LongT5 Base, the model dimension is 768 while the feedforward hidden dimension is 2048.  For CoLT5 Base we keep the same 768 model dimension and use 1024 and 8096 for light and heavy feedforward hidden dimensions, respectively.  Table 7 in the appendix has more details for each model size, but perhaps we should move this up to the main paper.\n\nQuestion B:  \nGreat suggestion regarding different ratios for different layers.  The architecture can indeed accommodate this, and although we didn’t investigate such configurations in this paper, we think it’s an interesting future direction.  The routing analysis visualizations for TriviaQA did indicate that later layers may be doing more meaningful routing, so perhaps using relatively larger important token ratios at later layers would be beneficial as you suggest.\n\nSuggestion A:  \nThanks, we have some additional space we can use to elaborate a bit on UL2 and MQA in case of acceptance.\n\nSuggestion B:  \nPerformance was nearly 0 for few-shot using the PEGASUS pre-trained checkpoints, such that the metrics would not look interesting, but we can highlight this in the text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper has a clear and well-organized structure, with a concise statement of objectives and clear description of the methodology.\n- The idea of integrating both a light route and a heavy route using a dynamic routing mechanism is thoughtfully designed and appears to be an innovative approach, particularly for long-range sequence-to-sequence models.\n- The extensive range of experiments conducted across various NLP tasks convincingly demonstrates the model's superior performance.\n- The inclusion of a thorough ablation study and analysis on the routed tokens provides valuable insights into the design decisions and inner workings of the model.", "reasons_to_reject": "My primary concern is that the selected important tokens ratio (m) in this mechanism remains static after pretraining. In some cases, particularly when training a multilingual version of this model, the ratio of important tokens could be dependent on the task/language. I wonder if it is possible to adjust this ratio during finetuning?", "questions_to_authors": "A. In the computation of FLOPs for the FFds, it's mentioned in lines 231 and 231 that r_L and r_H represent the dimension ratios relative to the standard T5. Considering this, it appears that the formula on line 235 should include squares for both r_H and r_L since they are applied to \"d.\" If that adjustment is made, the entire FLOPs formulation and final FLOPs formulas would likely change, although this change should not affect the experimentation results.\n\nB: Is it possible for this architecture to accommodate a more customized setting where each layer has its own ratio? For instance, the model could have larger local attention windows \"w\" and lower ratios \"m\" for the early layers, while the higher layers could have smaller local attention windows and higher important token ratios, considering their increased contextualization.", "paper_topic_and_main_contributions": "The authors present a transformer model designed for handling long-input texts. They split the computational process into two branches: a light one that processes all tokens and a heavy one that focuses on important tokens selected by a trainable routing network. By incorporating multi-query cross-attention and pretraining with UL2, they showcase the model's performance and efficiency in both few-shot and fine-tuning scenarios.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {"Typos_Grammar_Style_and_Presentation_Improvements": "A. It would be an improvement if some concepts have a brief description, like the UL2 objective and MQA.\n\nB. If PEGASUS is weaker in a few-shot setup (according to lines 474-475) better demonstrate in some experimental results."}}, {"review_id": "bhgJFUIOjq", "reviewer": "Reviewer_gBi6", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful review!\n\nRegarding algorithmic complexity, the quadratic cost in sequence length n results from how we selected the number of routed tokens m as a function of n for simplicity of analysis.  For example, with fine-tuning we typically set m=1024 for sequence length n=16384 and therefore replace m with n/16.  In practice, when scaling to sequence length 64k for NarrativeQA for example, we actually found we could reduce m to n/32, so the algorithmic analysis that sets m as a constant fraction of n could be considered more of a “worst case” for further length scaling, and we’re optimistic that scaling can look better than this in practice, as often many practical long-context applications have an increasing fraction of less relevant tokens the longer the context window.\n\nRegarding training from scratch, this was unfortunately necessary due to the differences in the new light/heavy branches, but we’ve seen concurrent related work on conditional adapters (https://arxiv.org/abs/2304.04947) which we consider a promising future direction to “adapt” the light or heavy branch starting from an existing Transformer checkpoint.\n\nThank you for the nice suggestion about Table 6 also.  We’ll be sure to highlight the best scores for easier comparison.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "The authors carefully set up experiments to show how their approaches both improve on training and inference time and how each of the algorithmic improvements compares to the LongT5 model presented as the baseline. Training speed increases of 35-75% and inference speed increases of 50-100% are shown as well as improvements in output quality when increasing input token sizes from 16k to 64k.", "reasons_to_reject": "The improvements to the model architecture do increase speed and performance by using computational 'tricks', but do not address overall algorithm complexity (e.g. the transformer algorithm used still has overall quadric complexity.) However, the authors architecture choices do have concrete runtime improvements when considering potential real-world use cases.\n\nThe authors state this type of model needs to be trained from scratch, which may be prohibitive for some compute-constrained research organizations. This can be mitigated by providing a pretrained base model to serve as a starting point.", "questions_to_authors": "", "paper_topic_and_main_contributions": "The authors present a series of improvements to the T5 language model architecture, focusing on both increased input context and improvements to speed for both training and inference. Specifically, they extend the LongT5 model with a routing mechanism that examines only a portion of the input sequence deemed importance by a learnable score with a 'heavy' component with more network resources and a lower-dimensional 'light' component that attends to the entire input sequence. Multi-query cross-attention is used instead of standard multi-head attention to improve inference speed and the UL2  (span corruption, prefix, causal LM) training objective is used in place of LongT5's PEGASUS sentence reconstruction objective.", "excitement": 4, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "It would be nice to see the \"best\" scores for each column highlighted in Table 6 (similar to Table 3.)"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a conditional computational graph based on the importance of the tokens. Only a pre-set number of important tokens would be sent to heavy modules and thus the inference would be much faster. 

Good parts:

1. The idea of using a light route and a heavy route with a dynamic routing mechanism is novel and interesting.

2. Experiments are solid to justify the claim.

3. The introduction motivates well.

Few potential issues:

1. Most experiments are performed on TPUs which limits its accessibility to mostly Googlers.

2. Description of method is mixed with many FLOP analysis, which is hard to digest. It's not very easy to understand what's the main proposed components and leave the FLOP analysis a separate session.
Decision: Accept-Main