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 函数衰减;有外部输入时,历史内容由脉冲响应核累积。二者都是重尾,但指数不同,不能统称为同一个幂律。
上述渐近关系要求固定 \(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)离散化。以下同时列出理论注入系数与实验实际采用的注入系数,避免把二者混成一个公式。
正文式 (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。写路由采用同样形式,但使用独立映射,避免把“适合保存”与“适合现在读取”强行绑定。
两个映射都是带可学习标量 sigmoid 门的线性投影,参数彼此独立。softmax 在模式轴上归一化,使先验和内容残差一起决定相对分配,而不是允许各模式幅度任意漂移。读先验有理论渐近系数的来源,写先验则并非由定理唯一决定:它是作者采用的共同寻址方案,由消融验证其效果。
这也说明 Frac 不是简单复制多个独立记忆槽。模式有共享的对数时间尺度语义,当前内容能修正这种默认偏好;阶数同时参与时间尺度控制与路由先验。将多尺度库、可学习控制和归一化存取拆开看,才能理解为什么“仅有多个状态”不足以复现完整模型。
4. 仿射扫描实现:长记忆先验不破坏并行预填充和单步解码
将实际注入因子乘以写权重,就得到每个模式对当前内容的写入系数。状态更新仍然是“旧状态乘保留率,加当前写入”,输出是更新后状态的读权重加权和;这些权重由输入而非递归状态产生,因此能够预先构建扫描参数。
仿射变换可以按结合律组合,训练和预填充因此使用分块并行扫描;解码只缓存各头的模式状态。主配置每层主递归状态为 \(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 — 提供可复用的长记忆先验,短上下文检索和解码速度仍需改进。