科学空间|Scientific Spaces

在前三篇文章中，我们较为详细地讨论了HiPPO和S4的大部分数学细节。那么，对于接下来的第四篇文章，大家预期我们会讨论什么工作呢？S5、Mamba乃至Mamba2？都不是。本系列文章主要关心SSM的数学基础，旨在了解SSM的同时也补充自己的数学能力。而在上一篇文章我们简单提过S5和Mamba，S5是S4的简化版，相比S4基本上没有引入新的数学技巧，而Mamba系列虽然表现优异，但它已经将$A$简化为对角矩阵，所用到的数学技巧就更少了，它更多的是体现了工程方面的能力。

这篇文章我们来学习一篇暂时还声名不显的新工作《State-Free Inference of State-Space Models: The Transfer Function Approach》（简称RFT），它提出了一个新方案，将SSM的训练、推理乃至参数化，都彻底转到了生成函数空间中，为SSM的理解和应用开辟了新的视角

基础回顾

首先我们简单回顾一下上一篇文章关于S4的探讨结果。S4基于如下线性RNN
\begin{equation}\begin{aligned}
x_{k+1} =&\, \bar{A} x_k + \bar{B} u_k \\
y_{k+1} =&\, \bar{C}^* x_{k+1} \\
\end{aligned}\label{eq:linear}\end{equation}

点击阅读全文...

前面我们用两篇文章《重温SSM（一）：线性系统和HiPPO矩阵》和《重温SSM（二）：HiPPO的一些遗留问题》介绍了HiPPO的思想和推导——通过正交函数基对持续更新的函数进行实时逼近，其拟合系数的动力学正好可以表示为一个线性ODE系统，并且对于特定的基底以及逼近方式，我们可以将线性系统的关键矩阵精确地算出来。此外，我们还讨论了HiPPO的离散化和相关性质等问题，这些内容奠定了后续的SSM工作的理论基础。

接下来，我们将介绍HiPPO的后续应用篇《Efficiently Modeling Long Sequences with Structured State Spaces》（简称S4），它利用HiPPO的推导结果作为序列建模的基本工具，并从新的视角探讨了高效的计算和训练方式，最后在不少长序列建模任务上验证了它的有效性，可谓SSM乃至RNN复兴的代表作之一。

基本框架

S4使用的序列建模框架，是如下的线性ODE系统：
\begin{equation}\begin{aligned}
x'(t) =&\, A x(t) + B u(t) \\
y(t) =&\, C^* x(t) + D u(t)
\end{aligned}\end{equation}

点击阅读全文...

书接上文，在上一篇文章《重温SSM（一）：线性系统和HiPPO矩阵》中，我们详细讨论了HiPPO逼近框架其HiPPO矩阵的推导，其原理是通过正交函数基来动态地逼近一个实时更新的函数，其投影系数的动力学正好是一个线性系统，而如果以正交多项式为基，那么线性系统的核心矩阵我们可以解析地求解出来，该矩阵就称为HiPPO矩阵。

当然，上一篇文章侧重于HiPPO矩阵的推导，并没有对它的性质做进一步分析，此外诸如“如何离散化以应用于实际数据”、“除了多项式基外其他基是否也可以解析求解”等问题也没有详细讨论到。接下来我们将补充探讨相关问题。

离散格式

假设读者已经阅读并理解上一篇文章的内容，那么这里我们就不再进行过多的铺垫。在上一篇文章中，我们推导出了两类线性ODE系统，分别是：
\begin{align}
&\text{HiPPO-LegT:}\quad x'(t) = Ax(t) + Bu(t) \label{eq:legt-ode}\\[5pt]
&\text{HiPPO-LegS:}\quad x'(t) = \frac{A}{t}x(t) + \frac{B}{t}u(t) \label{eq:legs-ode}\end{align}
其中$A,B$是与时间$t$无关的常数矩阵，HiPPO矩阵主要指矩阵$A$。在这一节中，我们讨论这两个ODE的离散化。

点击阅读全文...

前几天，笔者看了几篇介绍SSM（State Space Model）的文章，才发现原来自己从未认真了解过SSM，于是打算认真去学习一下SSM的相关内容，顺便开了这个新坑，记录一下学习所得。

SSM的概念由来已久，但这里我们特指深度学习中的SSM，一般认为其开篇之作是2021年的S4，不算太老，而SSM最新最火的变体大概是去年的Mamba。当然，当我们谈到SSM时，也可能泛指一切线性RNN模型，这样RWKV、RetNet还有此前我们在《Google新作试图“复活”RNN：RNN能否再次辉煌？》介绍过的LRU都可以归入此类。不少SSM变体致力于成为Transformer的竞争者，尽管笔者并不认为有完全替代的可能性，但SSM本身优雅的数学性质也值得学习一番。

