Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers

作者/机构: Albert Gu†, Isys Johnson†, Karan Goel†, Khaled Saab†, Tri Dao†, Atri Rudra‡, and Christopher Ré†
†Department of Computer Science, Stanford University
‡Department of Computer Science and Engineering, University at Buffalo, SUNY

A1 主要贡献

本文旨在解决机器学习中高效建模长序列数据(超过几千个时间步)的挑战。传统序列模型范式,如循环神经网络(RNNs)、卷积神经网络(CNNs)和神经微分方程(NDEs),各有其优缺点。RNNs是自然的状态化模型,但训练缓慢且存在优化难题(如梯度消失);CNNs训练可并行化,但推理成本高且上下文长度受限;NDEs理论上能处理连续时间和长期依赖,但效率极低。理想的模型应结合这些范式的优点:即可并行训练(卷积特性)、状态化推理(循环特性)、适应时间尺度(微分方程特性),并能高效处理极长序列。

核心问题与研究目标:
核心问题是如何构建一个模型家族,能够融合RNN、CNN和NDE的优点,同时克服它们在处理长序列时的局限性。研究目标是创建一个富有表达力、计算高效且能捕捉长期依赖的序列模型。

创新点与主要贡献:
为实现这一目标,本文提出了一种受控制系统启发的简单序列模型——线性状态空间层(Linear State-Space Layer, LSSL)。LSSL通过模拟一个线性的连续时间状态空间表示来映射序列 $u \to y$。
其核心贡献如下:
1. 提出LSSL模型:LSSL通过一个隐式状态 $x(t)$ 将一维函数或序列 $u(t)$ 映射到 $y(t)$,其动力学由以下线性连续时间状态空间方程定义:



其中,$A$ 控制系统演化,$B, C, D$ 是投影参数。LSSL模型具有三种视图(如图1所示),继承了三大模型家族的优点:
* 循环视图:可离散化为一个线性循环模型,实现高效的状态化推理。
* 卷积视图:其离散时间版本等价于一个卷积操作,可利用FFT进行高效的并行化训练。
* 连续时间视图:本质是一个微分方程,能处理不规则采样数据并适应不同时间尺度。

图1: (LSSL的三种视图) 线性状态空间层是一个映射 $u_t \in \mathbb{R} \to y_t \in \mathbb{R}$,其中每个特征 $u_t \to y_t$ 是通过一个参数 $\Delta_t$ 离散化一个状态空间模型 A, B, C, D 来定义的。底层的状态空间模型通过将状态矩阵A和时间尺度$\Delta_t$组合成一个转移矩阵$\bar{A}$来定义一个离散循环。(左) 作为一个隐式的连续模型,不规则间隔的数据可以通过使用不同的时间尺度$\Delta_t$来离散化相同的矩阵A来处理。(中) 作为一个循环模型,可以通过按时间计算层(即,一次一个垂直切片(ut, xt, yt), (ut+1, xt+1, yt+1), ...),通过展开线性循环来高效地执行推理。(右) 作为一个卷积模型,可以通过并行地深度计算层(即,一次一个水平切片(ut)t∈[L], (yt)t∈[L], ...),通过与特定的滤波器进行卷积来高效地执行训练。

  1. 理论联系与泛化:本文证明LSSL不仅未牺牲表达能力,反而泛化了RNN和CNN。理论上,任何一维卷积核都可以被LSSL近似。此外,本文揭示了RNN中的门控机制(gating)与ODE离散化中的步长 $\Delta_t$ 之间的深刻联系,表明一些流行的RNN模型可视为LSSL的特例。

  2. 解决长期依赖问题:为解决LSSL在处理长序列时继承的长期记忆难题,本文利用并推广了关于连续时间记忆的HiPPO理论【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】。通过引入一类特殊的可训练结构化矩阵A,使LSSL能够:

    • 数学上捕捉关于可学习度量的长期依赖关系。
    • 在特定计算模型下,通过新算法实现理论上的计算加速,即使在学习度量A和时间尺度 $\Delta_t$ 时也是如此。
  3. 卓越的经验验证:在多个时间序列基准测试中,由LSSL层堆叠而成的深度神经网络取得了当前最佳(SOTA)结果:

    • 在序列图像分类(如sCIFAR准确率提升超10%)和医疗健康回归任务(RMSE降低高达80%)上超越了最新的RNN、CNN和NDE模型。
    • 在序列长度高达38000的CelebA分类任务中,小型LSSL模型的性能接近于参数量10倍的专用ResNet-18视觉模型。
    • 在长度为16000的原始语音分类任务中,LSSL不仅比先前方法准确率高出24个点,训练时间缩短至1/5,甚至超越了在100倍短的预处理序列上使用手工特征的基线模型。

A3 背景知识

