下一代Transformer：RetNet结构可视化及Vision RetNet展望

2023-08-16 10:08

©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 的发展产生深远的推动作用。也期待其能为计算机视觉等领域带来新的建模范式和进展。

更多阅读