C o LT 5: Faster Long-Range Transformers with Conditional Computation (2023)

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.

## introduction many natural language processing tasks, such as summarization @xcite or question answering over long documents @xcite , require machine learning models to encode longform text. processing long documents with a transformer model is computationally expensive, both because attention cost scales quadratically with input length and because feedforward and attention 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 @xcite @xcite @xcite @xcite . 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 colt5 (conditional longt5), a new family of models that, building on top of longt5 @xcite , enables fast processing of long inputs by combining architecture improvements for both attention and feedforward layers. colt5 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. in particular, colt5 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 longt5 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 1 provides an overview of the colt5 conditional mechanism. finally, colt5 also includes two other modifications to the longt5 architecture. colt5 adds multi-query cross-attention @xcite , significantly speeding up inference. colt5 also employs the ul2 @xcite pre-training objective, which we demonstrate allows for in-context learning over long inputs. we show that colt5 performs much faster finetuning and inference with similar or better model quality, improving over longt5 on arxiv summarization @xcite and triviaqa question answering @xcite datasets and achieving sota on the scrolls benchmark @xcite . moreover, colt5 achieves further gains in quality and speed for tasks with extremely long inputs (64k tokens), with less-than-linear scaling of "focus" tokens. ## conclusion we propose colt5, a 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 applied only to a subset of important tokens selected by a learned router. we show that colt5 achieves stronger performance at any speed compared to longt5 on a variety of long-input datasets, and can effectively and efficiently make use of extremely long inputs up to 64k tokens.