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 \ANDGoogle Research Author contributions are outlined in Appendix A. Correspondence author:
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.
中文速览
长文档处理一直是Transformer模型的痛点：注意力机制随序列长度呈二次方增长，而前馈层和投影层又对每个词都做同等计算，导致计算量极为庞大。CoLT5（Conditional LongT5）的核心思路是"不是每个词都同等重要"——它在前馈层和注意力层各设一条轻量分支（处理全部词）和一条重量分支（只处理路由模块动态选出的少数重要词），让计算资源向真正关键的词集中，同时还引入多查询交叉注意力（multi-query cross-attention）加速推理并采用UL2预训练目标赋予模型上下文学习能力。实验表明，CoLT5在训练和推理速度上大幅领先LongT5，在长文档基准SCROLLS上达到最优性能，且能高效处理长达64k词的超长输入，随输入长度增加所需的"重点词"比例还能进一步降低。这项工作的意义在于，它以极低的额外参数代价，系统性地解决了长文档Transformer在注意力和前馈层两个维度上的计算瓶颈，为实用化的长文档理解提供了一条可扩展的路径。
原文 arXiv:2303.09752;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2303.09752v3