跳转至

Fractional State Space Transition for Long Sequence Modeling

会议: NeurIPS2026
arXiv: 2609.36314
代码: https://github.com/anasiri/frac-ssm
领域: LLM 高效推理 / 长序列建模
关键词: 分数阶动力学、状态空间模型、长上下文、指数和近似、选择性读写

一句话总结

Frac 用共享几何时间尺度上的有限指数模式近似分数阶长记忆,再加入逐 token 控制和独立读写路由,在保持有界递归状态的同时将 1.3B 模型的 LongBench 平均分从 GDN 的 16.0 提高到 17.9,但并非所有任务或速度指标都领先。

研究背景与动机

状态空间模型(SSM)把历史压入固定大小的状态,不必像自注意力一样持续保存整段上下文;代价是读者必须问:历史信息到底按照什么规律被遗忘?Mamba 类方法通常从常微分方程(ODE)及其离散化出发,单个模式产生指数衰减。选择性门控能够决定当前内容是否值得留下,却不能自动保证在跨多个数量级的距离上都保有合适的记忆权重。对于稀疏事件的长期累积,仅有某个特征时间尺度往往不够。

分数阶微分方程(FDE)提供另一种起点:状态变化依赖过去的加权轨迹,其 Mittag–Leffler 核具有渐近多项式尾部。但直接求解它需要访问历史,恰好破坏 SSM 希望保留的有限状态和高效解码。因此,本文不是给现有 SSM 加一个更大的历史缓存,而是把连续的多时间尺度记忆近似成有限个普通递归状态。

这一切入点也限定了结论的含义:有限指数和是在选定时间范围内近似重尾行为,不是在无限时间上把指数系统变成精确幂律;之后加入的 token 选择性也超出了固定系数 FDE 的严格近似定理。核心 idea:用有理论来源的几何多尺度记忆库规定遗忘几何,再让内容决定写入和读出的模式,而不是只在单一指数遗忘机制上优化门控。

方法详解

整体框架

Frac layer 接收隐藏表示,经过归一化和输入投影后进入 FracMixer,再由门控归一化和输出投影返回更新表示。Mixer 内部保留两条来源不同的分支:卷积前的信号生成控制变量,局部深度可分离因果卷积后的信号提供待写入内容。控制变量作用于共享时间尺度库,读写权重则把内容分配到这些模式并组合其输出。

具体流程是:几何多尺度记忆库提供从快到慢的基础时间尺度;区间冻结控制计算当前 token 的保留率;先验读写路由决定信息在哪些模式存取;仿射扫描实现完成状态递推和加权读出。读出还叠加可学习的直接馈通项 \(D u_t\),使局部内容不必全部经过长期状态。训练和预填充使用分块并行计算,解码则缓存模式状态并执行同一递推的单步形式。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    X["隐藏表示<br/>归一化与输入投影"] --> C["区间冻结控制"]
    B["几何多尺度记忆库"] --> C
    X --> U["局部因果卷积<br/>内容信号"]
    C -->|保留率与注入因子| R["先验读写路由"]
    U --> R
    R --> S["仿射扫描实现"]
    P["前一模式状态"] --> S
    S --> O["读出与直接馈通<br/>门控归一化及输出投影"]
    U -->|直接馈通| O

图中实线表示计算数据流;训练目标在下文说明,不是额外的推理模块。记忆库是共享基础参数,并不是从当前输入重新生成的历史缓存。

关键设计

1. 几何多尺度记忆库:把非局部历史变成有限个递归模式

理论起点是 Caputo 分数阶松弛方程,其中固定阶数和固定衰减参数决定连续系统的记忆规律。这里必须区分两种响应:无外部输入时,初始状态按 Mittag–Leffler 函数衰减;有外部输入时,历史内容由脉冲响应核累积。二者都是重尾,但指数不同,不能统称为同一个幂律。

