Sessa：选择性状态空间注意力

论文信息

标题: Sessa: Selective State Space Attention

作者: Liubomyr Horbatko

发布日期: 2026-04-20

arXiv ID: 2604.18580v1

PDF 链接: 下载 PDF

3 分钟速览

研究问题：这篇论文试图解决现代序列模型中一个基础问题：当注意力机制无法精确聚焦（扩散注意力）或状态空间模型无法维持长期记忆（冻结时间失败）时，现有架构都面临信息衰减过快的困境。作者希望设计一种新的序列混合器，在这些失败模式下仍能维持更慢的远端信息衰减。
核心方法：提出 Sessa 架构，将输入依赖的（input-dependent）注意力放置到反馈路径内部，实现单层内的 “多路径聚合”（many-path aggregation），而非传统 Transformer 的单次读取或 Mamba 的单链反馈（见第 3.2 节）。
关键结果：在扩散路由（diffuse routing）假设下，Sessa 能实现多项式衰减的记忆尾部 $O(\ell^{-\beta_{\text{tail}}})$ ，指数 $\beta_{\text{tail}} \in (0,1)$ ，这比注意力稀释的 $O(1/\ell)$ 和状态空间模型的指数遗忘更慢（见定理 8）。
主要局限：在短上下文语言建模任务上，Sessa 的性能略低于 Transformer 和 Mamba-2（见实验表 2），作者认为这是因为反馈机制在该场景下利用率低，分配了参数容量但未充分获益。
适合读者：关注序列模型架构设计、长上下文推理、记忆机制理论的研究者和工程师，特别是对 Transformer、状态空间模型（SSM）和递归网络有基础了解的读者。

论文背景和研究动机

现代基础模型在处理长上下文序列时面临一个核心挑战：如何从遥远的过去位置有效地检索和整合信息。目前主流的两种架构各有其失败模式。

Transformer 的自注意力机制是一种 “直接读取” 系统：每个位置 $t$ 通过查询-键匹配，一次性从可见的历史窗口读取信息。当注意力权重能够尖锐地聚焦在小部分相关令牌上时，这种方式效果很好。然而，当上下文充满竞争性的 “干扰项”（distractors），注意力变得扩散时，单个令牌的影响力会随着有效支持集大小 $S_{\text{eff}}(t)$ 而稀释，通常按 $O(1/S_{\text{eff}}(t))$ 衰减；在全前缀设定下，对于非常老的令牌会达到 $O(1/\ell)$ （其中 $\ell = t - \tau$ 为时滞）（见第 1 节）。

状态空间模型（SSM）如 Mamba 采用反馈系统：通过显式的状态递推传播信息。选择性的变体可以使转移矩阵在某些步长上接近单位矩阵（“冻结时间”），从而减缓衰减。但是当模型无法维持长期的冻结时间走廊（freeze‑time corridor）时 — 例如，当相关和不相关的位置产生相似的局部表征时 — 累积的离散化时间 $\sum \Delta_t$ 会随滞后线性增长，导致指数级遗忘（见命题 5）。

这些观察指向一组互补的失败模式：注意力遭受稀释，而选择性 SSM 遭受指数遗忘。Sessa 的动机在于，能否设计一种架构，在上述两种锐利检索都失效的扩散/失败冻结时间区间下，仍然维持更缓慢的信息衰减？

核心方法和技术细节

Sessa 是一个解码器（decoder）层，核心创新在于其混合器（mixer）的设计。一个 Sessa 块首先对输入 $x$ 进行层归一化和门控线性投影得到 $a$ 和门控向量 $g$ ，然后通过混合器计算 $s = \text{Mixer}(a)$ ，最后经门控和输出投影与残差连接得到输出 $y$ 。

混合器结构。混合器同时使用两种因果注意力机制：

