Sessa：选择性状态空间注意力新架构深度解析

背景与动机

现代序列建模主要被两大架构主导：Transformer 和结构化状态空间模型（SSM）【1†source】。Transformer 通过自注意力（self-attention）机制可以访问上下文中的任意元素，但当注意力变得扩散（即对大量过去token均匀关注）时，单个token的影响力会被稀释，典型情况下随距离呈 O(1/ℓ) 衰减【1†source】。相反，结构化SSM通过显式的循环状态传播信息，但即使像Mamba这样的选择性SSM在无法维持“冻结时间”时，其长程敏感性也会随滞后呈指数衰减【1†source】。因此，这两种机制在长上下文下都面临挑战：Transformer 会因为注意力稀释而难以保留远距离信息，而SSM则可能因状态遗忘而丢失长程依赖【1†source】。现有架构要么只通过单次注意力读取来检索过去信息，要么仅通过单一反馈链传播信息，缺乏在单层内对多种路径进行聚合的能力【1†source】。

核心创新：反馈路径中的注意力

Sessa（Selective State Space Attention）架构的核心创新在于将注意力机制置于反馈路径之中【2†source】。换言之，Sessa将输入依赖的注意力路由注入到循环反馈通路中，从而结合了直接读取输入与通过反馈状态聚合信息的双重机制【2†source】。这一设计使得在单个层内可以实现多路径聚合：过去token可以通过多条注意力路径影响未来状态，而不再局限于单次注意力读取或单一循环链【1†source】。具体而言，Sessa的解码器在每层采用单个门控MLP风格的模块，将因果前向注意力信号与反馈项相结合【5†source】。反馈项通过一个下三角线性系统（因果卷积）来混合过去的输出，实现多跳（multi-hop）路由【5†source】。这种反馈求解（feedback solve）机制正是Sessa实现其长程记忆特性的关键。

理论分析：长程记忆与选择性检索

Sessa通过反馈路径中的注意力，在理论上展现出显著的长程记忆优势。在给定假设和匹配条件下，Sessa能够实现幂律记忆尾，即影响随滞后 ℓ 呈 O(ℓ^-β) 衰减，其中 0<β<1【1†source】。这一衰减速率比Transformer的 O(1/ℓ) 和Mamba在失败冻结时间情况下的指数衰减都要慢得多【1†source】。换言之，Sessa的记忆衰减更加平缓，能够在长上下文中更久地保留信息。更重要的是，在相同条件下，只有Sessa在所比较的模型类中实现了灵活的选择性检索，包括那些影响力不随距离衰减的配置【1†source】。这意味着Sessa能够根据需要选择性地从过去提取信息，而不受长距离稀释或遗忘的限制。

幂律记忆尾

Sessa的反馈机制导致其传输核（kernel）具有特殊的下三角结构，从而产生幂律尾部。具体而言，Sessa的反馈可以表示为一个严格下三角矩阵 B_fb，其对应的传输核 K = (I - B_fb)^-1 也是下三角矩阵，对角线元素为1【5†source】。通过将反馈路由系数限制在适当范围内，可以证明对于远离源点的目标位置，K_t,k 的绝对值被上界 C_K(t-k)^-β_tail 所控制，其中 β_tail 介于0和1之间【5†source】。这表明影响随距离呈多项式衰减，而非Transformer的线性衰减或Mamba的指数衰减【5†source】。在均匀路由（最扩散的情况）和恒定正反馈的极端设定下，这种幂律衰减被证明是紧致的：影响恰好以 Θ(ℓ^-β_tail) 衰减【5†source】。因此，Sessa的幂律记忆尾不仅存在，而且在某些情况下是衰减速率的下界，无法再快衰减【5†source】。

灵活的选择性检索