微分方程的近似。任何微分方程 $\dot{x}(t) = f(t, x(t))$ 都有一个等价的积分形式 $x(t) = x(t_0) + \int_{t_0}^{t} f(s, x(s)) ds$。数值求解该方程的一种方法是存储对 $x$ 的某种近似,并在迭代方程时保持 $f(t, x)$ 内部的 $x$ 固定。例如,皮卡德迭代(Picard iteration)常用于证明ODE解的存在性,它通过迭代方程 $x_{i+1}(t) := x_i(t_0) + \int_{t_0}^{t} f(s, x_i(s)) ds$ 来找到一系列函数 $x_0(t), x_1(t), \ldots$ 来逼近积分方程的解 $x(t)$。

离散化。对于给定的离散时间序列 $t_i$,可以通过迭代方程 $x(t_{i+1}) = x(t_i) + \int_{t_i}^{t_{i+1}} f(s, x(s)) ds$ 来找到 $x(t_0), x(t_1), \ldots$ 的近似值。对右侧积分的不同近似方法会产生不同的离散化方案。本文特别关注一种称为广义双线性变换(generalized bilinear transform, GBT)的离散化方法,它专门用于形如公式(1)的线性ODE。给定步长 $\Delta_t$,GBT的更新规则是:


其中三个重要特例是:$\alpha=0$ 对应经典欧拉法,即一阶近似 $x(t+\Delta_t) = x(t) + \Delta_t \cdot x'(t)$;$\alpha=1$ 称为后向欧拉法;$\alpha=\frac{1}{2}$ 称为双线性方法,该方法能保持系统的稳定性【索引61,Performance recovery in digital implementation of analogue systems,2007,SIAM journal on control and optimization】。在第3.2节中,本文将展示后向欧拉法和皮卡德迭代与RNNs有关。而双线性离散化将是计算连续时间模型精确离散时间近似的主要方法。特别地,定义 $\bar{A}$ 和 $\bar{B}$ 为 $\alpha=\frac{1}{2}$ 时公式(3)中出现的矩阵,则离散时间状态空间模型为:

$\Delta_t$作为时间尺度。在大多数模型中,它们能捕捉的依赖长度大致与 $\frac{1}{\Delta_t}$ 成正比,因此步长 $\Delta_t$ 也被称为时间尺度。这是将连续时间ODE转换为离散时间循环的内在部分,并且在大多数基于ODE的RNN模型中,它是一个重要的、不可训练的超参数【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS;索引47,Unicornn: A recurrent model for learning very long time dependencies,2021,ICML;索引58,Legendre memory units: Continuous-time representation in recurrent neural networks,2019,NeurIPS】。然而,在3.2节中本文展示了经典RNN的门控机制实际上是学习 $\Delta_t$ 的一种形式。此外,当LSSL被看作一个CNN时,时间尺度 $\Delta_t$ 可以被视为控制卷积核的宽度(3.2节)。理想情况下,所有基于ODE的序列模型都应该能够自动学习合适的时间尺度。

连续时间记忆。考虑一个输入函数 $u(t)$、一个固定的概率度量 $\omega(t)$ 和一组N个基函数(如多项式)。在每个时刻 $t$, $u$ 在 $t$ 之前的历史可以被投影到这个基上,得到一个系数向量 $x(t) \in \mathbb{R}^N$,它代表了相对于给定度量 $\omega$ 的 $u$ 的历史的最优近似。将函数 $u(t) \in \mathbb{R}$ 映射到系数 $x(t) \in \mathbb{R}^N$ 的算子被称为关于度量 $\omega$ 的高阶多项式投影算子(High-Order Polynomial Projection Operator, HiPPO)。在一些特殊情况下,如均匀度量 $\omega = I_{\{[0, 1]\}}$ 和指数衰减度量 $\omega(t) = \exp(-t)$,Gu等人【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】证明了 $x(t)$ 满足一个微分方程 $\dot{x}(t) = A(t)x(t) + B(t)u(t)$(即公式(1)的形式),并推导出了矩阵 $A$ 的闭式解。他们的框架为设计处理长依赖的记忆模型提供了一种有原则的方法;然而,他们只证明了这几个特例。

A2 方法细节

3. 线性状态空间层 (LSSL)

本节定义了LSSL这一核心抽象,它是一个泛化了循环和卷积的模型家族。3.1节首先正式定义了LSSL,然后讨论了如何从多个视角来计算它。反之,3.2节展示了LSSL与最流行的RNN机制之间的关系。

3.1 LSSL的不同视角

LSSL的定义与计算。给定一个固定的状态空间表示 A, B, C, D,LSSL是通过离散化线性状态空间模型(1)和(2)定义的序列到序列映射。具体来说,一个LSSL层拥有参数 A, B, C, D 和 $\Delta_t$。它作用于一个输入 $u \in \mathbb{R}^{L \times H}$,该输入代表一个长度为L的序列,每个时间步都有一个H维的特征向量。每个特征 $h \in [H]$ 定义一个序列 $(u_t^{(h)})_{t \in [L]}$,它与一个时间尺度 $\Delta_{t_h}$ 结合,通过离散化的状态空间模型(4)+(5)来定义一个输出 $y^{(h)} \in \mathbb{R}^L$。

