TransformerFAM: Feedback attention is working memory
Dongseong Hwang Weiran Wang Zhuoyuan Huo Khe Chai Sim Pedro Mengibar Google LLC Mountain View, CA, USA
Abstract
While Transformers have revolutionized deep learning, their quadratic attention complexity hinders their ability to process infinitely long inputs. We propose Feedback Attention Memory (FAM), a novel Transformer architecture that leverages a feedback loop to enable the network to attend to its own latent representations. This design fosters the emergence of working memory within the Transformer, allowing it to process indefinitely long sequences. TransformerFAM requires no additional weights, enabling seamless integration with pre-trained models. Our experiments show that TransformerFAM significantly improves Transformer performance on long-context tasks across various model sizes (1B, 8B, and 24B). These results showcase the potential to empower Large Language Models (LLMs) to process sequences of unlimited length.
中文速览
传统Transformer因自注意力机制的平方级复杂度，无法高效处理超长序列，且滑动窗口注意力虽能缓解内存压力，却因感受野有限而丢失远端上下文信息。为此，研究者提出了TransformerFAM架构，其核心创新是在每个Transformer块中引入一组"反馈注意力记忆"（Feedback Attention Memory，FAM）单元，让模型在处理当前块时同时关注输入序列和自身上一步的隐状态，形成一个类似工作记忆的反馈回路——当前块信息被压缩进FAM并传递给下一块，使全局上下文得以跨块持续传播。该设计无需引入任何新权重，可直接复用预训练模型并通过LoRA微调激活，推理时计算复杂度为O(L)、内存复杂度为O(1)，理论上支持无限长序列。在1B、8B和24B规模的Flan-PaLM模型上，仅用5万步微调，TransformerFAM在多项长上下文任务中均显著优于基线，为大语言模型处理无限长输入提供了一条无需修改模型参数规模、即可落地的实用路径。
原文 arXiv:2404.09173;中英对照 + 大白话阅读 https://aha.fim.ai/paper/2404.09173v3