\[ {}^{C}D_t^{\alpha}h(t)=-\lambda h(t)+u(t),\qquad E_{\alpha}(-\lambda t^{\alpha})\sim\frac{t^{-\alpha}}{\lambda\Gamma(1-\alpha)},\qquad g_{\alpha,\lambda}(t)\sim\frac{\alpha\,t^{-\alpha-1}}{\lambda^2\Gamma(1-\alpha)}. \]

上述渐近关系要求固定 \(0<\alpha<1\)、\(\lambda>0\) 且 \(t\to\infty\);Mittag–Leffler 核并非在所有距离上都等于精确幂律。原文定理 1 把两种核表示成非负指数衰减的连续混合,定理 2 再截断时间尺度范围并在对数坐标下做求积。结果是一组共享的几何时间尺度 \(\tau_m=\tau_0q^{m-1}\),每个模式只需维护自己的第一阶状态。

为什么必须使用多尺度而不是单个慢衰减?单个模式只有一个主要记忆跨度;几何间隔把有限模式分配到多个数量级,使快模式处理局部变化、慢模式覆盖远距离依赖。实现采用 \(M=16\),基础时间尺度从 \(1\) 到 \(2^{17}\),而不是为每个新长度追加状态。该基础范围超过语言模型实验的最长评估长度,但实际有效时间尺度还会被控制变量缩放,因此端点覆盖本身不等于所有输入上的记忆保证。

定理对齐次核给出有限区间上的一致误差,对脉冲响应给出积分绝对误差。两种范数并不矛盾:后者在零点附近有可积奇异性,有限指数和在那里有界,不能要求一致点态近似。命题 3 用这两个误差控制固定系统、有限区间和有界输入下的状态近似;它不是对训练后的选择性网络提供任意长度的误差证书。

2. 区间冻结控制:用内容调节同一个记忆库的有效时间尺度

控制分支为每个头、每个 token 生成步长 \(\Delta_t\)、阶数 \(\alpha_t\) 和缩放参数 \(\lambda_t\)。步长和缩放参数经过 softplus 保持正值,阶数经过 sigmoid 落在 \((0,1)\)。基础时间尺度共享,但其有效值按 \(\lambda_t^{1/\alpha_t}\) 缩放;因此网络能够根据输入改变整个库的记忆跨度,而不是仅在固定模式之间选择一个输出。

每个 token 区间内冻结控制和输入后,模式是普通的标量线性系统,可以进行精确零阶保持(ZOH)离散化。以下同时列出理论注入系数与实验实际采用的注入系数,避免把二者混成一个公式。

\[ \tilde\rho_{t,m}=\exp\!\left(-\Delta_t\lambda_t^{1/\alpha_t}/\tau_m\right),\qquad \tilde\beta^{\mathrm{ZOH}}_{t,m}=\frac{1-\tilde\rho_{t,m}}{\lambda_t},\qquad \tilde\beta^{\mathrm{impl}}_{t,m}=1-\tilde\rho_{t,m}. \]

正文式 (20) 使用带除数的精确 ZOH 系数,附录 D 的 Algorithm 1 和 Remark 7 明确说明所有实验去掉了这个除数,以解耦时间尺度控制与写入幅度、改善训练稳定性。只有输入无关的固定 \(\lambda\) 才能通过常数状态/读出缩放建立二者的等价关系;对于 token 相关参数,不能直接宣称实验递推就是理论 FDE 的精确离散化。

实现还把 \(\Delta\in[10^{-4},1.0]\)、\(\lambda\in[0.25,4.0]\) 用于数值裁剪。对固定模式系统,“精确”指区间冻结条件下的局部 ODE 积分;它不表示逐 token 改变阶数后仍精确求解某个变阶 Caputo FDE。Remark 8 将完整网络定位为固定分数阶构造的架构推广和归纳偏置,这比把定理直接延伸到可学习网络更准确。

3. 先验读写路由:同一时间尺度坐标,独立的存取决策