多种计算视图。在计算上,离散时间的LSSL可以从多个角度来看待(如图1)。
- 作为循环模型。循环状态 $x_{t-1} \in \mathbb{R}^{H \times N}$ 携带了时间 $t$ 之前所有输入的上下文。当前状态 $x_t$ 和输出 $y_t$ 可以通过简单地遵循方程(4)+(5)来计算。因此,LSSL是一个循环模型,具有高效和状态化的推理能力,可以处理(可能无限长的)输入序列,而每个时间步仅需固定的计算和存储。
- 作为卷积模型。为简单起见,设初始状态 $x_{-1}=0$。那么方程(4)+(5)明确地得出:


于是,$y$ 便是(非循环的)卷积 $y = K_L(\bar{A}, \bar{B}, C) * u + Du$,其中卷积核为:

因此,LSSL可以被看作是一个卷积模型,其中整个输出 $y \in \mathbb{R}^{H \times L}$ 可以通过一次卷积一次性计算出来,这可以通过三个快速傅里叶变换(FFT)高效实现。

计算瓶颈。需要注意的是,(i)循环视图的瓶颈在于模拟公式(4)时与离散化状态矩阵 $\bar{A}$ 的矩阵向量乘法(MVM),而(ii)卷积视图的瓶颈在于计算Krylov函数 $K_L$(公式7)。在本节中,我们假设LSSL的参数是固定的,这意味着 $\bar{A}$ 和 $K_L(\bar{A}, \bar{B}, C)$ 可以被缓存以提高效率。然而,学习参数A和 $\Delta_t$ 将涉及反复重新计算这些值,这在实践中是不可行的。我们将在第4.2节重新讨论并解决这个问题。

3.2 LSSL的表达能力

模型表达能力的权衡。一个模型要同时具备循环和卷积的特性,人们可能会预期它在其他方面会受到限制。确实,虽然【索引12,Parallelizing legendre memory unit training,2021,ICML】和【索引44,Ckconv: Continuous kernel convolution for sequential data,2021,arXiv】也观察到某些循环可以用卷积替代,但他们指出,卷积是否能被循环替代并不明显。此外,LSSL是一个线性循环,而流行的RNN模型是每个时间步之间带有激活函数的非线性序列模型。我们现在将证明,LSSL出人意料地并没有受限的表达能力。

卷积是LSSL。关于状态空间系统(1)+(2)的一个众所周知的事实是,输出 $y$ 通过与系统的冲激响应 $h$ 进行卷积与输入 $u$ 相关联:$y(t) = \int h(\tau)u(t-\tau)d\tau$。反之,一个N阶有理函数的卷积核 $h$ 可以由一个大小为N的状态空间模型来表示【索引59,Linear state-space control systems,2007,Wiley Online Library】。因此,一个任意的卷积核 $h$ 可以通过有理函数(例如,通过Padé近似)来逼近,并由一个LSSL表示。特别地,对于使用HiPPO矩阵的LSSL(第2节和4.1节),LSSL与卷积的关系还有另一种直观的解释。考虑当A对应于均匀度量(在文献中称为LMU【索引58,Legendre memory units: Continuous-time representation in recurrent neural networks,2019,NeurIPS】或HiPPO-LegT【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】)的特殊情况。那么对于一个固定的 $dt$,方程(1)只是在大小为 $\frac{1}{\Delta_t}$ 的滑动窗口内记忆输入,而方程(2)则从这个窗口中提取特征。因此,LSSL可以被解释为自动学习具有可学习核宽度的卷积滤波器。

RNN是LSSL。我们展示了两个关于RNNs的可能具有更广泛意义的结论。第一个结论是,RNNs中无处不在的门控机制,通常被认为是平滑优化的启发式方法【索引28,Long short-term memory,1997,Neural computation】,实际上等同于步长或时间尺度 $\Delta_t$。
- 引理 3.1: 一个(一维)门控循环 $x_t = (1-\sigma(z))x_{t-1} + \sigma(z)u_t$,其中 $\sigma$ 是sigmoid函数,z是任意表达式,可以被看作是一维线性ODE $\dot{x}(t) = -x(t) + u(t)$ 的GBT($\alpha=1$)(即后向欧拉)离散化。
- 证明: 应用离散化需要一个正的步长 $\Delta_t$。参数化一个正函数的最简单方法是使用指数函数 $\Delta_t = \exp(z)$ 应用于任何表达式z。将此代入(3)式,并设 $A=-1, B=1, \alpha=1$,恰好可以得到上述门控循环的表达式。

深度线性RNN与ODE的关系。虽然引理3.1涉及使用离散化来近似连续系统,但第二个结果是关于使用皮卡德迭代(Picard iteration,第2节)来近似它们。粗略地说,深度线性RNN的每一层可以被看作是连续的皮卡德迭代 $x_0(t), x_1(t), \dots$ ,它们逼近一个由非线性ODE定义的函数 $x(t)$。这表明,我们通过使用线性而非非线性循环并不会损失建模能力,并且非线性可以被“移动”到深度神经网络的深度方向,从而在不牺牲表达能力的情况下提高速度。
- 引理 3.2: (无限)深的、堆叠的、阶数N=1且带有位置级非线性函数的LSSL层可以近似任何非线性ODE $\dot{x}(t) = -x + f(t, x(t))$。

