SXP 官网  

你的位置:SXP 官网 > GULF中文网 >

工作记忆涌现!Google全新升级反馈注意力机制FAM,解锁Transformer长序列处理能力|算法|key|上下文|transformer

发布日期:2025-01-22 13:39    点击次数:103

毫不夸张的说,Transformer架构改变了现代深度学习的整体格局,为人工智能领域激发出了诸多富有活力的智能应用。但是它的注意力计算复杂度呈平方级别,这一直被工业界和学术界所诟病。为了解决这一问题,解除Transformer无法处理无限长输入序列的限制,Google的研究人员提出了一种全新的Transformer架构,称为反馈注意力记忆网络(Feedback Attention Memory,FAM),FAM基于一种反馈循环机制,能够使网络关注自身的潜在表示,并且促进了Transformer中工作记忆的涌现,使其能够处理无限长的序列。值得关注的是,TransformerFAM是一种轻量级的架构改进,它不需要引入新的参数,就可以完美兼容现有的大型语言模型,并赋予它们处理长序列的新能力。本文作者在各种参数规模(1B、8B 和 24B)上的实验都表明,FAM可以显著提升模型在各种长上下文任务上的性能。论文题目: TransformerFAM: Feedback attention is working memory 论文链接: https://arxiv.org/abs/2404.09173一、引言Transformer架构强大的扩展能力展现出了一种类似于摩尔定律的发展趋势,这种扩展模式将模型尺寸的增加与性能提升联系起来,并且在文本、图像、音频和视频多种模态和领域中表现优异。但Transformer也存在重大缺陷,即处理长序列时会遇到计算和内存开销的瓶颈。其次,它会忘记注意力窗口之前上下文中的信息,这与理论上可以无限传播信息的LSTM不同。现有的方法在处理长序列时通常会引入滑动窗口注意力(Block Sliding Window Attention,BSWA)或稀疏注意力(Sparse Attention,SA)等技术,但这类方法会在序列生成过程中忽略超出有效感受野的信息,如上图(d)所示,Receptive Field表示全局感受野,而BSWA的感受野远远无法涵盖这个范围。为了解决上述问题,本文提出了一种名为 Feedback Attention Memory (FAM) 的新 Transformer 架构。FAM在BSWA的基础上引入了一种全新的反馈机制,使得 Transformer 能够同时关注当前输入和内部工作记忆的表示,从而解决处理长序列的问题。二、本文方法2.1 BSWA在引入TransformerFAM之前,本文作者首先阐述了Block Sliding Window Attention (BSWA)的工作机制,因为TransformerFAM是在BSWA的基础上构建的。鉴于标准Transformer使用全局自注意力机制会导致计算复杂度呈现平方级别,SWA(基础的滑动窗口注意力)的优化思想是将注意力限制在一个滑动窗口内,从而将复杂度降低到,其中是窗口大小,但SWA会丢失窗口之外的信息。BSWA在SWA的基础上做了改进,引入了两个超参数:块大小(block size):指定一个块中包含多少个token,本文设置为1024。记忆段数(memory segment):指定在当前块之前保留多少个块的key和value,允许当前块关注之前块的信息。BSWA的自注意力计算过程可以形式化表示如下:其中 表示当前块的索引, 表示记忆段数,首先拼接记忆段和当前块的key和value,随后再扩展后的key和value上计算注意力。可以看出,查询向量 不仅可以关注当前块,还可以关注之前的 个块,记忆段的引入扩大了感受野,但仍有限制,BSWA的具体算法流程如下所示。虽然BSWA已经可以帮助Transformer可以处理比窗口更长的序列,但其引入的记忆段仍有长度限制,当输入序列非常长时,模型难以保留所有的关键信息。2.2 Feedback Attention Memory (FAM)为了解决BSWA的记忆段限制,本文在其基础上引入了一种称为反馈注意力记忆的FAM机制,FAM中包含了多个关键步骤,首先通过设置集成式注意力(Integrated Attention)来同时处理输入序列的上下文和反馈注意力记忆,在对FAM进行更新时,作者设计了逐块更新 (Block-Wise Updates)和信息压缩 (Information Compression)策略,前者会在输入序列在块之间传递时对FAM进行更新,而后者对FAM进行更新时,会以前一个时间步的FAM作为条件对当前块的信息进行压缩更新。上图展示了FAM与BSWA之间的注意力计算对比,可以观察到,TransformerFAM的输入查询主要涉及当前块、内存段和先前的FAM(图中的绿色线条)。FAM的查询向量是动态生成的,主要复制于前一个时间步的FAM(图中的蓝色线条),随后对当前块进行压缩编码来更新FAM。编码后的FAM会被反馈送回模型,与下一步的输入一起参与注意力计算。这种反馈循环使得信息能够在近似无限的范围内压缩和传播,这可以被理解为工作记忆。TransformerFAM的详细算法过程可以表示如下:从中我们可以看出,与BSWA相比,TransformerFAM的注意力计算包括了两个并行分支:输入查询分支输入查询向量 不仅关注当前输入块的 key/value,还关注记忆段和前一时间步的 FAM 的 key/value,具体计算过程如下:这样一来,当前输入就可以融合之前时间步的全局上下文信息(由 FAM 表示)和局部窗口信息(由记忆段表示)。输出 因此包含了全局和局部的上下文信息。FAM 查询分支与此同时,前一时间步的 FAM 查询向量 也参与到自注意力计算中,关注当前输入块和自身,计算过程表示如下:从上式我们可以注意到,FAM 查询向量 是如何生成的:它是直接复制自前一时间步的 FAM 输出,这样可以保持 FAM 的动态性和连续性。通过注意力计算 ,当前输入信息被编码并更新到新的 FAM 状态 中。整体来看,FAM 查询首先将当前输入块的信息进行压缩编码,并基于先前的 FAM 状态来更新到新的 FAM 中。此外还需要考虑的是FAM的引入会不会对模型整体训练效率造成不利影响,这是因为Transformer架构本身相比RNN更擅长利用GPU的并行性,RNN对输入序列的感知具有很强的因果关系,而Transformer仅在输入层和下一层之间具有因果关系。作者提到,FAM的反馈机制并不会消除Transformer架构的优势,FAM的因果关系只存在于区块之间,FAM记忆的长度相比输入序列的长度要小的多(例如长度1024的输入序列仅需额外处理64长度的FAM),性能开销可以忽略不计。三、实验效果由于TransformerFAM并未引入新的权重参数,因此本文作者直接使用了现有的大型语言模型权重,通过微调(fine-tuning)的方式来评估TransformerFAM的性能表现。具体来说,作者使用了谷歌的Flan-PaLM模型[1]的1B,8B,24B版本,每个模型都执行了50,000次迭代微调,优化方法使用LoRA[2],只对Transformer的QKV和FFN层更新参数。3.1 评估任务为了全面评估TransformerFAM在处理长序列任务上的性能,作者选取了以下几类常用的大模型评测基准:1. PassKey Retrieval 任务PassKey Retrieval是一个检索关键信息的基准任务,模型需要从非常长的文本(最长260k tokens)中提取出最开始给定的密钥信息,下图展示了一个密钥信息的模板,作者使用该任务快速检验模型在长序列上的基本能力。2. 长文本理解任务 (Long Context Tasks)Long Context Tasks包括NarrativeQA,PG-19,ScrollsQasper,ScrollsQuality,XLSum等5项任务,需要模型理解长达数千甚至数十万tokens的上下文,并回答相关问题。这类任务更能检验模型压缩长序列信息的能力。3. GPT-3 Benchmark 任务为了进一步检测TransformerFAM是否会影响短序列的处理能力,作者在GPT-3提供的Ranking和生成任务上进行了评估,涉及的序列长度通常在2K tokens以内。3.2 实验结果PassKey Retrieval任务的实验效果如下图所示,MX表示BSWA内存段的数量,FAM 表示具有 0 个内存段的 TransformerFAM。TransformerFAM在最长260k tokens的上下文长度下,能够100%正确地从中识别出密钥信息。而TransformerBSWA的性能则会随着上下文长度的增加而显著下降。长文本理解任务的实验效果如下图所示,TransformerFAM在所有长文本理解任务上都显著优于TransformerBSWA。尤其是在ScrollsQasper和NarrativeQA这两个需要处理数千甚至数十万tokens上下文的任务中,TransformerFAM展现出了极大的优势。此外,作者还对模型在8B和24B等不同参数规模下的性能表现进行了对比,如下表所示,结果显示TransformerFAM的优势会随着模型规模的增长而扩大,说明该架构具有一定的可扩展性。尽管TransformerFAM的主要目标是处理长序列,但在GPT-3 Benchmark上的实验结果表明,它在短序列任务上的表现也相当出色。如下表所示,与TransformerBSWA相比,TransformerFAM在GPT-3 Ranking和生成任务上均有提升。作者分析,这可能是由于TransformerFAM通过反馈注意力机制更高效地利用了上下文表示,从而获得了更优的全局表征能力。为了对FAM的内部工作机制进行探索,作者设计了一个可视化实验。具体来说,作者可视化了一个1B FAM模型每层每个头的注意力图,如下图所示。其中,沿左边缘的亮点表示参与到 FAM 的块输入,而沿下边缘的亮点表示FAM对相应块的信息进行压缩,总体来看,块输入的信息主要基于FAM,而FAM会对相应位置的块信息进行选择性压缩。这表明,在处理序列的初始阶段,FAM主要存储与输入序列无关的通用语义知识。但随着处理的进行,FAM会逐步聚焦于输入序列的具体语义,编码更多与之相关的信息。四、总结在本文的结尾,作者致敬了诺兰导演的电影《记忆碎片》,片中的主角患有顺行性遗忘症,这意味着他不记得过去10分钟内发生的任何事情,但他的长期记忆完好无损,他必须将重要信息记在纸上或者直接纹在身体上。本文作者认为目前的LLMs也类似于这种状态,根据缩放定律,我们有理由相信,LLMs能够记住整个互联网上的信息,但它们的短期记忆却受到了注意力窗口的限制。本文引入的反馈注意力机制FAM是解决这一限制的一个很好的尝试,FAM允许模型同时关注当前输入和内部工作记忆的表示,从而实现更加高效的长上下文推理效果。参考资料[1] Chung, Hyung Won u.a.(2022): Scaling instruction-finetuned language models.[2] Hu, Edward J / Shen, Yelong / Wallis, Phillip / Allen Zhu, Zeyuan / Li, Yuanzhi / Wang, Shean / Wang, Lu / Chen, Weizhu(2021): Lora: Low-rank adaptation of large language models.llustration From IconScout By WOOBRO LTD-The End-扫码观看!本周上新!“AI技术流”原创投稿计划TechBeat是由将门创投建立的AI学习社区(www.techbeat.net) 。 社区上线500+期talk视频,3000+篇技术干货文章,方向覆盖CV/NLP/ML/Robotis等;每月定期举办顶会及其他线上交流活动,不定期举办技术人线下聚会交流活动。我们正在努力成为AI人才喜爱的高质量、知识型交流平台,希望为AI人才打造更专业的服务和体验,加速并陪伴其成长。投稿内容// 最新技术解读/系统性知识分享 //// 前沿资讯解说/心得经历讲述 //投稿须知稿件需要为原创文章,并标明作者信息。我们会选择部分在深度技术解析及科研心得方向,对用户启发更大的文章,做原创性内容奖励投稿方式发送邮件到[email protected]或添加工作人员微信(chemn493)投稿,沟通投稿详情;还可以关注“将门创投”公众号,后台回复“投稿”二字,获得投稿说明。关于我“门”将门是一家以专注于数智核心科技领域的新型创投机构,也是北京市标杆型孵化器。 公司致力于通过连接技术与商业,发掘和培育具有全球影响力的科技创新企业,推动企业创新发展与产业升级。将门成立于2015年底,创始团队由微软创投在中国的创始团队原班人马构建而成,曾为微软优选和深度孵化了126家创新的技术型创业公司。如果您是技术领域的初创企业,不仅想获得投资,还希望获得一系列持续性、有价值的投后服务,欢迎发送或者推荐项目给我“门”:[email protected]点击右上角,把文章分享到朋友圈

Powered by SXP 官网 @2013-2022 RSS地图 HTML地图

Copyright Powered by365站群 © 2013-2024