前向注意力：标准的多头自注意力变体，用于产生前向信号 $f_t = \sum_{j=0}^{t} \alpha^{\text{fwd}}_{t,j} v_j$ 。这一部分引入 RoPE 位置编码，负责当前令牌对可见上下文的直接读取。
反馈注意力：在严格过去（ $j < t$ ）上计算注意力权重 $\alpha^{\text{fb}}_{t,j}$ ，并与一个可学习的标量增益 $\gamma_t \in (-1, 1)$ 相乘，构造成严恪下三角的反馈矩阵 $[B_{\text{fb}}]_{t,j} = \gamma_t \alpha^{\text{fb}}_{t,j}$ 。注意反馈路径不使用任何位置编码 — 其时间不对称性由强制因果施加的严格下三角结构自然提供。

反馈求解。混合器的核心运算不是简单地对值进行加权求和，而是求解一个下三角线性系统：

(I - B_{\text{fb}}) s = f

等价的逐次代入形式为 $s_t = f_t + \gamma_t \sum_{j=0}^{t-1} \alpha^{\text{fb}}_{t,j} s_j$ 。这一结构具有关键的理论性质：由于 $B_{\text{fb}}$ 在有限长度上严格下三角，它满足 $B_{\text{fb}}^T = 0$ （幂零性）。因而解可以展开为诺伊曼级数：

s = (I - B_{\text{fb}})^{-1} f = \sum_{k=0}^{T-1} B_{\text{fb}}^k f

其中 $B_{\text{fb}}^k f$ 对应经过 $k$ 步反馈路由的贡献。在索引层面， $(B_{\text{fb}}^k)_{t,\tau}$ 是从 $\tau$ 到 $t$ 的所有长度为 $k$ 的时间路径上的权重乘积之和（见第 2.1 节）。

这种多跳路由是 Sessa 区别于 Transformer（单跳）和 Mamba（单一路径上的多跳）的结构性特征。即使单条边的权重很小，路径的总数与滞后长度相关，求解过程聚合了所有跳数和所有可能路径的贡献，这正是产生多项式衰减尾部分析的来源。

理论分析框架。作者通过固定路由下的雅可比矩阵来隔离记忆机制的特性：对 Transformer， $\partial y_t / \partial v_\tau$ 直接等于注意力权重 $\alpha^{\text{fwd}}_{t,\tau}$ ；对 Sessa， $\partial s / \partial f$ 为 $(I - B_{\text{fb}})^{-1}$ 。在扩散假设下（注意力权重 $\leq c_2/t$ ，且增益 $\gamma_t \leq \gamma_{\max} < 1$ ），论文证明了解的影响力 $|y_{\tau+\ell}|$ 按 $\ell^{-\beta_{\text{tail}}}$ 衰减，其中 $\beta_{\text{tail}} = 1 - \gamma_{\max} c_2 \in (0,1)$ （定理 8）。在 $\alpha_{t,j}=1/t$ 且 $\gamma_t=\gamma$ 的均匀扩散情形下，给出了闭合形式解 $y_{\tau+\ell} = \Theta_\tau(\ell^{-(1-\gamma)})$ （附录推论 F.2）。

创新点和贡献

架构创新。Sessa 首次将输入依赖的注意力明确放在反馈回路内部，构建了一个 “输入依赖的反馈系统”（见第 1 节和第 3 节）。这与传统的 Transformer（输入依赖的直接读取）和选择性 SSM（输入依赖的链状反馈）形成完备对比。

记忆特性理论。论文系统刻画了三种架构在匹配的失败模式下的长期记忆行为：

Transformer： $O(1/\ell)$ 稀释；
Mamba（失败冻结时间）： $O(e^{-c\ell})$ 指数遗忘；
Sessa： $O(\ell^{-\beta_{\text{tail}}})$ 多项式衰减，且指数可随结构参数调节（命题 9、10）。

灵活选择性检索定理。在有限值域族上，深度 Sessa 网络可以实现给定指数的检索剖面，包括衰减、恒定和递增的剖面，而匹配深度下的扩散 Transformer 和失败冻结时间 Mamba 不能实现恒定或递增的检索行为（定理 12 与命题 13）。这一结果不仅展示了 Sessa 在衰减速度上的优势，更揭示了一种定性上的检索灵活性差异。