LSSL与流行RNN的统一视角。许多最流行且有效的RNN变体,如LSTM【索引28,Long short-term memory,1997,Neural computation】、GRU【索引14,Empirical evaluation of gated recurrent neural networks on sequence modeling,2014,arXiv】、QRNN【索引5,Quasi-recurrent neural networks,2016,arXiv】和SRU【索引33,Simple recurrent units for highly parallelizable recurrence,2017,arXiv】,都涉及一个隐藏状态 $x_t \in \mathbb{R}^H$,其中H个隐藏单元被独立地“门控”。应用引理3.1,它们实际上也近似于引理3.2中形式的ODE。因此,LSSL和这些流行的RNN模型可以被看作都在近似相同类型的底层连续动态过程,通过在深度方向上使用皮卡德近似,在时间方向上使用离散化(门控)。附录C给出了精确的陈述和证明。

3.3 深度LSSL

基础LSSL层的扩展。基础LSSL定义为在长度为L的一维序列上从 $\mathbb{R}^L \to \mathbb{R}^L$ 的序列到序列映射,由参数 $A \in \mathbb{R}^{N \times N}, B \in \mathbb{R}^{N \times 1}, C \in \mathbb{R}^{1 \times N}, D \in \mathbb{R}^{1 \times 1}, \Delta_t \in \mathbb{R}$ 参数化。对于一个隐藏维度为H(即特征维度大于1)的输入序列,我们只需将参数B, C, D, $\Delta_t$ 增加一个H维度进行广播。这H个副本中的每一个都是独立学习的,因此有H个不同版本的一维LSSL独立地处理每个输入特征。总的来说,独立的LSSL层是一个序列到序列的映射,其接口与标准序列模型层(如RNNs, CNNs, Transformers)相同。

深度LSSL架构。在深度神经网络中,完整的LSSL架构的定义与标准序列模型(如深度ResNets和Transformers)类似,包括将LSSL层与归一化层和残差连接堆叠起来。完整的架构细节在附录B中有描述,包括A和$\Delta_t$的初始化、计算细节以及其他架构细节。

4. 结合LSSL与连续时间记忆

在第3节中,我们介绍了LSSL模型,并展示了它共享卷积和循环的优点,同时还对它们进行了泛化。现在,我们讨论并解决其主要局限性,特别是处理长期依赖(第4.1节)和高效计算(第4.2节)。

4.1 将长期依赖融入LSSL

LSSL的局限性。LSSL的通用性意味着它们可能会继承循环和卷积在处理长期依赖方面的问题(第1节)。例如,从循环的角度看,重复乘以A可能会遭受梯度消失问题的困扰【索引39,On the difficulty of training recurrent neural networks,2013,ICML;索引44,Ckconv: Continuous kernel convolution for sequential data,2021,arXiv】。我们凭经验证实,使用随机状态矩阵A的LSSL作为通用序列模型实际上并不有效(第5.4节)。

利用HiPPO框架解决局限性。然而,这些数学化的连续时间模型的一个优点是它们在理论上是可分析的,并且可以推导出特定的A矩阵来解决这个问题。特别是,HiPPO框架(第2节)描述了如何根据一个度量 $\omega$ 在连续时间内记忆一个函数【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】。这个将函数映射到其过去连续表示的算子表示为hippo($\omega$),并且在三个特例中被证明具有方程(1)的形式。然而,这些矩阵是不可训练的,因为没有其他已知的A矩阵是hippo算子。为了解决这个问题,我们在理论上解决了【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】中的一个开放问题,证明了对于任何度量 $\omega$,hippo($\omega$) 都会得到一个具有结构化矩阵A的公式(1)。
- 定理 1 (非正式): 对于任意度量 $\omega$,最优记忆算子hippo($\omega$)具有 $\dot{x}(t) = Ax(t) + Bu(t)$ (1) 的形式,其中A是一个低循环宽度(low recurrence-width, LRW)【索引17,A two-pronged progress in structured dense matrix vector multiplication,2018,SIAM】的状态矩阵。
对于覆盖经典正交多项式(OPs)【索引52,Orthogonal Polynomials,1967,American Mathematical Society】(特别是对应于雅可比和拉盖尔多项式)的度量,结构性更强。
- 推论 4.1: 对于对应于经典OPs的 $\omega$,hippo($\omega$) 是3-拟可分(3-quasiseparable)的。

结构化矩阵A的引入。虽然超出了本节的范围,我们提到LRW矩阵是一类具有线性时间矩阵向量乘法(MVM)的结构化矩阵【索引17,A two-pronged progress in structured dense matrix vector multiplication,2018,SIAM】。我们在附录D中定义了这个类并证明了定理1。拟可分矩阵是一类相关的结构化矩阵,具有额外的算法特性。我们在附录D.3的定义4中定义了这些矩阵并证明了推论4.1。定理1告诉我们,使用特定结构化矩阵类中的状态矩阵A的LSSL将带有连续时间记忆的理论解释。理想情况下,我们希望能够在这个类中自动学习最好的A;然而,这会遇到我们接下来要解决的计算挑战(第4.2节)。目前,我们将LSSL-fixed或LSSL-f定义为其中A矩阵固定为由【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】规定的HiPPO矩阵之一的模型。