尽管我们说SSM起源于S4，但在S4之前，SSM有一篇非常强大的奠基之作《HiPPO: Recurrent Memory with Optimal Polynomial Projections》（简称HiPPO），所以本文从HiPPO开始说起。

点击阅读全文...

近年来，RNN由于其线性的训练和推理效率，重新吸引了不少研究人员和用户的兴趣，隐约有“文艺复兴”之势，其代表作有RWKV、RetNet、Mamba等。当将RNN用于语言模型时，其典型特点就是每步生成都是常数的空间复杂度和时间复杂度，从整个序列看来就是常数的空间复杂度和线性的时间复杂度。当然，任何事情都有两面性，相比于Attention动态增长的KV Cache，RNN的常数空间复杂度通常也让人怀疑记忆容量有限，在Long Context上的效果很难比得上Attention。

在这篇文章中，我们表明Causal Attention可以重写成RNN的形式，并且它的每一步生成理论上也能够以$\mathcal{O}(1)$的空间复杂度进行（代价是时间复杂度非常高，远超平方级）。这表明Attention的优势（如果有的话）是靠计算堆出来的，而不是直觉上的堆内存，它跟RNN一样本质上都是常数量级的记忆容量（记忆瓶颈）。

点击阅读全文...

近年来，线性RNN由于其可并行训练以及常数推理成本等特性，吸引了一定研究人员的关注（例如笔者之前写的《Google新作试图“复活”RNN：RNN能否再次辉煌？》），这让RNN在Transformer遍地开花的潮流中仍有“一席之地”。然而，目前看来这“一席之地”只属于线性RNN，因为非线性RNN无法高效地并行训练，所以在架构之争中是“心有余而力不足”。

不过，一篇名为《Parallelizing Non-Linear Sequential Models over the Sequence Length》的论文有不同的看法，它提出了一种迭代算法，宣传可以实现非线性RNN的并行训练！真有如此神奇？接下来我们一探究竟。

求不动点

原论文对其方法做了非常一般的介绍，而且其侧重点是PDE和ODE，这里我们直接从RNN入手。考虑常见的简单非线性RNN：
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t)\label{eq:rnn}\end{equation}

点击阅读全文...

当前，像ChatGPT之类的LLM可谓是“风靡全球”。有读者留意到，几乎所有LLM都还是用最初的Multi-Head Scaled-Dot Attention，近年来大量的Efficient工作如线性Attention、FLASH等均未被采用。是它们版本效果太差，还是根本没有必要考虑效率？其实答案笔者在《线性Transformer应该不是你要等的那个模型》已经分析过了，只有序列长度明显超过hidden size时，标准Attention才呈现出二次复杂度，在此之前它还是接近线性的，它的速度比很多Efficient改进都快，而像GPT3用到了上万的hidden size，这意味着只要你的LLM不是面向数万长度的文本生成，那么用Efficient改进是没有必要的，很多时候速度没提上去，效果还降低了。

那么，真有数万甚至数十万长度的序列处理需求时，我们又该用什么模型呢？近日，Google的一篇论文《Resurrecting Recurrent Neural Networks for Long Sequences》重新优化了RNN模型，特别指出了RNN在处理超长序列场景下的优势。那么，RNN能否再次辉煌？

点击阅读全文...

本来笔者已经决心不玩RNN了，但是在上个星期思考时忽然意识到RNN实际上对应了ODE（常微分方程）的数值解法，这为我一直以来想做的事情——用深度学习来解决一些纯数学问题——提供了思路。事实上这是一个颇为有趣和有用的结果，遂介绍一翻。顺便地，本文也涉及到了自己动手编写RNN的内容，所以本文也可以作为编写自定义的RNN层的一个简单教程。

注：本文并非前段时间的热点“神经ODE”的介绍（但有一定的联系）。

RNN基本

什么是RNN？

众所周知，RNN是“循环神经网络（Recurrent Neural Network）”，跟CNN不同，RNN可以说是一类模型的总称，而并非单个模型。简单来讲，只要是输入向量序列$(\boldsymbol{x}_1,\boldsymbol{x}_2,\dots,\boldsymbol{x}_T)$，输出另外一个向量序列$(\boldsymbol{y}_1,\boldsymbol{y}_2,\dots,\boldsymbol{y}_T)$，并且满足如下递归关系
$$\boldsymbol{y}_t=f(\boldsymbol{y}_{t-1}, \boldsymbol{x}_t, t)\tag{1}$$
的模型，都可以称为RNN。也正因为如此，原始的朴素RNN，还有改进的如GRU、LSTM、SRU等模型，我们都称为RNN，因为它们都可以作为上式的一个特例。还有一些看上去与RNN没关的内容，比如前不久介绍的CRF的分母的计算，实际上也是一个简单的RNN。

说白了，RNN其实就是递归计算。

点击阅读全文...