TransformerFAM: Feedback attention is working memory

📄 arXiv: 2404.09173v3 📥 PDF

作者: Dongseong Hwang, Weiran Wang, Zhuoyuan Huo, Khe Chai Sim, Pedro Moreno Mengibar

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-04-14 (更新: 2024-05-07)

备注: 26 pages, 12 figures, 14 tables

💡 一句话要点

提出Feedback Attention Memory以解决Transformer处理长序列问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 反馈注意力 工作记忆 长序列处理 Transformer 深度学习 大型语言模型 自然语言处理

📋 核心要点

  1. 现有Transformer模型在处理长序列时面临二次方注意力复杂度的挑战，限制了其应用范围。
  2. 本文提出的反馈注意力记忆（FAM）架构通过反馈机制使网络能够利用自身的潜在表示，形成工作记忆。
  3. 实验结果显示，TransformerFAM在不同规模的模型上（1B、8B和24B）显著提升了长上下文任务的性能。

📝 摘要（中文）

尽管Transformer在深度学习领域取得了革命性进展，但其二次方注意力复杂度限制了其处理无限长输入的能力。本文提出了一种新颖的Transformer架构——反馈注意力记忆（FAM），通过反馈回路使网络能够关注自身的潜在表示。这一设计促进了Transformer内部工作记忆的形成，使其能够处理无限长的序列。TransformerFAM不需要额外的权重，能够与预训练模型无缝集成。实验表明，TransformerFAM在各种模型规模（1B、8B和24B）上显著提升了Transformer在长上下文任务中的表现，展示了赋能大型语言模型（LLMs）处理无限长度序列的潜力。

🔬 方法详解

问题定义：本文旨在解决现有Transformer模型在处理长序列时的二次方注意力复杂度问题，导致其无法有效处理无限长输入的痛点。

核心思路：通过引入反馈注意力机制，FAM架构使得Transformer能够关注自身的潜在表示，从而形成工作记忆，提升对长序列的处理能力。

技术框架：TransformerFAM的整体架构包括输入嵌入层、反馈注意力模块和输出层。反馈注意力模块通过反馈回路连接潜在表示，增强模型的记忆能力。

关键创新：FAM的核心创新在于引入反馈机制，使得Transformer能够在不增加额外权重的情况下，利用自身的潜在表示进行长序列处理，这与传统的注意力机制有本质区别。

关键设计：在设计上，FAM不需要额外的参数设置，保持了与预训练模型的兼容性。损失函数和网络结构经过优化，以确保在长上下文任务中实现最佳性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TransformerFAM在长上下文任务中显著提升了性能，尤其是在1B、8B和24B模型规模上，均表现出优于传统Transformer的效果，展示了其在处理无限长度序列方面的潜力。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、长文本生成、对话系统等。通过提升Transformer在长序列处理上的能力，FAM能够为大型语言模型提供更强的支持，推动智能助手、自动翻译等技术的发展，具有重要的实际价值和未来影响。

📄 摘要（原文）

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.