4.2 LSSL的理论上高效的算法

训练A和$\Delta_t$的挑战。虽然A和$\Delta_t$是LSSL中最关键的参数,它们分别控制状态空间(参见4.1节)和时间尺度(参见2和3.2节),但在一个朴素的LSSL中训练它们是不可行的。特别是,3.1节指出,要分别计算循环和卷积视图,需要对A进行高效的矩阵向量乘法(MVM)和Krylov函数(7)计算。然而,前者似乎涉及矩阵求逆(3),而后者似乎需要将A自乘L次。

结构化矩阵带来的计算效率。在本节中,我们证明,将A限制在拟可分矩阵类中(推论4.1),这不仅赋予了LSSL理论上记忆长依赖的能力,同时也赋予了它计算效率。首先,众所周知,拟可分矩阵具有高效的(线性时间)MVM【索引40,Computing with quasiseparable matrices,2016,ACM】。我们证明它们也具有快速的Krylov函数计算,从而允许通过卷积进行高效训练。
- 定理 2: 对于任何k-拟可分矩阵A(常数k)和任意的B, C,Krylov函数 $K_L(A, B, C)$ 可以在准线性时间和空间 $O˜(N+L)$ 以及对数深度(即可并行化)内计算。操作计数是在精确算术模型下,不考虑位复杂性或数值稳定性。

定理2的意义与局限。我们强调定理2并非显而易见。举例来说,对于一个通用矩阵A,展开(7)式需要 $O(LN^2)$ 的时间。即使A是高度结构化的,具有线性计算复杂度,也需要 $O(LN)$ 的操作和线性深度。深度可以通过平方技术(批量乘以 $A, A^2, A^4, \dots$)来减少,但这又需要 $O(LN)$ 的中间存储。实际上,定理2的算法相当复杂(附录E),涉及对多项式矩阵进行分治递归,并利用了(7)式与幂级数 $C(I-Ax)^{-1}B$ 相关的观察。除非另有说明,完整的LSSL指的是A满足推论4.1的LSSL。总之,在这个结构化矩阵族内学习,通过定理1同时赋予LSSL长程记忆能力,并通过定理2在理论上是计算可行的。我们注意到一个警示,即定理2是基于精确算术而非浮点数的,因此更多地被视为一个概念验证,证明LSSL在理论上可以是计算高效的。我们在第6节中更多地评论LSSL的局限性。

A4 实验环境与结果

实验环境

  • 数据集:
    • 序列图像分类: 序列MNIST (sMNIST, 长度784), 排列MNIST (pMNIST, 长度784), 序列CIFAR (sCIFAR, 长度 1024)。这些是测试模型捕捉长期依赖能力的流行基准。
    • 医疗时间序列回归: BDIMC医疗数据集,包含长度为4000的时间序列,用于预测呼吸率(RR)、心率(HR)和血氧饱和度(SpO2)。
    • 极长序列分类:
      • 序列CelebA: 新构建的任务,将178x218的图像展平为长度38800的序列,进行4种面部属性分类。
      • Speech Commands (SC): 1秒音频剪辑分类。使用了两种设置:原始信号(SC-Raw, 长度16000)和预处理的MFCC特征(SC-MFCC, 长度161)。
    • 不规则采样时间序列: CharacterTrajectories数据集,用于测试在数据缺失情况下的性能。
  • 模型架构:
    • LSSL: 主要使用两种尺寸:
      • LSSL small (约20万参数): 6层, H=128, N=128, M=1。
      • LSSL large (约200万参数): 4层, H=256, N=256, M=4。
    • 基线模型:
      • 长序列模型: CKConv (连续时间CNN), UnICORNN (ODE-RNN), Neural Controlled/Rough Differential Equations (NCDE/NRDE)。
      • 视觉模型: ResNet-18 (用于CelebA对比)。
      • 其他RNN、CNN基线模型。
  • 硬件与软件配置:
    • 硬件: 实验在GPU上进行,对于最大规模的实验(语音和高分辨率图像)使用了多GPU训练。
    • 软件:
      • 优化器: Adam优化器。
      • 学习率调度: 在验证集指标平台期(plateau)时学习率衰减5倍。
      • 实现: LSSL的实现利用了cuSPARSE库在GPU上进行高效的三对角矩阵求解。

实验结果

5.1 图像和时间序列基准测试

  • 实验内容: 在sMNIST, pMNIST, sCIFAR序列图像分类和BDIMC医疗回归任务上评估LSSL性能。
  • 实验结果:

    • 在sCIFAR上,LSSL取得了SOTA,准确率比之前最佳结果高出10个百分点以上 (Table 1)。
    • 在所有BDIMC数据集上,LSSL将均方根误差(RMSE)降低了超过三分之二 (Table 2)。
    • 所有这些结果都是在参数量比之前的SOTA少至少5倍的情况下实现的。

    表1: (逐像素图像分类。) (顶部) 本文方法。(中部) 循环基线。(底部) 卷积+其他基线。

    表2: (生命体征预测。) 预测呼吸率(RR)、心率(HR)和血氧(SpO2)的RMSE。*表示我们自己运行以完善最强基线的结果。

