Scaling Transformer to 1M tokens and beyond with RMT
Aydar Bulatov1, Yuri Kuratov2,1, Yermek Kapushev2, Mikhail Burtsev3
Abstract
A major limitation for the broader scope of problems solvable by transformers is the quadratic scaling of computational complexity with input size. In this study, we investigate the recurrent memory augmentation of pre-trained transformer models to extend input context length while linearly scaling compute. Our approach demonstrates the capability to store information in memory for sequences of up to an unprecedented two million tokens while maintaining high retrieval accuracy. Experiments with language modeling tasks show perplexity improvement as the number of processed input segments increases. These results underscore the effectiveness of our method, which has significant potential to enhance long-term dependency handling in natural language understanding and generation tasks, as well as enable large-scale context processing for memory-intensive applications.
中文速览
大语言模型的自注意力机制会随输入长度平方级增长计算量，严重限制了它处理超长文本的能力。为此，研究者提出了循环记忆变换器（Recurrent Memory Transformer，RMT），通过在输入序列中加入少量可训练的"记忆令牌"，并让这些令牌的输出在文本分段顺序处理时像接力棒一样传递下去，从而在完全不修改预训练模型内部结构的前提下，将计算复杂度从平方级压缩到线性级。实验中，研究者将RMT套在BERT和GPT-2等预训练模型外层，并设计了一套可扩展到超长序列的事实记忆与推理合成任务加以验证；结果显示，RMT能在长达两百万个词元（token）的序列上准确检索和存储关键信息，刷新了已知深度神经网络所能处理的最长序列纪录，且理论上没有进一步扩展的技术障碍。这项工作的意义在于：它让现有预训练大模型以"插件式"方式获得近乎无限的长程记忆能力，同时大幅降低推理时的算力和显存消耗，为长文档理解、多轮对话及大规模上下文处理等应用打开了新的可能。
原文 arXiv:2304.11062;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2304.11062v2