【AI学习】Mamba学习(十五):关于S4模型计算优化的再学习

前面理解了S4模型,但是对于具体的优化算法,还是没有完全理解透彻。现在补充学习。

S4 = SSM + HiPPO + Structured Matrices

具体方案:Structured State Spaces (S4)

简单总结:用HiPPO-LegS的矩阵形式初始化SSM,解决训练的稳定性问题。但是,基于卷积的并行化训练,依然处理复杂度很高,需要简化。

对角化

如何简化计算?最直接的思维就是矩阵A的对角化。前面的公式5的卷积形式,矩阵A的L次幂,如果是对角矩阵,计算就可以大大简化。

文章3.1节,给出了这种对角化动机的目的。

但是文章接着说,不幸的是,由于数值问题,对角化的简单应用在实践中不可行。尽管有其他方法,数值上也不稳定。

为什么数值不稳定,简单看一下,以N=3为例,对角化的结果是:

"HiPPO 矩阵的对角化涉及到的矩阵元素在状态大小 N 增大时会呈指数级增长,这使得对角化在数值上变得不稳定和不可行"。

那怎么办,3.2节给出方法:将HiPPO 矩阵A分解为正态矩阵和低秩矩阵的和。这样处理后获得了一个反对称矩阵,"重点来了,反对称矩阵不单单一定可以对角化,它一定可以被正交矩阵(复数域叫做酉矩阵)对角化!酉矩阵一般数值稳定性都非常好"。

S4 Parameterization: Normal Plus Low-Rank

前面的讨论意味着我们应该进行共轭计算通过条件良好的矩阵 V 。理想的情况是,当矩阵 A 可由完全条件(即酉)矩阵对角化时。根据线性代数的谱定理,这正是normal矩阵的一类。然而,这类矩阵具有限制性;特别是它不包括HiPPO 矩阵(2)。

我们观察到,尽管 HiPPO 矩阵不是normal矩阵,但它可以被分解为normal和低秩矩阵的和。然而,这本身仍然没有用:与对角矩阵不同,对这个幂进行累加(在(5)中)仍然很慢,也不容易优化。我们通过同时应用三种新技术克服了这一瓶颈。

上面的描述还是有点抽象,分步骤来看:

矩阵A分解为Normal Plus Low-Rank (NPLR)形式

上面说的公式(2)就是下面的LegS形式,对于的低秩r=1

论文中在附录C.1NPLR Representations of HiPPO Matrices进行了说明。

已知 HiPPO 矩阵A可以表示为:


"重点来了,可以看到这是一个反对称矩阵,所以它一定可以(在复数域中)对角化!于是我们就将A分解为了可对角化矩阵与低秩矩阵之和!可能有读者质疑,原本 A就 一定是可对角化矩阵,但还是有数值稳定性问题,难道这个反对称矩阵的对角化不用担心数值稳定性问题吗?

重点的重点来了,反对称矩阵不单单一定可以对角化,它一定可以被正交矩阵(复数域叫做酉矩阵)对角化!酉矩阵一般数值稳定性都非常好,所以不用担心这个问题,这也就是为什么我们不直接对角化 ,而绕一圈来构建反对称矩阵的原因。"

这样就将矩阵A转换为了正规矩阵+低秩矩阵的形式。

但是,论文中指出:然而,这本身仍然没有用:与对角矩阵不同,对这个幂进行累加(在(5)中)仍然很慢,也不容易优化。

公式5见下面:

生成函数

论文中指出,利用截断的生成函数!

这里的生成函数如何理解?熟悉卷积运算的就知道,卷积运算计算量大,可以先做FFT,在频域变成乘法,然后IFFT。这是利用FFT的简化卷积运算经常使用的方法。只不过,这里傅立叶变换所需要的实际是"截断生成函数",将无限长度截断为L。

直观地说,生成函数基本上将SSM卷积滤波器从时域转换为频域。重要的是,它保留了相同的信息,并且可以从其生成函数的评估中恢复所需的SSM卷积滤波器。

这样 就把矩阵的幂的问题转化为矩阵求逆。

Woodbury Correction

然后,通过Woodbury恒等式来解决低秩问题。

虽然DPLR矩阵由于低秩项而不能有效地求幂,但它们可以通过众所周知的Woodbury恒等式有效地反转。

有了这个Woodbury恒等式,我们可以将DPLR矩阵A上的SSM生成函数转换为仅在其对角分量上的生成函数。


Cauchy Kernel

柯西矩阵计算是数值分析中一个研究得非常好的问题,有快速算法和基于著名的快速多极方法(FMM)的快速数值算法。

下面的卷积核的计算,第2、3步骤就是计算Lemma C.3中的K(z)

然后又因为截断为L长度的生成函数,需要下面的第4、5步骤。

具体第三章涉及的公式推导,可以参见苏神在文章《重温状态空间模型SSM:HiPPO的高效计算(S4)》中的详细推导。

S4模型的处理示意图

在S5的论文中,给出了S4模型的处理示意图。

相关推荐
大千AI助手1 小时前
指数分布:从理论到机器学习应用
人工智能·机器学习·参数估计·概率密度函数·mle·指数分布·累积分布函数
MATLAB代码顾问1 小时前
MATLAB绘制多种混沌系统
人工智能·算法·matlab
搬砖的小码农_Sky1 小时前
人形机器人:Tesla Optimus的AI集成细节
人工智能·ai·机器人
做运维的阿瑞1 小时前
2025 年度国产大模型「开源 vs. 闭源」深度评测与实战指南
人工智能·低代码·开源
渡我白衣1 小时前
深度学习入门(三)——优化算法与实战技巧
人工智能·深度学习
可触的未来,发芽的智生1 小时前
触摸未来2025.10.10:记忆的种子,当神经网络拥有了临时工作区,小名喜忆记系统
人工智能·python·神经网络·机器学习·架构
极客BIM工作室1 小时前
演化搜索与群集智能:五种经典算法探秘
人工智能·算法·机器学习
繁花与尘埃1 小时前
HTML5简介与基本骨架(本文为个人学习笔记,内容整理自哔哩哔哩UP主【非学者勿扰】的公开课程。 > 所有知识点归属原作者,仅作非商业用途分享)
笔记·学习·html5
Guheyunyi1 小时前
消防管理系统如何重构现代空间防御体系
大数据·运维·人工智能·安全·信息可视化·重构
东方芷兰1 小时前
LLM 笔记 —— 04 为什么语言模型用文字接龙,图片模型不用像素接龙呢?
人工智能·笔记·深度学习·语言模型·自然语言处理