选择性检索指的是模型能够根据输入动态地选择关注哪些过去的信息，而不是对所有历史信息一视同仁地衰减。Sessa通过将注意力放入反馈路径，使得模型可以累积来自多个路由深度和多个时间路径的信息【5†source】。这种多路径聚合机制赋予Sessa在长程依赖上的灵活性：它可以在需要时保持对某些远距离信息的影响不衰减（即非衰减配置），也可以在其他情况下让影响逐渐衰减【1†source】。理论分析表明，在所考虑的模型类中，只有Sessa能够同时实现幂律衰减和非衰减配置【1†source】。Transformer的注意力权重本质上对所有旧token都呈 O(1/ℓ) 衰减，无法保持恒定影响；Mamba在选择门打开时能够避免指数衰减，但一旦无法维持“冻结时间”，敏感性就会迅速丧失【1†source】。相比之下，Sessa即使在注意力路由非常扩散、缺乏明显“冻结”区间的情况下，仍可通过反馈累积保持长程影响，从而实现真正的选择性检索【1†source】。

与Transformer和Mamba的对比

为了更直观地理解Sessa的优势，可以将其与Transformer和Mamba进行对比：

• 记忆衰减速率：Transformer的记忆影响随距离 ℓ 呈 O(1/ℓ) 衰减，Mamba在无法保持冻结时间时影响呈指数衰减 O(e^-λℓ)【1†source】。Sessa则可实现 O(ℓ^-β) 的衰减，其中 0<β<1【1†source】。这意味着对于非常大的 ℓ，Sessa的影响衰减最慢，能够更长时间地保留信息【1†source】。
• 选择性检索能力：Transformer对所有旧token的影响力都会随距离稀释，无法选择性地保持某些信息不衰减；Mamba虽然可以通过选择性门控在短期内保持信息，但一旦进入非冻结区间，长程敏感性会急剧下降【1†source】。Sessa则通过反馈机制，即使在注意力非常扩散的情况下，依然能够通过多条路径累积影响，从而实现对特定信息的选择性保持【1†source】。
• 实现机制：Transformer依赖全局自注意力，计算开销随序列长度二次增长；Mamba通过选择性状态空间模型在线性时间内处理序列，但需要额外的机制来保持长程信息。Sessa则在单层内结合了全局注意力（用于直接读取当前输入）和局部反馈（用于聚合过去信息），在保持线性时间复杂度的同时，实现了超越两者的记忆特性【1†source】【5†source】。

实验评估

在实验中，Sessa的上述理论优势得到了验证。在长上下文基准测试上，Sessa在匹配架构和训练预算的情况下取得了最强的性能【1†source】。这意味着在需要处理长距离依赖的任务中，Sessa能够更有效地保留和利用过去的信息。与此同时，在短上下文语言建模任务上，Sessa仍然与Transformer和Mamba风格的基线模型保持竞争力【1†source】。这表明Sessa并非通过牺牲短期性能来换取长期记忆，而是在长短期任务上都表现良好。这种均衡的性能使其在实际应用中具有吸引力，特别是在需要同时处理长短期依赖的场景下。

架构细节与实现

Sessa的架构设计旨在作为标准解码器块的即插即用替代【4†source】。它使用单个门控MLP风格的块来包裹一个循环混合器，而非传统Transformer中交替的注意力块和MLP块【5†source】。该混合器结合了标准的因果前向注意力信号和一个反馈项，反馈项通过下三角矩阵乘法来混合过去的输出，从而实现多跳路由【5†source】。为了限制反馈的增益并保证稳定性，Sessa引入了一个超参数 γ_max，要求反馈增益 |γ_t| 始终小于 γ_max（γ_max<1）【8†source】。这种设计保证了反馈系统的稳定性，并防止梯度爆炸等问题。Sessa的实现基于PyTorch，并支持可选的FlashAttention加速【8†source】。通过简单地替换现有解码器层为Sessa层，模型即可获得更强的长程信息保持能力。

结论

Sessa通过将注意力嵌入反馈路径，创造性地解决了Transformer和SSM在长上下文下的局限性。它证明了在单层内聚合多路径信息的可行性，并从理论上和实验上都展示了更慢的长程衰减和更灵活的选择性检索能力【1†source】。这一新架构为序列模型在处理长距离依赖时的信息保持和检索提供了新的思路，有望在需要长上下文理解的自然语言处理、时间序列预测等任务中取得突破。随着对Sessa及其变体的进一步研究，我们期待看到更多结合反馈机制与注意力优势的创新模型出现，推动序列建模能力的发展。