有了多个模式,还要决定什么内容值得写入慢状态,以及当前输出应从哪些模式读取。理论求积系数在大时间尺度部分具有阶数控制的衰减结构,因此 Frac 将负的“阶数乘对数时间尺度”作为读路由先验,再加入从卷积后内容生成的残差 logits。写路由采用同样形式,但使用独立映射,避免把“适合保存”与“适合现在读取”强行绑定。

\[ c_{t,m}=\operatorname{softmax}_m\!\left(-\alpha_t\log\tau_m+g_{\mathrm{read}}(u_t)_m\right),\qquad b_{t,m}=\operatorname{softmax}_m\!\left(-\alpha_t\log\tau_m+g_{\mathrm{write}}(u_t)_m\right). \]

两个映射都是带可学习标量 sigmoid 门的线性投影,参数彼此独立。softmax 在模式轴上归一化,使先验和内容残差一起决定相对分配,而不是允许各模式幅度任意漂移。读先验有理论渐近系数的来源,写先验则并非由定理唯一决定:它是作者采用的共同寻址方案,由消融验证其效果。

这也说明 Frac 不是简单复制多个独立记忆槽。模式有共享的对数时间尺度语义,当前内容能修正这种默认偏好;阶数同时参与时间尺度控制与路由先验。将多尺度库、可学习控制和归一化存取拆开看,才能理解为什么“仅有多个状态”不足以复现完整模型。

4. 仿射扫描实现:长记忆先验不破坏并行预填充和单步解码

将实际注入因子乘以写权重,就得到每个模式对当前内容的写入系数。状态更新仍然是“旧状态乘保留率,加当前写入”,输出是更新后状态的读权重加权和;这些权重由输入而非递归状态产生,因此能够预先构建扫描参数。

\[ s_{t,m}=\tilde\rho_{t,m}s_{t-1,m}+\tilde\beta^{\mathrm{impl}}_{t,m}b_{t,m}u_t,\qquad h_t=\sum_{m=1}^{M}c_{t,m}s_{t,m}. \]

仿射变换可以按结合律组合,训练和预填充因此使用分块并行扫描;解码只缓存各头的模式状态。主配置每层主递归状态为 \(16\times16\times256=65{,}536\) 个标量,规模不随上下文增长。该统计不包括短卷积缓冲和架构特有的辅助缓存,也不等于整个模型的运行显存。

Triton 实现融合扫描参数构建、块内状态构造、块间状态传递和读写混合,并把直接馈通并入输出计算。块内传播不使用累计指数的比值,而在累计对数衰减中取差,再计算指数,减少下溢等数值问题。多个模式增加的成本由模式数控制,不能仅凭“序列长度线性”就断言比所有注意力实现更快。

一个完整示例

考虑合成任务中很早出现的一个正事件,后面隔了很长一段背景 token,末尾才判断所有事件的加权和符号。事件到达时,局部因果卷积生成内容;控制分支决定当步各模式的保留率;写路由把内容分配到快慢状态。之后每个背景 token 都继续更新状态,快模式中的事件贡献较快消失,慢模式仍可能保留较弱但非零的贡献。

到末尾时,独立读路由组合这些状态,再由任务头预测符号。模式组合能在选定范围内产生比单一指数更广的时间尺度覆盖,但如果后续控制持续强烈遗忘,早期事件仍会丢失。这个例子解释的是数据流,不是假定训练后所有 token 都共享同一个固定幂律核,也不保证精确恢复任意多条历史事件。

损失函数 / 训练策略

语言模型使用标准自回归下一 token 预测目标,没有额外的分数阶方程监督损失。主实验从头训练约 1.3B 参数模型,在去重 FineWeb-Edu 中采样 100B tokens,使用 Llama-2 的 32K 词表和 4K 完全打包序列。Frac 为 48 层、隐藏维度 2048、16 个头;比较模型约匹配参数量,但层数和内部扩展不完全相同。

优化器为 AdamW,初始学习率 \(6\times10^{-4}\),预热比例 10%,权重衰减 0.1,全局批量 11M tokens,并使用 FSDP 与混合精度。消融另外使用 390M 模型和 30B tokens,在 4K 上训练、16K 上评估,不能把这些困惑度直接当作 1.3B 主模型结果。