内建位置信息。论文证明，即使反馈注意力不使用显式位置编码，Sessa 的反馈求解本身可以产生绝对位置的、分离的位置编码信号，并且可以让一个令牌级连续映射从中恢复出位置索引（推论 4.13）。这意味着理论上 Sessa 可以在前向注意力使用 RoPE（只提供相对位置）的同时，在反馈路径内生绝对位置信息。

通用逼近性质。作为旁证，论文证明了带宽度扩展器的 Sessa 网络可以以任意精度逼近任何紧集上的连续因果映射（定理 14）。

实验结果分析

为保持可比性，实验采用参数规模、优化器和训练步数完全一致的三组模型，唯一的区别在于混合器类型（Sessa、多头自注意力、Mamba-2 混合器）。

长上下文任务。论文聚焦于两个合成基准：

SymbolSoup：一种分类任务，要求模型从长噪声序列中识别两个有特定风格的块，且这两个块的顺序随机排列。
扩散 MQAR：修改自 MQAR 的多查询关联召回任务，引入了多令牌键、结构性干扰项以及测试时更大的检索延迟。

在这两项任务中，Sessa 都取得了最优成绩。在 SymbolSoup 上分类精度达 $0.8601$ ，而 Transformer 为 $0.7921$ ，Mamba-2 为 $0.0500$ （机会水平）；在扩散 MQAR 上令牌精度为 $0.1541$ ，对应对照组的 $0.1222$ 和 $0.0021$ （见表 1）。Mamba-2 在这两个任务上均未收敛，作者认为这印证了理论中的失败冻结时间分析（见第 5.1.1 节）。

短上下文语言建模。在 SimpleStories 语料库上，Sessa 的验证困惑度为 $8.3700$ ，略逊于 Transformer（ $7.6701$ ）和 Mamba-2（ $7.7229$ ）（见表 2）。消融实验显示，移除反馈分支后，困惑度降至 $8.0902$ ，精度回升，支持了作者的假设：在短上下文场景下，为反馈分支分配的参数容量未得到充分利用。

实践建议

Sessa 的当前实现采用全前缀注意力，计算复杂度为 $O(T^2)$ 。因此，在极长上下文直接部署全精度 Sessa 并不现实。基于论文特性，以下建议值得在实践中考虑：

长上下文检索场景优先：Sessa 的核心优势体现在需要综合大量远端碎片化证据的任务中（如长文档问答、长程关联记忆）。如果你的任务中注意力分布天然扩散（例如需要比较多个文本段落），Sessa 的理论优势更可能转化为实际增益。
混合架构的启示：即使不直接使用 Sessa，其核心思想 — 将注意力作为反馈算子 — 可以融入现有架构设计。你可以尝试在 Transformer 的某些层中引入一个带标量增益的反馈注意力旁路，并求解递推形式的输出，从而在局部实现多跳聚合。
内生位置编码的利用：论文证明反馈回路自身可以生成绝对位置信号。在设计长上下文模型时，可以有意省略反馈路径中的显式位置编码，仅依赖其内生时间不对称性，这可能带来更强的长度外推能力 — 因为不存在传统位置编码的长度限制。
短上下文优化：Sessa 在短上下文上表现稍弱，但这可以通过结构上的条件性激活改善。例如，可以设计门控机制，在检测到短上下文或局部模式时动态减小反馈增益 $\gamma_t$ ，使模型回退到类似前向注意力的行为。
计算效率优化：尽管论文未涉及这一点，但递推求解 $(I - B_{\text{fb}}) s = f$ 的逐次代入实现可以实现 $O(T^2)$ 复杂度，且不改变 Sessa 的记忆特性。若需扩展到更长序列，可以考虑引入稀疏反馈注意力（如滑动窗口 + 跨步采样），此时只要仍保留对历史状态的足够路径多样性，多项式衰减的理论性质有望部分保留 — 但这需要额外的理论验证。