©PaperWeekly 原创 · 作者 | 岳廷
微软和清华大学的研究人员在论文《Retentive Network: A Successor to Transformer for Large Language Models》中提出了一种新的模型架构 RetNet,同时实现 Transformer 模型的高效并行训练、低延迟高吞吐量的推理以及良好的建模性能,被誉为“打破不可能三角形”。
论文链接:
https://arxiv.org/pdf/2307.08621.pdf
代码链接:
https://github.com/microsoft/unilm/tree/master/retnet
本文将详细介绍 RetNet 的技术原理、实验结果以及未来在 CV 领域应用前景展望。
RetNet技术原理
RetNet 的核心创新在于提出了一种新的序列建模机制——多尺度保留(Multi-Scale Retention,MSR)其核心为 Retention 机制,而 Retention 机制又可以表示为并行形式、递归形式和分块递归形式三种计算范式。Retnet整体结构见下图:
1.1 保留(Retention)机制
Retention 机制是 RetNet 的关键组成部分。给定输入序列 ,, 是隐含层维度。
通过投影操作 ,将其映射到一维函数。再考虑一个序列建模的映射问题 。
将 通过 映射到 。为简单起见 , 表示 .
以循环迭代的方式建立映射关系:
其中 A 是对角矩阵, 是内容相关投影向量:
 是可学习参数,得到 。将 代入 ,公式 1 重新如下:
其中
是 Transformer 中著名的 xPos,用于相对位置编码。将 简化成标量,得到以下公式:
这个公式对于训练样本来说,很容易并行化。
RetNet 的表达能力主要来源于保留机制中状态向量的设计。相比仅使用元素积的操作,通过递归保留了高维的状态 ,从而不会损害模型的表示能力。同时,这个设计又使得并行和递归两种表示成为可能。
1.2 Retention机制的三种表示形式
Retention 机制从循环网络开始推导,但由于其设计的特殊性,可以有三种表示形式,三种表示形式之间的关系和在训练推理中如何切换使用,下面将进行详细说明:
1. 并行表示
这是最基础的表示形式,通过 position encoding 实现内容相关的全连接。可以高效并行计算,适合训练
2. 递归表示
这是通过状态迁移的递归方式进行建模,可以顺序地进行编码,适合推理
3. 分块递归表示
这是在并行表示和递归表示之间的折衷,对输入进行分块,每个块内部使用并行表示并行计算,而块与块之间使用递归表示传递信息。适合非常长的序列建模
对应伪代码如下:
三者之间的关系是:
  • 并行表示是最基础的形式,递归表示和分块递归表示都可以从中推导出来。
  • 递归表示专注于顺序建模,用于推理。
  • 分块递归表示融合了并行表示的并行计算和递归表示的递归连接,以兼顾计算速度和建模能力。
在训练和推理中,常见的使用方式是:
  • 训练阶段主要使用并行表示,以并行地处理训练数据,进行高效训练。
  • 推理阶段切换为递归表示,以降低内存和延迟,提高实际使用的效率。
  • 对于特别长的序列,可以在训练阶段采用分块递归表示,以平衡训练速度和建模能力。