实验关键数据

主实验

下面汇总原文表 2–4 的平均分,均为越高越好。LongBench 平均混合 F1、ROUGE、准确率和代码编辑相似度等任务得分;短上下文平均排除了两列困惑度;真实召回检索采用截断至 2K 的输入,并不是长距离外推实验。

模型 LongBench,14 任务 短上下文,8 项平均 真实召回检索,6 任务
Transformer 14.7 55.3 38.5
Mamba2 12.0 53.8 29.9
GDN 16.0 54.2 30.9
Mamba3-SISO 14.4 54.6 35.1
Mamba3-MIMO 11.4 55.3 36.8
Frac 17.9 55.1 36.4

LongBench 相比 GDN 提高的是 1.9 个分数点,不是相对提升 1.9%。作者统计 Frac 在 14 项中赢得 8 项,不是全部获胜;例如 LCC 为 11.0,明显低于 Transformer 的 26.0,2WikiMQA 为 12.2,也低于 Mamba3-SISO 的 16.3。本文的优势不能概括成所有长文本功能都更强。

合成重尾任务单层约 200K 参数,训练长度为 512,外推到 128K,表 6 使用 10 次运行。128K 时 Frac 准确率为 \(63.8\pm1.9\),GDN 为 \(58.4\pm3.8\);但训练长度处 Frac 的 \(96.5\pm2.5\) 低于 Mamba3 的 \(98.6\pm0.7\)。这支持“远距离退化更慢”,不支持“每个长度都最优”。MADLab 的四项平均为 75.4,对比 Mamba3 的 74.8,饱和的 ICR 和 N-ICR 未纳入主表。

语言模型 NIAH 是另一套实验:4K 训练,测试至 64K;图 3 显示总体退化更慢,但重复填充文本的 S-NIAH-1 上 GDN 是重要例外。HG38 DNA 实验使用约 7M 参数、1K–64K 多种训练长度,并在各自相同长度测试;它考察上下文长度扩展的收益,不是短训练长度向长测试长度的外推。

消融实验

原文表 9:390M 参数、30B tokens、4K 训练,以下三列均为 16K 长文档困惑度,越低越好。

配置 ProofPile PG19 GovReport
完整 Frac 47.2 28.4 11.3
\(\alpha=1\) 61.3 36.2 19.7
去掉写先验 68.5 41.6 18.2
纯多尺度库 76.3 47.1 17.2
去掉读写 softmax 80.4 36.9 15.5
自由学习时间尺度 54.1 34.3 13.5
去掉 \(D\) 70.9 46.4 23.6
\(M=8\) 51.4 33.9 14.5
\(M=32\) 47.9 29.1 10.6

“纯多尺度库”同时固定阶数、去掉两侧先验并改变时间尺度初始化,属于组合消融,不能把全部差距归因于其中一个因素。单独自由学习时间尺度也变差,说明几何布局有价值;去掉直接馈通在 GovReport 上损失尤其明显。增加到 32 模式并非三列都改善,作者报告其代价为参数增加 5%、计算慢 7%。

附录 F 表 8 提供系统成本的边界:以下单位都是 tokens/s,预填充批量为 1,解码平均生成 64 tokens;Transformer 的解码列在 16K 上下文测得。

模型 预填充 1K 预填充 4K 预填充 16K 解码 训练 4K
Transformer 55,839 79,173 56,297 62 218,184
Mamba2 15,362 41,060 67,571 34 278,620
GDN 10,455 35,481 58,384 20 215,357
Mamba3-SISO 19,517 44,998 74,314 27 238,924
Mamba3-MIMO 13,692 30,710 46,707 25 200,994
Frac 16,213 49,946 77,490 33 241,598

