手写 Mamba 模型：从状态空间模型到选择性 SSM 的完整实现

一、引言Transformer 之后的下一站自 2017 年 Attention Is All You Need 发表以来Transformer 架构统治了自然语言处理、计算机视觉、音频处理等几乎所有序列建模领域。其核心机制——自注意力Self-Attention——通过计算序列中任意两个位置之间的注意力权重实现了强大的长距离依赖建模能力。然而这个看似完美的设计有一个根本性缺陷计算复杂度与序列长度呈二次关系 O(L²)。当序列长度从 512 增长到 128K 时自注意力的计算量暴增 64000 倍。这让 Transformer 在处理超长序列如整本书、完整基因组、高分辨率视频帧时面临巨大的计算开销和内存瓶颈。FlashAttention、稀疏注意力等技术虽然在一定程度上缓解了这个问题但并未从根本上改变二次复杂度的本质。2023 年底Albert Gu 和 Tri Dao 联合提出了Mamba架构引起了深度学习社区的广泛关注。Mamba 的核心是选择性状态空间模型Selective State Space Model, Selective SSM它将序列建模的计算复杂度降到了线性 O(L)同时在语言建模等任务上达到了媲美甚至超越 Transformer 的性能。更令人惊叹的是Mamba 的推理阶段具有类似于 RNN 的状态压缩特性——生成时只需维护一个固定大小的隐藏状态而非像 Transformer 那样需要缓存所有历史 Key-Value 对。这意味着 Mamba 的推理吞吐量可以比同规模 Transformer 高出数倍。2024 年同一团队又推出了Mamba-2通过结构化状态空间对偶Structured State Space Duality, SSD理论将 SSM 与注意力机制统一在了一个数学框架下进一步提升了训练速度和性能。本文将从零开始带你一步步推导并实现 Mamba 的核心组件。我们将从经典控制理论中的状态空间模型出发理解 S4Structured State Space Sequence Model的离散化和卷积表示再到 Mamba 的三大创新选择性机制、硬件高效扫描算法、以及简化的架构设计。最后我们会用 PyTorch 手写一个可运行的 mini Mamba 实现。二、状态空间模型从控制理论到深度学习2.1 连续系统的状态空间表示状态空间模型State Space Model, SSM起源于控制理论用于描述动态系统的输入-输出行为。一个连续线性时不变系统可以用以下微分方程组描述h(t) A · h(t) B · x(t) y(t) C · h(t) D · x(t)其中-x(t) ∈ ℝ^D输入信号如文本的 token 嵌入-h(t) ∈ ℝ^N隐藏状态系统内部的状态N 是状态维度-y(t) ∈ ℝ^D输出信号-A ∈ ℝ^{N×N}状态转移矩阵描述了状态如何随时间演化-B ∈ ℝ^{N×D}输入投影矩阵输入如何影响状态-C ∈ ℝ^{D×N}输出投影矩阵状态如何映射到输出-D ∈ ℝ^{D×D}前馈连接通常可忽略或作为残差连接第一个方程称为状态方程它描述了隐藏状态随时间的连续演化第二个方程称为输出方程它将隐藏状态映射回输出空间。直觉理解可以把状态空间模型想象成一个水箱系统。x(t) 是流入水箱的水流h(t) 是水箱中的水量A 决定了水自然会蒸发/渗漏的速率B 决定了流入的水中有多少真正进入水箱C 决定了水箱中的水量如何影响下游如出水口的水压。SSM 的美妙之处在于它用一个低维状态来总结整个输入序列的历史。2.2 离散化从连续到离散深度学习处理的是离散的 token 序列因此我们需要将连续时间 SSM 转化为离散时间形式。常用的方法是零阶保持Zero-Order Hold, ZOH离散化其核心假设是在每个时间步内输入信号 x(t) 保持恒定。给定步长 Δ可理解为采样间隔或token 的时间分辨率ZOH 离散化的结果是h_k exp(Δ·A) · h_{k-1} (Δ·A)^{-1}(exp(Δ·A) - I) · Δ·B · x_k这看起来有点复杂但我们可以通过定义离散化参数来简化Ā exp(Δ·A) B̄ (Δ·A)^{-1}(exp(Δ·A) - I) · Δ·B那么离散 SSM 就变成了h_k Ā · h_{k-1} B̄ · x_k y_k C · h_k D · x_k在实际实现中为了简化计算常用前向欧拉法Euler method的一阶近似h_k (I Δ·A) · h_{k-1} Δ·B · x_k不过使用矩阵指数的 ZOH 离散化理论精度更高且对于某些需要长时间记忆的任务表现更好。2.3 SSM 的两种计算模式离散 SSM 可以以两种模式进行计算分别对应训练和推理阶段。递推模式Recurrent Modefor k in range(L): h_k Ā · h_{k-1} B̄ · x_k y_k C · h_k这类似于 RNN计算复杂度是 O(L·N²)其中 L 是序列长度N 是状态维度。在推理时自回归生成这种模式非常高效——只需维护一个 N 维的状态向量每步计算量恒定。卷积模式Convolutional Mode通过展开递推公式我们可以发现输出 y_k 实际上是输入 x 与一个SSM 卷积核的卷积。以初始状态 h_{-1}0 为例y_0 C · B̄ · x_0 y_1 C · Ā · B̄ · x_0 C · B̄ · x_1 y_2 C · Ā² · B̄ · x_0 C · Ā · B̄ · x_1 C · B̄ · x_2 ...可以写成y K * x, 其中 K (C · Ā^k · B̄)_{k0}^{L-1}K 是一个长度为 L 的卷积核每个元素是一个 D×D 矩阵。这种模式可以利用 FFT 在 O(L·log L) 时间内并行计算非常适合训练时处理完整的序列。这种递推-卷积对偶是 SSM 的核心优势之一训练时用卷积看完整序列做并行计算推理时用递推做高效自回归生成。三、S4结构化状态空间序列模型2019 年到 2022 年间研究者们将 SSM 引入深度学习序列建模形成了 S4Structured State Space Sequence Model及其变体如 DSS, S5, H3 等。S4 的核心贡献是解决了 SSM 在长序列上的两个关键问题。3.1 问题一HiPPO 初始化随机初始化的 A 矩阵无法捕获长距离依赖。S4 的解决方案是使用HiPPOHigh-order Polynomial Projection Operators矩阵来初始化 A。HiPPO 矩阵本质上是对一个在线勒让德多项式逼近过程的离散化它能让隐藏状态 h_k 有效地存储整个输入序列的历史信息。最常用的 HiPPO 矩阵是 LegSLegendre State形式A_{ij} -((2i1)(2j1))^{1/2} · min(i,j) mod 2 * ...简化形式实际上HiPPO 矩阵的一个重要性质是它产生的状态维度可以看成是用勒让德多项式对输入历史进行投影低维分量存储近期信息高维分量存储远期信息。3.2 问题二结构化的高效计算直接使用 N×N 的密集矩阵 AN 通常为 64-256在计算上效率低下。S4 的关键洞察是通过特定的参数化A 可以表示为对角矩阵 低秩矩阵或更简单的对角矩阵的形式从而将矩阵乘法复杂度从 O(N²) 降低到 O(N)。S4 通过Lanczos 算法和柯西矩阵求逆等数值技巧将 SSM 的核 K 的计算从 O(N²L) 优化到 O(NL)。在后续的简化版 S4D 中作者直接将 A 参数化为对角矩阵A diag(a_1, a_2, ..., a_N)这就意味着 Ā 也是对角阵Ā diag(exp(Δ·a_1), exp(Δ·a_2), ..., exp(Δ·a_N))这样状态更新变成了逐元素的乘法h_k^{(i)} exp(Δ·a_i) · h_{k-1}^{(i)} Δ·b_i · x_k3.3 从 S4 到 Mamba 的桥梁S4 及其变体在长序列任务如 LRA 基准测试、音频波形建模上取得了巨大成功但在语言建模上却不敌 Transformer。原因在于S4 的 A, B, C 参数是固定的时不变/线性时不变系统这意味着同一个卷积核对所有输入 token 一视同仁——它无法像注意力机制那样根据输入内容选择性关注特定 token。这就引出了 Mamba 的核心创新点。四、Mamba 的核心创新选择性状态空间模型Mamba 在 2023 年底横空出世它对传统 SSM 做了三方面的关键改进。4.1 选择性机制Selection MechanismMamba 的核心洞见是让 SSM 的参数依赖于输入。具体来说Mamba 将 B, C, Δ 从固定参数变为输入的函数B(x) linear_B(x) # ℝ^D → ℝ^{N×D} C(x) linear_C(x) # ℝ^D → ℝ^{D×N} Δ(x) softplus(linear_Δ(x)) # ℝ^D → ℝ^D 正值而 A 矩阵仍然保持固定用 HiPPO 初始化这样状态转移的结构就保留了但输入和输出的门控由内容决定。为什么选择性机制如此重要让我们通过一个经典的选择性复制任务来说明假设输入是 aaa...aabaaa...aa其中只有一个 b。任务要求输出所有 a 的数量。对于一个固定的 SSMS4它无法区分重要和不重要的 token ——所有的 a 都以同样的方式被记住。而具备选择性的 Mamba 可以在遇到 b 时通过 Δ 参数动态调整状态更新的门控速率对 b 之前的 token 快速遗忘对 b 之后的新信息重点编码。从数学上讲Δ 控制了离散化的步长。当 Δ 较大时状态更新更快更关注当前输入当 Δ 较小时状态更平稳更多保留历史。这种输入依赖的 Δ 让模型学会了什么时候记住什么时候忘记。4.2 硬件高效算法并行扫描Parallel Scan选择性机制破坏了一个重要的性质卷积等价性。因为 B 和 C 现在依赖于输入不同位置的卷积核不同无法再使用全局的 FFT 卷积加速。最直接的实现方式退化为串行的递推扫描Recurrent Scan复杂度 O(L·N²)效率极低。Mamba 的第二个关键创新是使用并行前缀和扫描算法Parallel Prefix Sum / Associative Scan来高效计算 SSM 的前向传播具体来说就是双调扫描Blelloch Scan。核心思想是SSM 递推公式h_k Ā_k · h_{k-1} B̄_k · x_k实际上是一个结合性运算。我们将状态更新表示为二元运算符 ⊕(h_{k-1}, x_k) ⊕ x_{k1} (h_k, x_{k1}), 其中 h_k Ā_k·h_{k-1} B̄_k·x_k注意这里 Ā_k exp(Δ_k·A)B̄_k Δ_k·B(x_k)每个时间步的 Ā_k 和 B̄_k 都不同。由于 ⊕ 是结合性的我们可以使用树形归约Tree Reduction来并行计算扫描阶段Down-sweep将序列分成两半分别计算组内的累积转换合并阶段Up-sweep将相邻组的累积转换合并这样计算复杂度就从 O(L) 的串行扫描降低到了 O(log L) 的并行树扫描。在实际的 CUDA 实现中Mamba 将序列分块block每个 block 内做串行扫描block 之间做并行扫描充分利用 GPU 的共享内存和线程层级。这种硬件感知的并行扫描是 Mamba 训练效率的关键。4.3 简化的网络架构除了 SSM 本身的改进Mamba 还设计了一个简洁的整体架构。与 Transformer 的复杂结构Multi-Head Attention Feed-Forward带 Layer Norm 和 Residual不同Mamba 的每个块仅包含线性投影将输入投影到 2×ED 维度ED 扩展维度通常为 2D~4D1D 卷积在 SSM 之前加一个简单的因果卷积kernel size4帮助模型感知局部上下文SiLU 激活使用 Swish/SiLU 激活函数选择性 SSM输入依赖的 SSM 层残差连接 LayerNorm实际上Mamba 块的结构类似于门控 MLP SSM的组合def mamba_block(x): residual x x norm(x) # 扩展 门控 x_proj linear_up(x) # D → 2*ED x_conv silu(conv1d(x_proj)) # 局部卷积 激活 x_gate, x_state chunk(x_conv, 2, dim-1) # 选择性 SSM y selective_ssm(x_state) # 门控输出 y y * x_gate y linear_down(y) # ED → D return y residual这种简化设计让 Mamba 可以堆叠到很大深度例如 2.8B 参数、64 层而不会出现训练不稳定问题。五、手写实现Mini Mamba 从零开始现在让我们用 PyTorch 实现一个简化但可运行的 Mini Mamba。我们会实现完整的 SSM 前向传播包括选择性机制和并行扫描。5.1 SSM 核心实现首先实现离散化和选择性 SSM 的前向传播import torch import torch.nn as nn import torch.nn.functional as F import math class SelectiveSSM(nn.Module): 选择性状态空间模型核心 参数: d_model: 模型维度 D d_state: SSM 状态维度 N (通常 16-64) dt_rank: Δ投影的中间维度 def __init__(self, d_model, d_state16, dt_rankauto): super().__init__() self.d_model d_model self.d_state d_state # dt_rank 决定 Δ 的瓶颈维度 self.dt_rank math.ceil(d_model / 16) if dt_rank auto else dt_rank # A 矩阵固定的 HiPPO 初始化对角形式 # 这里用 log 空间参数化确保 A 的实部为负稳定系统 self.A_log nn.Parameter( torch.log(torch.arange(1, d_state 1, dtypetorch.float32)) .unsqueeze(0).repeat(d_model, 1) ) # 可学习的 D 参数Skip connection self.D nn.Parameter(torch.ones(d_model)) # 输入依赖的投影层 self.dt_proj nn.Linear(self.dt_rank, d_model, biasTrue) self.B_proj nn.Linear(d_model, d_state, biasFalse) self.C_proj nn.Linear(d_model, d_state, biasFalse) # dt 的中间编码x → dt_rank → dt通过 dt_proj self.dt_enc nn.Linear(d_model, self.dt_rank, biasFalse) def forward(self, x): x: (batch, seq_len, d_model) 返回: (batch, seq_len, d_model) batch, seq_len, d_model x.shape # 1) 计算输入依赖的参数 dt_enc self.dt_enc(x) # (B, L, dt_rank) delta F.softplus(self.dt_proj(dt_enc)) # (B, L, d_model) B self.B_proj(x) # (B, L, d_state) C self.C_proj(x) # (B, L, d_state) # 2) 离散化ZOH 对角简化 # A -exp(A_log)保证负实部 → 稳定系统 A -torch.exp(self.A_log.float()) # (d_model, d_state) # 对于对角 AZOH 离散化简化为元素乘法 # Ā exp(Δ · A)B̄ Δ · B对角 A 的近似 delta delta.unsqueeze(-1) # (B, L, d_model, 1) A_expanded A.unsqueeze(0).unsqueeze(0) # (1, 1, d_model, d_state) # Ā exp(Δ * A) 逐元素 delta_A delta * A_expanded # (B, L, d_model, d_state) A_bar torch.exp(delta_A) # (B, L, d_model, d_state) # B̄ Δ * B 简化为元素乘法 B_expanded B.unsqueeze(2) # (B, L, 1, d_state) delta_B delta * B_expanded # (B, L, d_model, d_state) B_bar delta_B # B̄ ≈ Δ·B ZOH 一阶近似 # 注意对于严格的 ZOH需要计算 (e^{ΔA} - I)A^{-1}ΔB # 这里做了一阶简化。完整版本可以参考官方实现。 # 3) 并行扫描递推扫描实现 # h_t Ā_t · h_{t-1} B̄_t · x_t # y_t C_t · h_t D · x_t # 初始化状态 h torch.zeros(batch, d_model, self.d_state, devicex.device) outputs [] # 转为 (d_model, d_state) 维度 # 注意这里的 x 在 SSM 空间中的维度是 d_model for t in range(seq_len): # x_t: (B, d_model) x_t x[:, t, :] # (B, d_model) # 状态更新 h_t Ā · h_{t-1} B̄ · x_t # A_bar[t]: (B, d_model, d_state), B_bar[t]: (B, d_model, d_state) # x_t 需要扩展到 (B, d_model, 1) 与 B_bar 做逐元素乘 h A_bar[:, t, :, :] * h B_bar[:, t, :, :] * x_t.unsqueeze(-1) # 输出 y_t C_t · h_t D · x_t # C[t]: (B, d_state), h: (B, d_model, d_state) y_t torch.einsum(bds,bs-bd, h, C[:, t, :]) # (B, d_model) y_t y_t self.D * x_t outputs.append(y_t) return torch.stack(outputs, dim1) # (B, L, d_model)上面的循环实现虽然直观但效率较低。实际训练中应该使用 CUDA 级别的并行前缀扫描。我们稍后会实现一个更高效的版本。5.2 完整 Mamba 块实现有了 SelectiveSSM我们可以组合完整的 Mamba 块class MambaBlock(nn.Module): 完整的 Mamba 块 架构: Norm → LinearUp → Conv1d → SiLU → SelectiveSSM → Gating → LinearDown → Residual def __init__(self, d_model, d_state16, d_conv4, expand_factor2, dt_rankauto): super().__init__() self.d_model d_model self.d_state d_state self.d_conv d_conv self.expand_factor expand_factor self.d_inner int(expand_factor * d_model) # 输入归一化 self.norm nn.LayerNorm(d_model) # 投影升维到内部维度 self.in_proj nn.Linear(d_model, self.d_inner * 2, biasFalse) # 因果 1D 卷积 self.conv1d nn.Conv1d( in_channelsself.d_inner, out_channelsself.d_inner, kernel_sized_conv, paddingd_conv - 1, groupsself.d_inner, biasTrue, ) # 激活函数 self.act nn.SiLU() # 选择性 SSM self.ssm SelectiveSSM(self.d_inner, d_state, dt_rank) # 输出投影降维回 d_model self.out_proj nn.Linear(self.d_inner, d_model, biasFalse) def forward(self, x): x: (batch, seq_len, d_model) residual x x self.norm(x) # 1) 升维 门控拆分 x_proj self.in_proj(x) # (B, L, 2*d_inner) x_proj x_proj.transpose(1, 2) # (B, 2*d_inner, L) x_conv, x_gate x_proj.chunk(2, dim1) # 各 (B, d_inner, L) # 2) 因果卷积 激活 x_conv self.conv1d(x_conv)[:, :, :x.shape[1]] # 裁掉 padding x_conv self.act(x_conv) # (B, d_inner, L) # 3) SSM 处理 x_ssm x_conv.transpose(1, 2) # (B, L, d_inner) x_ssm self.ssm(x_ssm) # (B, L, d_inner) x_ssm x_ssm * self.act(x_gate.transpose(1, 2)) # 门控乘 SSM 输出 # 4) 降维输出 y self.out_proj(x_ssm) # (B, L, d_model) return y residual5.3 并行前缀扫描实现上面的串行扫描只适合理解和调试。真正的 Mamba 训练需要一个并行前缀扫描associative scan实现。这里给出一个基于递推展开的 batched 版本def selective_ssm_parallel(A_bar, B_bar_x, C): 并行扫描实现选择性 SSM 参数: A_bar: (batch, seq_len, d_model, d_state) 状态转移矩阵 B_bar_x: (batch, seq_len, d_model, d_state) 输入投影 C: (batch, seq_len, d_state) 输出投影矩阵 返回: y: (batch, seq_len, d_model) batch, seq_len, d_model, d_state A_bar.shape # 关联扫描h_t A_t * h_{t-1} Bx_t # 使用前缀和递推展开 # 方式1递推展开适合短序列 h torch.zeros(batch, d_model, d_state, deviceA_bar.device) outputs [] for t in range(seq_len): h A_bar[:, t] * h B_bar_x[:, t] # (B, d_model, d_state) y_t torch.einsum(bds,bs-bd, h, C[:, t]) # (B, d_model) outputs.append(y_t) return torch.stack(outputs, dim1) # 方式2使用 torch.associative_scanPyTorch 2.2 实验性 API # 实际训练中推荐集成 CUDA selective_scan 算子官方 Mamba 实现使用了自定义的 CUDA kernelselective_scan_fwd它将序列分块到 GPU blocks每个 block 内做扫描block 之间用并行前缀扫描合并用共享内存存储扫描中的中间结果为反向传播保存必要的中间状态六、Mamba-2结构化状态空间对偶2024 年中同一团队推出了 Mamba-2。它的核心贡献是结构化状态空间对偶SSD从理论上统一了 SSM 和注意力机制。6.1 SSD 理论SSM Attention 的另一种形式SSD 的核心发现是当 SSM 的状态维度 N1 时SSM 的输出等价于一种线性注意力Linear Attention。让我们来看一下这个对偶关系。当 A 是标量时d_state1SSM 的递推变为h_k a_k · h_{k-1} b_k · x_k y_k c_k · h_k展开后y_k c_k · Σ_{i0}^{k} (∏_{ji1}^{k} a_j) · b_i · x_i写为矩阵形式y (下三角矩阵) · (输入), 其中下三角矩阵的元素是 M[k, i] c_k · (∏ a_j) · b_i (for i ≤ k)这其实就是一种具有结构化掩码的注意力机制——权重矩阵是低秩的rank-1 分解且呈下三角形式。更一般地当 N 1 时SSM 相当于一种结构化状态空间注意力其权重矩阵的秩受 N 限制。Mamba-2 利用这个对偶关系在训练时可以使用矩阵乘法而不是扫描算法来计算 SSM 的前向传播这使得 Mamba-2 可以利用高度优化的 cuBLAS 矩阵乘法库在 GPU 上获得更好的硬件利用率。6.2 Mamba-2 的实际改进在实际架构上Mamba-2 做了以下简化去掉 1D 卷积SSD 理论让 SSM 本身就能捕获足够的信息不再需要额外的局部卷积简化的参数化A 直接使用标量N1B 和 C 也采用了更简洁的形式Grouped SSM类似于分组注意力GQA的思路多个头共享 SSM 参数更大的状态维度由于计算效率提升可以支持更大的 d_state实验表明Mamba-22.8B在语言建模上与同等规模的 Transformer 持平在长上下文任务上甚至略有优势而训练速度比 Mamba-1 快了 2-3 倍。七、Mamba 的现状与展望7.1 当前生态截至 2025-2026 年Mamba 已经在以下领域展现出强大的竞争力语言建模Mamba-2 在 Pile、Dolma 等基准上的困惑度perplexity已接近 LLaMA 级别的 Transformer长上下文Mamba 对 256K 的超长序列可以不做任何特殊处理直接推理视觉Vision MambaVim、VMamba 将 SSM 应用于图像分类、分割、检测多模态Cobra、Video Mamba 等扩展到视觉语言模型基因组学借助线性复杂度直接在完整染色体序列上训练音频Mamba 在语音识别、音乐生成等任务上表现优异7.2 挑战与局限性Mamba 并非万能解决方案训练生态不成熟缺乏像 HuggingFace 那样完善的训练工具链自定义 CUDA kernel 维护成本高在小规模上优势不明显当序列长度 1024 时Mamba 的性能优势有限涌现能力目前还没有足够大规模的 Mamba 模型如 70B来验证涌现能力混合架构的趋势越来越多的研究表明Mamba Attention 混合如 Jamba、Samba可能是最优解7.3 实践建议什么时候选择 Mamba- 需要处理超长序列16K tokens的场景- 推理吞吐量要求极高的部署场景- 需要低延迟的流式处理什么时候选择混合架构- 需要顶尖的语言理解和生成质量- 预算充足可以采用 Attention 层 Mamba 层的混合堆叠八、总结本文从状态空间模型的基本原理出发逐步推导了从连续 SSM、离散化、S4 到 Mamba 选择性 SSM 的完整技术路线。我们用 PyTorch 手写了一个可运行的 Mini Mamba 实现涵盖了选择性机制、输入依赖参数化和硬件感知扫描的核心思想。Mamba 的核心贡献可以概括为1.线性计算复杂度O(L) 而非 O(L²)理论上可以处理任意长度的序列2.内容感知的选择性通过学习何时记住、何时遗忘解决了固定 SSM 在语言建模中的根本痛点3.硬件感知实现通过并行扫描算法将理论优势转化为实践中的速度提升4.统一理论视角Mamba-2 的 SSD 理论证明了 SSM 与注意力的内在联系Mamba 代表了一种新的序列建模范式——它告诉我们Attention 并非序列建模的唯一选择线性复杂度的结构化状态空间模型同样可以达到顶尖性能。随着训练基础设施的完善和混合架构的成熟SSM 系列的模型将在 AI 基础设施、长序列处理、边缘计算等场景中扮演越来越重要的角色。 实战指南推荐如果你想进一步了解如何将 Mamba、DeepSeek、Dify 等模型和工具结合起来构建完整的 AI 应用请参考 DeepSeek Dify 智能体全栈开发实战指南涵盖从模型部署、Prompt 工程到 Agent 工作流编排的完整教程。参考文献Gu, A., Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces.Gu, A., et al. (2024). Mamba-2: Transformers are SSMs, and SSMs are Transformers.Gu, A., et al. (2022). Efficiently Modeling Long Sequences with Structured State Spaces (S4).Gu, A., et al. (2020). HiPPO: Recurrent Memory with Optimal Polynomial Projections.Smith, J., et al. (2023). Simplified State Space Layers for Sequence Modeling (S5).Dao, T., et al. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms (SSD).