一文理解RetNet-CSDN博客

作者：梁德澎 | 来源：GiantPandaCV

前言

微软研究院最近提出了一个新的 LLM 自回归基础架构 Retentive Networks （RetNet）[1,4]，该架构相对于 Transformer 架构的优势是同时具备:训练可并行、推理成本低和良好的性能，不可能三角。

论文中给出一个很形象的示意图，RetNet 在正中间表示同时具备三个优点，而其他的架构 Linear Transformer、Recurrent Network 和 Transformer 都只能同时具备其中两个有点。

接下来看一下论文给出的 RetNet 和 Transformer 的对比实验结果:

当输入序列长度增加的时候，RetNet 的 GPU 显存占用一直是稳定的和权值差不多，而 Transformer 则是和输入长度成正比。

首先看红色线和紫色线，都是输入长度在 8192 下，RetNet 和 Transformer 推理延时的对比。

可以看到当 batch size 增加的时候， RetNet 的推理延时也还是很稳定，而 Transformer 的推理延时则是和 batch size 成正比。

而 Transformer 即使是输入长度缩小到 1024 ，推理延时也还是比 RetNet 要高。

RetNet 架构解读

RetNet 架构和 Transformer 类似，也是堆叠层同样的模块，每个模块内部包含两个子模块：一个 multi-scale retention（MSR）和一个 feed-forward network (FFN)。

下面详细解读一下这个 retention 子模块。

首先给定一个输入序列：

其中表示序列的长度。然后输入序列首先经过 embedding 层得到词嵌入向量：

其中表示隐含层的维度。

Retention 机制

首先对给定输入词嵌入向量序列中的每个时间步的向量都乘以权值得到：

然后同样有类似 Transformer 架构的 Q 和 K 的投影：

其中是需要学习的权值。

接着假设现在有一个序列建模的问题，通过状态将映射为向量。首先来看论文中给出的映射方式定义：

其中是一个矩阵，表示时间步对应的投影则。同样表示时间步对应的投影。

那么上面公式中的计算公式是怎么得出来呢，下面详细解释一下，首先将展开：

其中表示单位矩阵(主对角线元素为1，其余元素为0的方阵)。然后我们假定为初始状态元素为全0的矩阵，则有：

再继续上述推导过程：

所以根据上述推导过程和条件归纳可得：

然后我们来看一下矩阵是什么，论文中定义了是一个可对角化的矩阵，具体定义为：

其中都是维的向量，是一个可逆矩阵，而要理解首先得复习一下欧拉公式 [2]：

其中表示任意实数，是自然对数的底数，是复数中的虚数单位，也可以表示为实部，虚部的一个复数，欧拉公式[2]建立了指数函数、三角函数和复数之间的桥梁。

而这里是一个维向量:

则也就是将向量元素两两一组表示分别表示为复数的实部和虚部：

然后就是一个对角矩阵，对角元素的值就对应将和转成复数向量相乘再将结果转回实数向量的结果。

关于复数向量相乘可以参考文章：

一文看懂 LLaMA 中的旋转式位置编码（Rotary Position Embedding）

现在我们知道了矩阵的构成就能得到：

这里因为是可逆矩阵则有性质

其中为单位矩阵，则将次方展开：

就是个矩阵相乘，中间相邻的都消掉了，所以可得：

然后我们回到计算的公式：

接着论文中提出把吸收进和也就是和分别用和替代当作学习的权值，那么可得:

接着将公式简化，将改为一个实数常量，那么可得：

在继续推导前，先来仔细看一下，借助欧拉公式展开：

然后复习一下三角函数的性质[3]：

则有：

转为复数形式表示就是：

刚好就对应的共轭

所以可得：

其中表示共轭转置操作。

Retention 的训练并行表示

首先回顾单个时间步的输出的计算公式如下：

而所有时间步的输出是可以并行计算的，用矩阵形式表达如下：

其中，而表示两个矩阵逐元素相乘，和每一行对应一个时间步的 q 和 k 向量。

而每一行对应向量。就是对应矩阵的共轭，也就是将矩阵每一行改为复数的共轭形式。

而矩阵是一个下三角矩阵，其中第行第列的元素计算方式：

Retention 的推理循环表示

推理阶段的循环表示论文中定义如下:

怎么理解呢，还是先回顾单个时间步的输出的计算公式：

上述公式最后一步和推理阶段循环表示公式中各个元素的对应关系是：

对应论文中的图示：

图中的表示 GroupNorm。

可以看到在推理阶段，RetNet 在计算当前时间步的输出只依赖于上一个时间步产出的状态矩阵。

其实就是把计算顺序改了一下，先计算的和的相乘然后一直累加到状态矩阵上，最后再和相乘。

而不是像 Transformer 架构那样，每个时间步的计算要先算和前面所有时间步的相乘得到 attention 权值再和相乘求和，这样就需要一直保留历史的和。

Gated Multi-Scale Retention

然后 RetNet 每一层中的 Retention 子模块其实也是分了个头，每个头用不同的参数，同时每个头都采用不同的常量，这也是 Multi-Scale Retention 名称的来由。

则对输入， MSR 层的输出是：

其中，，是激活函数用来生成门控阈值，还有由于每个头均采用不同的，所以每个头的输出要单独做 normalize 之后再 concat。

参考资料

[1] https://arxiv.org/pdf/2307.08621.pdf
[2] https://en.wikipedia.org/wiki/Euler's_formula
[3] https://en.wikipedia.org/wiki/List_of_trigonometric_identities
[4] https://github.com/microsoft/torchscale/blob/main/torchscale/architecture/retnet.py

—END—

高效学习3D视觉三部曲

第一步加入行业交流群，保持技术的先进性

目前工坊已经建立了3D视觉方向多个社群，包括SLAM、工业3D视觉、自动驾驶方向，细分群包括：

[工业方向]三维点云、结构光、机械臂、缺陷检测、三维测量、TOF、相机标定、综合群；

[SLAM方向]多传感器融合、ORB-SLAM、激光SLAM、机器人导航、RTK|GPS|UWB等传感器交流群、SLAM综合讨论群；

[自动驾驶方向]深度估计、Transformer、毫米波|激光雷达|视觉摄像头传感器讨论群、多传感器标定、自动驾驶综合群等。

[三维重建方向]NeRF、colmap、OpenMVS、MVSNet等。

[无人机方向]四旋翼建模、无人机飞控等。

除了这些，还有求职、硬件选型、视觉产品落地等交流群。

大家可以添加小助理微信: dddvisiona，备注：加群+方向+学校|公司, 小助理会拉你入群。

添加小助理微信： dddvisiona，拉你入群

第二步加入知识星球，问题及时得到解答

2.1 「3D视觉从入门到精通」技术星球

针对3D视觉领域的视频课程（三维重建、三维点云、结构光、手眼标定、相机标定、激光/视觉SLAM、自动驾驶等）、源码分享、知识点汇总、入门进阶学习路线、最新paper分享、疑问解答等进行深耕，更有各类大厂的算法工程人员进行技术指导。与此同时，星球将联合知名企业发布3D视觉相关算法开发岗位以及项目对接信息，打造成集技术与就业、项目对接为一体的铁杆粉丝聚集区，6000+星球成员为创造更好的AI世界共同进步，知识星球入口：「3D视觉从入门到精通」

学习3D视觉核心技术，扫描查看，3天内无条件退款