关键发现

  • 长上下文收益和短上下文保持应分别讨论:Frac 的 55.1 仍低于两个 55.3,2K 召回检索的 36.4 也低于 38.5 和 36.8,不能把“竞争力”写成全面领先。
  • 16K 预填充最高为 77,490,但 1K 和 4K 的 Transformer 更快;Frac 解码 33 低于 Transformer 的 62 和 Mamba2 的 34,训练也不是最快。
  • 表 10 对固定齐次核做单纯形约束下的系数拟合,8、16、32 模式的平均最大绝对误差依次为 \(1.34\times10^{-3}\)、\(6.82\times10^{-4}\)、\(6.45\times10^{-4}\)。这是特定归一化时间区间上的核拟合,不是学得读写路由的端到端误差。

亮点与洞察

  • 把记忆规律当作架构轴。 遗忘曲线、关联更新规则和记忆容量是不同问题。本文主要改第一项,因此有机会与 delta-rule 更新互补,而不是替代所有关联记忆机制。
  • 理论指导结构,但不冒充完整网络证书。 非负扩散表示解释了几何时间尺度和读先验;写先验、softmax 路由以及实际注入幅度仍是架构选择。分开叙述能看清理论究竟约束了什么。
  • 多模式不必意味着更大主递归状态。 当前配置的模式库比比较模型的主状态更小,说明更合理的时间尺度分配可能比盲目增加状态元素更重要。但该观察不能推广成所有实现上的总显存优势。

局限与展望

  • 有限库依赖模式数与时间尺度范围,范围外依赖或需要迅速遗忘的任务可能收益下降。固定有限指数和在足够远处仍受最慢指数模式支配,不提供无限时间的精确幂律。
  • 固定系统的理论不能直接推出 token 选择性模型是精确变阶 FDE;实际注入系数又主动偏离正文 ZOH 形式。未来可显式分析控制变化和路由对核误差、稳定性的影响。
  • 1.3B 主比较约匹配参数量,但并非层数、状态容量和实现优化都严格一致。消融没有为每个小改动提供主任务上的多种规模复验,不能从一个困惑度表推断普适因果关系。
  • 核近似实验仅针对无输入的固定齐次系统,且在有限归一化时间区间上优化系数;并未同时验证选择性输入响应、零点附近奇异性或真实训练路由。
  • 自定义 Triton 实现和基线内核的优化程度不同,吞吐结论受硬件、批量与序列长度制约。后续应扩大硬件和批量覆盖,并单独报告缓存、卷积缓冲及完整运行显存。

相关工作与启发

  • vs Mamba / Mamba2 / Mamba3:这些方法强化选择性、计算结构或离散化表达能力;Frac 从分数阶核推导时间尺度模式。Frac 的多个模式不是 Mamba3 的 MIMO 秩投影,二者不能仅凭“多路状态”混称。
  • vs GDN / delta-rule 模型:后者用关联更新提升存取精度,Frac 更强调不同距离的保留规律。将分数阶时间尺度与纠错式关联更新结合,是作者提出但尚未在本文验证的方向。
  • vs RetNet / Mixture-of-Memories:多头固定衰减与多个独立记忆都能增加时间尺度或容量;本文额外引入分数阶来源的几何模式和逐 token 存取先验,其收益由组合消融和单因素消融共同支持。
  • vs FADE:FADE 通过历史状态与迭代求解器处理分数阶积分方程,Frac 以有限指数和获得固定大小递归状态,更直接服务于 decoder 语言模型。代价是有限范围近似以及从严格 FDE 到选择性架构的理论距离。

评分

  • 新颖性: 4/5 — 将分数阶记忆核落实为可扫描的选择性序列模块,结构来源清楚。
  • 实验充分度: 4/5 — 覆盖合成外推、语言模型、DNA、消融与吞吐,但并非所有评测都验证外推。
  • 写作质量: 4/5 — 附录明确解释理论和实现注入系数的差异;正文的百分比措辞需要按分数点理解。
  • 价值: 4/5 — 提供可复用的长记忆先验,短上下文检索和解码速度仍需改进。