所以三种表示可以根据实际需要,灵活切换使用,以发挥各自的优势。并行表示用于训练,递归表示用于推理,分块递归表示用于长序列建模。这种转换自如的设计是 RetNet 的关键创新之处。
1.3 门控多尺度保留模块(Gated Multi-Scale Retention) 
在原 Transformer Self-Attention 块中,Retnet 使用多头保留机制构建多尺度保留模块。具体来说,对输入 ,计算:
这里 表示第 i 个头对应的保留率。不同头使用不同的 实现了多尺度建模。
然后将所有头连接起来,并进行门控及投影:
实现了非线性门控。将所有头拼接后再进行线性投影以混合特征。
此外,本文还使用 GroupNorm 对不同 head 的输出进行层内标准化:
因为各个 head 参数不共享,方差统计特性不同,需要分头标准化。所有公式整理如下:
到此为止,就构建了多尺度保留模块。它充分发挥了保留机制的能力,组合了多尺度建模、门控非线性化、分头标准化等有效手段,成为一个高效且富有表达力的模块。
1.4 总体网络结构
最终,本文使用标准的 Transformer 式结构来构建整个 RetNet 模型:
其中堆叠 L 个 MSR 模块和 FFN 模块。LN 表示 LayerNorm 层,FFN 为前馈全连接网络。
这种结构充分融合了 Multi-Scale Retention 模块的创新与 Transformer 架构的优点。
需要说明的是,在训练过程中,本文使用保留机制的并行表示来进行高效的并行计算。而在部署过程中,则切换到递归表示,以实现低延迟和高吞吐量的推理。对于非常长的序列,本文还可以利用分块递归表示来进一步优化训练与建模。
上述转换自如的多表示能力,使得 RetNet 可以同时兼具 Transformer 的高效训练与递归网络的高效推理,这是其最大的创新之处。
1.5 RetNet实验结果
论文中在多个语言建模任务上进行了大量实验,与 Transformer 及其他高效变体进行比较。结果展现了 RetNet 在多个方面的优势。Retnet 的主要优势,是推理优势,因此本文主要关注其推理性能对比。
从下面的图中可以看到,在不同序列长度下,RetNet 的 GPU 内存占用、延迟和吞吐量都明显优于 Transformer。尤其是内存仅占 Transformer 的 30%,延迟降低 8 倍之多。这使其非常适合实际部署,在大模型成为主流的情况下这个数据,使得 retnet 非常有可能成为下一代 Transformer 替代品,因为 Retnet 大大降低了推理消耗
RetNet在CV任务中的应用前景
从本文的实验数据看,替代 Transformer 几乎板上定钉的事,但其在 CV 领域的应用前景如何?结合 Vision Transformer,展望 Vision Retnet,可以从以下几个方面探讨 RetNet 在视觉领域的应用前景:
1. 更好的全局注意力建模
图像中的物体通常具有整体的结构和形状,理解这些整体结构需要建模长程空间依赖。例如,识别一匹马需要结合头、身体、四肢等不同部分。但 Transformer 在较长的图像序列上表现较差。
RetNet 通过保存历史注意力分布明确地建模长程依赖。在处理图像时,可以将图像拆解为 Patch 序列,然后应用 RetNet。RetNet 会记忆各 Patch 之间的相关性,比如身体相关的 Patch。在处理后续 Patch 时可以利用这些记忆,从而建模整体结构。
2. 提高计算效率
传统的 CNN 和 Transformer 需要一次性处理整张图像,计算量非常大。
而 RetNet 可以将图像分割成小的 Patch 序列,然后递归地处理这个 Patch 序列。具体做法是: 
1. 将图像分割成 n*n 的小 Patch,形成序列(比如16*16) 
2. 对第一个 Patch,计算自注意力得到注意力分布 A1
3. 对第二个 Patch,同时参考上一步的 A1 和当前 Patch,计算新的注意力分布 A2
4. 如此递归处理所有 Patch,逐步构建完整的注意力分布
这种递归计算方式充分利用了前面的注意力分布,避免了重复计算。
相比一次性处理整张图像,这种 Patch 序列的递归计算可以显著降低计算量。根据 RetNet 论文,在相同配置下,计算量可以降低约 60%。
此外,这种 Patch 序列表示也带来了其他优势,比如可以自然地融合多尺度信息,以及增强对遮挡的鲁棒性。
总之,RetNet 通过图像 Patch 序列化可以实现分块递归计算,这是其提高计算效率的关键所在。这种设计充分发挥了 RetNet 的优势,值得在视觉任务中广泛应用。
3. 增强对遮挡的鲁棒性
图像中经常存在遮挡,遮挡的区域会提供错误的视觉信息。但 RetNet 通过记忆历史可以弥补这些错误,即使输入被遮挡,也可以用记忆进行推理。例如,处理一个被遮挡的人脸图像时,RetNet 可以回想起完整人脸的注意力分布,减少遮挡对识别的影响。
4. 支持基于注意力的图像生成
RetNet 的记忆机制不仅可以用于分析,也可以应用到生成任务中。在生成图像时,可以用记忆的注意力分布来引导生成更加连贯、自然的图片。例如生成连贯的文字图像时,可以用记忆的注意力分布来确定当前生成的内容与前文的关联性。
5. 与置信度相关的门控机制
对图像中不同区域可以赋予置信度。对于置信度高的区域,利用新的注意力分布;对于置信度低的区域,保留更多历史注意力信息。例如生成图像时,对已生成清晰区域利用当前分布;对仍不清晰的区域保留更多历史记忆,以生成更连贯的图像。
6. 多任务学习
不同视觉任务可以共享历史注意力记忆。例如同时进行分类、分割、检测等任务,记忆的注意力分布可以在任务间共享,实现多任务协同。
总之,RetNet 的工作记忆机制可以显著增强 Vision Transformer 对长程依赖的建模能力,并在计算效率、鲁棒性等方面带来优,有望成为视觉领域的重要工具。
总结
总体来说,RetNet 是一个非常有前途的工作。它对传统 Transformer 框架进行了巧妙而原创的改进,使之可以“兼收并蓄”,在多个维度上实现更优的 trade-off。这种融会贯通和创新的设计思想,为构建下一代高效建模模型提供了重要参考和启发。
随着模型和数据规模不断扩大,如何在表达能力、训练效率、部署效率之间找到最佳平衡,将是所有预训练语言模型和基础模型架构需要面对的核心问题。RetNet 给出了一个非常有前景的解决方案。相信其技术思想和创新将对未来 Transformer 的发展产生深远的推动作用。也期待其能为计算机视觉等领域带来新的建模范式和进展。
更多阅读
#投 稿 通 道#
 让你的文字被更多人看到 
如何才能让更多的优质内容以更短路径到达读者群体,缩短读者寻找优质内容的成本呢?答案就是:你不认识的人。
总有一些你不认识的人,知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁,促使不同背景、不同方向的学者和学术灵感相互碰撞,迸发出更多的可能性。 
PaperWeekly 鼓励高校实验室或个人,在我们的平台上分享各类优质内容,可以是最新论文解读,也可以是学术热点剖析科研心得竞赛经验讲解等。我们的目的只有一个,让知识真正流动起来。
📝 稿件基本要求:
• 文章确系个人原创作品,未曾在公开渠道发表,如为其他平台已发表或待发表的文章,请明确标注 
• 稿件建议以 markdown 格式撰写,文中配图以附件形式发送,要求图片清晰,无版权问题
• PaperWeekly 尊重原作者署名权,并将为每篇被采纳的原创首发稿件,提供业内具有竞争力稿酬,具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道:
• 投稿邮箱:[email protected] 
• 来稿请备注即时联系方式(微信),以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信(pwbot02)快速投稿,备注:姓名-投稿
△长按添加PaperWeekly小编
🔍
现在,在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
点击「关注」订阅我们的专栏吧
·
·
继续阅读
阅读原文