5.2 极长序列的语音和图像分类

  • 实验内容: 在原始长序列数据上测试LSSL的能力,包括长度为16000的原始语音信号(Speech Commands)和长度为38000的序列化图像(CelebA)。
  • 实验结果:

    • Speech Commands: 在原始信号上训练时,LSSL不仅比之前的方法准确率高出20个点以上,训练时间减少到1/5,而且其性能甚至超过了所有使用预处理后的短序列(长度161)的基线模型,证明了其自动学习特征的能力优于手工特征工程 (Table 4)。
    • Sequential CelebA: 小型LSSL模型的性能非常接近一个参数量是其10倍的专用ResNet-18图像分类架构,准确率差距仅为2.16个百分点 (Table 3)。这首次证明了通用序列模型处理此类任务的可能性。

    表3: (序列CelebA分类。)

    表4: (原始语音分类;时间尺度变换。) (顶部): 原始信号(长度16000);1 → f 表示测试时采样率变化为f倍。(底部): 先前工作中使用的预处理MFCC特征(长度161)。✗表示计算上不可行。

5.3 LSSL作为循环、卷积和连续时间模型的优势

  • 实验内容: 验证LSSL是否继承了三种模型范式的优点,包括收敛速度和时间尺度适应性。
  • 实验结果:

    • 收敛速度: LSSL结合了强大的序列数据归纳偏置(来自循环和NDE特性)和高效的并行训练(来自卷积视图),因此收敛速度极快。在所有基准测试中,LSSL-f达到先前SOTA性能所需的时间(无论是按epoch数还是挂钟时间衡量)都远少于之前的最佳模型 (Table 5)。
    • 时间尺度适应: 在测试时改变语音信号的采样率时,LSSL只需调整其$\Delta_t$参数即可适应,并且性能仍优于其他基线在没有时间尺度变化时的表现 (Table 4)。在CharacterTrajectories数据集上,LSSL在处理缺失数据时也表现出与最佳基线相当的竞争力 (Appendix F, Table 8)。

    表5: (LSSL的建模和计算优势。) 在每个基准类别中,我们比较了LSSL-f达到先前SOTA(PSoTA)结果以及接近SOTA目标所需的epoch数(ep.)。我们还报告了达到PSoTA相对于先前最佳模型的挂钟时间。

5.4 LSSL消融实验:学习记忆动态和时间尺度

  • 实验内容: 验证LSSL能够学习的关键参数A(记忆动态)和$\Delta_t$(时间尺度)对其性能的重要性。
  • 实验结果:
    • 记忆动态 A: 使用随机初始化的A矩阵的LSSL表现很差(例如,在pMNIST上准确率仅62%),证实了引入HiPPO理论的必要性。从LSSL-f(固定A)到完整的LSSL(可训练A),性能有一致的提升,表明学习记忆度量是有益的。
    • 时间尺度 $\Delta_t$: LSSL学习$\Delta_t$的能力是其对RNN门控机制的关键泛化。实验表明,如果$\Delta_t$设置不当,LSSL-f在sCIFAR上的准确率仅为49.3%。附录F中的额外结果显示,学习$\Delta_t$和学习A是两个正交的性能提升来源,并且$\Delta_t$的值在训练过程中会发生显著变化以更好地拟合数据。

A5 结论

本文介绍了一种受物理系统基本表示启发的简单且有原则的模型——线性状态空间层(LSSL)。理论和实验证明,LSSL泛化并继承了现代主流时间序列模型家族(循环、卷积、连续时间模型)的优点。通过结合关于连续时间记忆的新理论,LSSL解决了长期记忆这一主要限制,并在具有极长序列的困难任务上取得了显著的经验效果。

相关工作与模型简洁性:
LSSL与多条关于循环、卷积、连续时间模型以及长依赖序列模型的研究路线相关(详见附录A)。本文的模型非常简洁,由相同的LSSL层和简单的逐点非线性模块堆叠而成,对超参数不敏感,且能以比基线模型高得多的学习率进行训练。

局限性:
1. 推理效率: 尽管LSSL的循环表示理论上可以实现高效推理,但本文并未在实验中验证这一点。后续工作已证明在某些应用中确实可以加速推理。
2. 算法实用性: 定理2中提出的快速算法在后续研究中被发现数值不稳定,不适用于硬件。因此,该算法贡献主要作为理论上的概念验证,表明LSSL存在快速算法,但寻找实用、快速且数值稳定的算法仍是一个开放问题。
3. 空间效率: LSSL和LSSL-fixed都存在较大的空间开销,在处理长度为L的序列时需要 $O(NL)$ 而非 $O(L)$ 的空间,这源于使用维度为N的隐状态表示。因此,LSSL可能空间效率不高,在大型实验中需要多GPU训练。这些计算和空间复杂度的根本问题在本文的后续工作中得到了解决。

