Mamba：6 线性RNN

庞德公

已于 2024-05-12 20:22:05 修改

阅读量1.2k

点赞数 41

分类专栏：新技术解读文章标签： rnn 人工智能深度学习

于 2024-05-12 18:46:16 首次发布

本文链接：https://blog.csdn.net/Janexjy/article/details/138763101

版权

新技术解读专栏收录该内容

21 篇文章 5 订阅

订阅专栏

若在阅读过程中有些知识点存在盲区，可以回到如何优雅的谈论大模型重新阅读。另外斯坦福2024人工智能报告解读为通识性读物。若对于如果构建生成级别的AI架构则可以关注AI架构设计。技术宅麻烦死磕LLM背后的基础模型。当然最重要的是订阅跟随“鲁班模锤”。

Mamba自从出道就一直被拿来和Transformer对比，文章巨多，带着标题的爆炸性字样“颠覆Transfomer”，“全面包围”等等。文章深浅不一，小编认为其背后的设计思路以及演化过程更加重要。为此本专题将从山上有座庙开始，沿着RNN、SSM、S4、HIPPO一路走来。期间抛弃冗杂的文辞，以更接地气的模式协助读者打开新的逻辑思维。

线性RNN

2023年的<Resurrecting Recurrent Neural Networks for Long Sequences>将RNN拉回了赛道。文中提出的LRU通过简单的线性函数替换神经网络，然后在两个自注意力层之间用神经网络连接起来以弥补非线性处理能力的不足。会不会有点拗口，先来看看下面的结构图：

左边是线性RNN引入新单元的架构Deep Linear Recurrent Unit(LRU)，这个模型通过堆叠LRU块而构成可扩展，每两层之间是非线性投影，模型还使用了跳跃连接（类似残差网络）和归一化方法。在LRU单元内部，不同时刻的状态变化成为了线性变换，（会不会很眼熟。）

循环操作被调整为纯线性，就可以解决之前RNN的两个痛点。它在 O(log(n)) 时间内并行计算应用于n个向量的线性递推。然后通过特殊的参数设计用来解决训练中的梯度问题（左上部分）。

线性替代

线性循环神经网络是什么？其实用线性函数替换神经网络即是，大白话就是将上述的激活函数去掉。

论文作者做了实验得出结论其实线性化对于块状结构其实已经足够。

值得一提的是线性RNN配合MLP（非线性映射）在sequence to sequence map中表现尚可。下图中的隐藏状态能够还原出足够的信息。

不仅仅如此，线性化的好处在于并行化和矩阵对角化，甚至之前的魔幻矩阵A也可以不通过HiPPO模式进行预设，而是通过训练学习得到。

将并行化和矩阵对角化结合在训练速度方面的确超过了s4和s5。

梯度消失与爆炸

先来看个场景，简化模型且假设第一个输入之后的每个输入都是0。那么按照循环网络，在每一步，前一个输出都会乘以权重向量得到后一个，n步之后得到了 $w^{n}$ 与x的乘积。

这个时候假如按照梯度稍微的调整权重矩阵，那么再次训练的时候得到了如下的式子。

两者的差值是一个恒定的缩放因子 $w^{n-1}$ 乘以x1。这里存在的问题是随着n变大，这项要么变得非常小，要么变得非常大，具体取决于w中的值是小于还是大于1。（想象一下 $1.2^{10000}$ 和 $0.8^{10000}$ 的数值大小）

无论哪种情况，这都是一个问题：若梯度非常大，则权重变化太大，已经学习的现有功能就会被破坏。若梯度非常小，则权重变化不够，根本学不到任何东西。这被称为梯度消失和爆炸问题。

值得一提的是，这对于常规神经网络来说可能不是问题，因为在每一层中都会使用不同的权重。有些层的权重可以小于1，有些层的权重可以大于 1，只要梯度保持大致相同的大小，神经网络就能够学习。而且目前也有很多的权重配置方法可以确保训练过程中的梯度稳定。

但对于循环结构，在每个步骤中使用相同的权重，只有一种稳定的配置，就是都为1。任何偏离1的情况都会导致梯度呈指数增长或衰减，当权重为复数的时候使用其绝对值。那么我们如何解决梯度消失和爆炸问题呢？

只要循环权重为1并且输入为0，RNN梯度就是稳定的。线性RNN论文中，作者建议在这种稳定状态下初始化。具体来说，以复极坐标形式 $ae^{ib}$ 参数化权重，其中a是幅度，b是角度。然后通过此 $e^{-e^{a}}$ 函数调整，将a的幅度限制为小于1，该函数始终输出0到1之间的数字，而不是常规的以0 为中心的正态分布中随机采样。

$e^{-e^{a}}$ 均匀分布在0.999和1之间。b均匀分布在0和π/10弧度之间。这确保了在初始化时，权重都非常接近1。

当然输入的信号要乘以另一个可学习参数Δ，初始化为 $\sqrt{1-e^{-e^{a}}}$ ，因为 $e^{-e^{a}}$ 接近于1，所以这个参数是非常小的数字。这确保了初始化时输入都接近0。

虽然模型在训练的过程中很难保证梯度的稳定性，但是按照这种方法进行初始化则在实践中已经证明，它能够学会数万步的上下文。

后话

计算速度飞快的线性RNN，可以通过训练应付极长的上下文。在线性RNN的原始论文中，它在远程竞技场基准评估了该模型。

例如在PathX任务中，模型必须根据图像是否包含两个圆之间的完整虚线路径对图像进行分类。只不过图像被展平为16000 像素的长序列。线性RNN在长距离领域实现了最先进的性能，在所有任务中平均比Transformer高出约33%。

受控制理论启发的状态空间模型的初始化与线性RNN初始化大致相同。不管如何，当听到状态空间模型时，都会想到线性RNN。虽然线性RNN在长距离竞技场基准测试中确实表现出色，但这并不意味着它们是良好的语言模型。对于自然语言的建模，线性RNN的性能比Transformer差得多，尤其是和各种最先进的语言模型对比。

正如Mamba论文中所指出的，其原因是线性RNN无法选择性地忘记输出向量中的信息。如果权重接近 0，则每次输入后输出向量将设置为0，实际上模型将始终立即忘记当前输入之前的任何内容。如果权重接近1，则输出向量在与权重相乘时不会改变，因此输出向量将累积来自观察到的所有输入的信息。然而对于自然语言模型而言，需要它具备能够根据看到的输入来决定何时存储信息以及何时忘记信息的能力。不过整体而言，线性RNN将RNN拉回了生死线。