结论与未来工作:
现代深度学习模型在处理语音、视频和医疗时间序列等长时序数据时面临挑战。本文希望通过提出简单、有原则、工程化程度较低的模型,为这些应用带来新的可能性。像素级图像分类实验表明,无需特殊技巧的通用序列模型也能达到早期卷积网络的性能。语音实验则展示了学习优于手工特征的可能性。未来的潜在应用包括在预训练的状态空间特征之上训练其他下游模型。

A6 附录

A. 相关工作

与HiPPO的关系。LSSL与HiPPO框架【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】及其前身LMU【索引58,Legendre memory units: Continuous-time representation in recurrent neural networks,2019,NeurIPS】关系最密切。这些方法都定义了形如(1)的动力学并融入RNN架构,但存在两个主要限制:一是状态矩阵A和时间尺度$\Delta_t$因理论和计算限制而无法训练;二是(1)式是1-D到N-D的映射,需要将状态投影回1-D,造成了1-D的状态瓶颈,限制了模型表达力。LSSL通过保留线性循环(4)并用状态空间表示的第二部分(5)进行下投影,而非使用传统RNN架构。为避免1-D特征瓶颈,它独立计算H个1-D到1-D的映射副本,形成H维序列模型。本文通过新理论解决了表达力问题,允许通过展示HiPPO框架的广义理论结果来训练A矩阵,即存在一个参数化的结构化状态空间类是HiPPO算子。LSSL通过针对这些结构化矩阵的新算法(定理2)在计算问题上取得了进展,尽管该算法后来被发现不实用。

与连续时间CNN的关系。CKConv是我们所知的唯一连续时间CNN例子,也是我们实验中最强的基线之一。CKConv将卷积核参数化为一个从[0, 1]到$\mathbb{R}$的隐式函数,允许在任何分辨率下采样。其后继者FlexConv【索引43,Flexconv: Continuous kernel convolutions with differentiable kernel sizes,2021,arXiv】能学习宽度可变的卷积核,这与LSSL在使用某些HiPPO基时(第3.2节)的卷积解释相似。

与连续时间RNN的关系。RNN与连续时间模型的联系由来已久,近年来基于动力系统或ODE的CT-RNN模型激增。这些工作可分为几类:
1. 理论分析: 从连续时间角度分析RNN的表达力,如研究RNN与动力系统的对应关系【索引22,Approximation of dynamical systems by continuous time recurrent neural networks,1993,Neural networks】、稳定性和动力学【索引62,A comprehensive review of stability analysis of continuous-time recurrent neural networks,2014,IEEE Transactions on Neural Networks and Learning Systems;索引29,Gated recurrent units viewed through the lens of continuous time dynamical systems,2019,Frontiers in computational neuroscience】。
2. 解决梯度消失: 从动力系统分析出发设计RNN以对抗梯度消失,如AntisymmetricRNN【索引7,Antisymmetricrnn: A dynamical system view on recurrent neural networks,2019,ICLR】、iRNN【索引30,Rnns incrementally evolving on an equilibrium manifold: A panacea for vanishing and exploding gradients?,2020,ICLR】、LipschitzRNN【索引20,Lipschitz recurrent neural networks,2021,ICLR】。
3. 基于显式ODE的模型: 为满足特定性质而引入显式底层ODE,如UnICORNN【索引47,Unicornn: A recurrent model for learning very long time dependencies,2021,ICML】及其前身coRNN【索引46,Coupled oscillatory recurrent neural network (cornn): An accurate and (gradient) stable architecture for learning long time dependencies,2021,ICLR】,它们离散化一个二阶ODE。其他模型包括LTC【索引27,Liquid time-constant networks,2021,AAAI】及其后继CfC【索引26,Closed-form continuous-depth models,2021,arXiv】,以及LMU【索引58,Legendre memory units: Continuous-time representation in recurrent neural networks,2019,NeurIPS】和HiPPO【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS】。
4. 神经ODE家族: 神经ODE【索引10,Neural ordinary differential equations,2018,NeurIPS】被应用于连续时间,催生了ODE-RNN【索引45,Latent ordinary differential equations for irregularly-sampled time series,2019,NeurIPS】、GRU-ODE-Bayes【索引16,Gru-ode-bayes: Continuous modeling of sporadically-observed time series,2019,NeurIPS】、ODE-LSTM【索引32,Learning long-term dependencies in irregularly-sampled time series,2020,arXiv】等模型。NCDE【索引31,Neural controlled differential equations for irregular time series,2020,arXiv】和NRDE【索引37,Neural rough differential equations for long time series,2021,ICML】是其内存高效版本,适用于极长时序。

门控机制。一些工作已观察到门控机制与阻尼动力系统的关系【索引54,Can recurrent neural networks warp time?,2018,ICLR】,例如LTC【索引27,Liquid time-constant networks,2021,AAAI】和iRNN【索引30,Rnns incrementally evolving on an equilibrium manifold: A panacea for vanishing and exploding gradients?,2020,ICLR】。相比之下,引理3.1展示了一个更强的结果:sigmoid门控不仅是受任意值域为(0,1)的单调函数启发,其精确形式直接源于对阻尼ODE的离散化。

B. 模型细节

B.1 (M)LSSL计算

计算优化。第3.1节指出LSSL的某些计算成本高昂。当LSSL固定A和$\Delta_t$参数时(例如,在不训练它们或在推理时),可以通过缓存特定计算来规避这些困难。这种情况适用于LSSL-f。此时,其他状态空间矩阵C和D构成了固定转移LSSL的 $O(HN)$ 个可训练参数。我们假设存在一个该系统的黑盒推理算法,即与$\bar{A}$的矩阵向量乘法(附录E.2给出了一个特定结构化类的实现示例)。然后我们计算并缓存:
- 转移矩阵 $\bar{A}$:通过将黑盒$\bar{A}$ MVM算法应用于单位矩阵I来计算。
- Krylov矩阵 K($\bar{A}, \bar{B}$):


该矩阵通过指数化的平方技术(即批量乘以$\bar{A}, (\bar{A})^2, (\bar{A})^4, \dots$)以并行化方式计算。

训练与推理。在推理时,模型可以用$\bar{A}$进行循环展开。在训练时,卷积核 $K_L(\bar{A}, \bar{B}, C)$ (方程(7))通过矩阵乘法 $C \cdot K(\bar{A}, \bar{B})$ 计算,然后与输入u进行卷积。表7提供了这个固定A, $\Delta_t$版本的LSSL的更详细复杂度。如第6节所述,这种缓存算法相当快,但主要缺点是物化Krylov矩阵(8)需要 $O(NL)$ 而非 $O(L)$ 的空间。

B.2 A的初始化

HiPPO-LegS初始化。LSSL将(1)中的参数A初始化为HiPPO-LegS算子,该算子是为了解决一个特定的连续时间记忆问题而推导出来的。这个矩阵 $A \in \mathbb{R}^{N \times N}$ 是:


注意,LSSL-f是具有不可训练A(和$\Delta_t$)的LSSL,因此其A固定为上述矩阵。

B.3 $\Delta_t$的初始化

多时间尺度初始化。LSSL与最相关的先前工作的一个区别是,包含了投影(2)使得该层成为一个1维到1维的映射,而不是1-D到N-D【索引24,Hippo: Recurrent memory with optimal polynomial projections,2020,NeurIPS;索引58,Legendre memory units: Continuous-time representation in recurrent neural networks,2019,NeurIPS】。这使我们能够连接H个该映射的副本(以计算为代价,参见4.2节和附录D)。即使$\Delta_t$在LSSL-f中不被训练,这H个副本也允许通过为每个副本设置不同的$\Delta_t$来考虑多个时间尺度。具体来说,我们在一个范围 $[\Delta_{t_{min}}, \Delta_{t_{max}}]$ 内对$\Delta_t$进行对数均匀初始化(即,$\Delta_t$在此范围内初始化,使得$\log\Delta_t$是均匀分布的)。最大值和最小值通常选择相差100倍,以便数据集中序列的长度包含在此范围内。每个模型和数据集的具体值见附录F。我们没有将其作为超参数进行搜索,但注意到调整它可以为我们的实验带来额外的性能提升。

B.4 深度神经网络架构

整体架构。我们实验中使用的深度LSSL模型只是将LSSL层在一个简单的深度神经网络架构中堆叠起来。我们注意以下架构细节:
- 通道(Channels):状态空间模型(1)+(2)接受一维输入u,但不必严格返回一维输出y。通过使(2)中的矩阵维度为 $C \in \mathbb{R}^{M \times N}, D \in \mathbb{R}^{M \times 1}$,输出y的维度将是M而不是1。我们称M为模型中的通道数。
- 前馈层(Feedforward):当前LSSL定义存在两个缺点:1) H个输入特征不相互作用;2) 如果通道维度M>1,则无法应用残差连接。我们通过在LSSL之后引入一个形状为 $H \cdot M \to H$ 的逐点前馈层来解决这两个问题。这同时混合了隐藏特征,并在必要时将输出投影回维度H。在LSSL和这个前馈投影之间还可以选择性地加入非线性;我们在模型中固定使用GeLU激活函数。这种H个特征上的并行卷积后跟一个逐点线性映射的因式分解与深度可分离卷积【索引13,Xception: Deep learning with depthwise separable convolutions,2017,CVPR】非常相似。
- 残差和归一化(Residuals and normalization):为堆叠多层LSSL,我们使用非常标准的深度神经网络架构,即残差连接和层归一化(前置或后置),风格类似于标准的Transformer架构。使用前置还是后置归一化根据数据集和模型是否过拟合来选择。
- 参数量(Parameter count):一个LSSL模型的总参数量为 $M \cdot H \cdot (H+N)$。我们主要使用了两种模型尺寸:
- LSSL small (约200K参数): 6层, H=128, N=128, M=1。
- LSSL large (约2M参数): 4层, H=256, N=256, M=4。

... (由于篇幅限制,附录C, D, E, F的详细缩写内容在此省略,其核心思想已在正文和摘要中体现。这些附录主要包含详细的数学证明、算法推导和实验设置,遵循方法细节的缩写原则进行处理。)