Wall Attention: 基于对角门控的长度泛化新方法

位置编码已成为长上下文泛化的关键瓶颈。RoPE——几乎所有主要前沿模型采用的主流方法——是数据无关的 → 无论实际出现哪些标记，相同的两个位置总是获得相同的偏置。这阻碍了长度泛化，使得模型无法在不训练越来越长数据的情况下访问长上下文。

与此同时，线性RNN领域的一条并行研究路线悄然发展出一个更强大的想法。对角遗忘门（例如GLA、RWKV-7和Kimi Linear）允许不同特征通道以不同速率衰减，使模型具备细粒度、可学习、内容相关的记忆与遗忘能力。这已成为最新一代线性RNN架构的关键要素之一。

在研究如何将对角门控引入softmax注意力时，我们发现了一个根本性障碍。有限维度的对角门无法直接作用于指数核的无限维特征空间 → 我们需要一个提升。

最有意思的是，这个障碍可以通过诱导作用框架得到清晰解决。该框架通过将任何输入空间线性算子在嵌入之前作用（而非之后）来将其提升至完整特征空间。事实证明，这统一了FoX、PaTH和Wall，使它们成为同一构造的特例。

今天，我们发布Wall Attention，它将对角遗忘门推广到了softmax注意力。Wall实现了最先进的长度外推能力，在预训练中优于RoPE和FoX，可以从4k零样本泛化到160k+上下文，并且是一个即插即用的替代方案，配有开源Triton内核用于训练和解码——其中WallDecode的解码性能与FA3解码相当。

代码：github.com/tilde-research/wall-attention-release

问题：Softmax注意力无法遗忘

标准的softmax注意力是排列等变的。对于纯粹的注意力操作而言，上下文不是时间线；它是一个静态的标记包。为了解决这个问题，我们注入位置编码（PE），如RoPE或ALiBi，以引入近因偏差。

无论标记以何种顺序引入，标准softmax注意力都会产生相同的输出。

无论标记以何种顺序引入，标准softmax注意力都会产生相同的输出。

但像RoPE这样的方法是数据无关的。相同的两个位置获得完全相同的旋转惩罚，无论它们代表什么文本。然而，自然语言的信息密度变化很大。模型应该能够查看一个标记并决定：“这是一个语义思想的结束，我应该清空我的工作记忆”，或者“这是一个关键实体，我需要永远记住它“。

现代线性RNN（如Mamba、GLA和KDA）通过数据相关的对角遗忘门解决了这个问题。它们维护一个循环状态，并在写入新信息之前应用学习的、每通道的衰减。这正是它们具有选择性记忆的原因。

具有日益复杂的门控和更新策略的架构。

具有日益复杂的门控和更新策略的架构。

问题是：我们如何将对角门控引入softmax注意力？之前的工作，如FoX，尝试通过将衰减吸收到类似ALiBi的距离惩罚中来实现标量门控。我们寻求引入对角门控。

FoX公式。

FoX公式。

认识Wall Attention——一种将对角门控完全推广到softmax注意力的方法，是RoPE的完整替代品，也是零样本长度外推的新SOTA。

Wall公式。

Wall公式。

无限维度上的对角遗忘门

要理解为什么推导Wall并非易事，请看展开的门控线性RNN：

St=AtSt-1+φ(kt)vt⊤

展开后得到非归一化的注意力权重

wi,j=⟨φ(qi),(∏r=j+1i Ar)φ(kj)⟩

为了恢复标准softmax注意力，特征映射φ是无限维的。

对于标量门控，累积乘积只是一个数字。它可以直接从内积中提取出来，得到

exp⁡(q⊤k+log⁡ Fij)

这就是FoX公式：一个简单的加性偏置。

但我们想要一个对角门控来按通道控制遗忘。

从根本上说，你无法直接参数化一个无限维的对角矩阵。因此，标量吸收技巧完全失效。

诱导作用框架

我们不是尝试将对角矩阵直接提升到特征空间，而是让它作用于输入空间，并通过φ诱导出作用。我们定义诱导作用，使得在嵌入之前对输入进行门控等于对嵌入状态进行门控：

A~ φ(x) := φ(A x)

关键的是，这个诱导作用是一个群同态。这意味着累积诱导转移等于累积乘积的诱导作用！当我们通过指数核技巧传递它时，无限维的数学崩溃成一些极其美妙和简单的东西：

wij=exp⁡(qi⊤ Aj→i kj)

当特化为对角门控时，

Aj→i=diag(∏r=j+1i gr)

定义：Wall Attention

Wall注意力允许模型在注意力机制中设置动态的每通道遗忘率，从而实现选择性保留。

ot=∑j softmaxj(∑n Fij,n qi,n kj,n)vj=∑j softmaxj(∑n (∏r=j+1i gr,n) qi,n kj,n)vj

然而，我们也可以将Wall重写为一种高效的因子化格式。将门控的对数空间前缀和定义为

Pt=∑u≤t log⁡ gu

我们在标准注意力之前重新缩放查询和键：

q~i=exp⁡(Pi)⊙qi, k~j=exp⁡(-Pj)⊙kj ot=Attn(q~, k~, v)

就是这样。这就是Wall Attention。

示例可视化：Wall中的每通道选择性保留。

示例可视化：Wall中的每通道选择性保留。

Wall的快速训练与推理

如果你天真地编码因子化Wall Attention——在HBM中计算Q̃和K̃，并将它们传递给FlashAttention——模型会崩溃。

为什么？因为Pt是单调的累积和。在序列长度8192处，Pt的大小超过160。2^(Pi)下溢，2^(-Pj)上溢bf16的限制，导致灾难性抵消，尽管最终重构的得分是完全有界的。

为了解决这个问题，我们编写了一个自定义的Triton内核。

1. 每块锚点： 我们不进行全局重新缩放，而是引入局部锚点R。

exp⁡2(Pi - Pj)=exp⁡2(Pi - R)·exp⁡2(R - Pj)

这将对数限制在单个序列块内的最大累积门控值，数学上保证了稳定性。

2. 融合的门梯度： 在反向传播中，计算门前缀dP的梯度需要一个额外的累加器。这增加了我们的寄存器压力，迫使Triton编译器静默地减半我们的块大小，降低了SM占用率。

关键洞察？我们发现

dPj=-ln⁡2·Kj⊙dKj

当内循环结束时，Kj和dKj都已经驻留在SRAM中。通过在后循环中融合门梯度，我们释放了寄存器，恢复了128x64的块大小，并达到了张量核心的饱和。

3. 解码与Flash Attention 3竞争

对于推理，存储所有标记的Pt会膨胀KV缓存。相反，我们使用块锚定引用将门控吸收到缓存键中。内循环不需要额外的门算术。我们的WallDecode内核在不同序列长度下与FA3解码相当。

在不同序列长度下，WallDecode实现了与FA3解码相当的处理能力。

在不同序列长度下，WallDecode实现了与FA3解码相当的处理能力。

实证结果

小规模预训练

在我们测试的所有规模上，Wall明显优于RoPE和FoX。Wall在我们之前的记录Aurora 1.1B的预训练token与性能曲线上设定了新的SOTA，在关键基准测试上与强大的开源SLM相匹配。

顶部：不同位置编码策略的下游评估。Wall（无位置编码）获得最强性能。底部：与使用30-500倍更多数据训练的公开可用模型以及前代Aurora模型相比的标记效率。

顶部：不同位置编码策略的下游评估。Wall（无位置编码）获得最强性能。底部：与使用30-500倍更多数据训练的公开可用模型以及前代Aurora模型相比的标记效率。

零样本长度外推

预训练的收益很好，但Wall是为长度而生的。我们取出我们的1B模型——它们仅在最大上下文窗口为4,096个token的情况下训练——并在极端上下文外推条件下进行评估。

Wall（无位置编码）稳定地泛化到160,000+个token，并且零样本性能不下降。

长度泛化基准测试。

长度泛化基准测试。

在“大海捞针“（NIAH）基准测试中，每个RoPE和FoX模型都在8k token处崩溃。Wall在16k token处保持强劲的检索能力，并在更长距离上继续表现良好。Wall在LongBench v1上也比RoPE获得了约9%的相对改进。

双模态通道的发现

当我们打开“墙壁“查看其如何路由信息时，我们发现了令人着迷的现象。

由于Wall按通道设置门控，我们追踪了整个65k token文档中的保留分数。在初始化时，所有通道开始大致相同（完全打开）。但随着训练的进行，Wall学会了将其头维度分成两个不同的群体：

• 静态记忆通道： 一些通道将其保留率锁定为恰好1.0（零方差）。它们从不遗忘。它们充当无条件的长期记忆，其行为与普通注意力完全相同。

• 动态遗忘通道： 其他通道变得高度反应。它们的保留率在每一步之间剧烈波动，完全依赖于文本的语义内容。在段落中断或语义转换处，它们迅速关闭，清除通道的工作记忆，然后再次打开以吸收下一个思想。

每个点代表一个通道，按其保留率的均值和标准差绘制。颜色对应模型中的不同层。

每个点代表一个通道，按其保留率的均值和标准差绘制。颜色对应模型中的不同层。

Wall端到端地学习了一个多时间尺度的记忆层次结构，将位置编码的结构先验直接卸载到序列的潜在动态中。

攀登“墙壁“

我们从识别一个数学病理学问题开始（如何将对角矩阵投影到无限维特征空间？），构建了一个严格框架来解决它（对称代数上的诱导作用），提出了一个新的注意力机制，并工程化了一个稳定、硬件感知的内核，使其能够规模化。

• 充分性： Wall优于我们测试的所有位置变体。它不需要在顶部叠加RoPE。它证明了对角门控是语言模型中时间理解的充分框架。

• 硬件对齐： Wall保留了普通注意力的令人尴尬的并行结构。它可以无缝地插入GQA和MLA架构，并具有高效的解码公式。

• 诱导作用桥接RNN和Softmax注意力： 据我们所知，诱导作用框架以前所未有的方式正式连接了驱动现代线性RNN的机制和标准softmax注意力。

Wall的未来是广阔的。从WallMLA，到内核改进，再到将开源RoPE模型升级训练为Wall模型，我们非常期待开源社区对Wall Attention的运用。

阅读论文：https://blog.tilderesearch.com/blog/wall-attn

获取内核：github.com/tilde-research/wall-attention-release

Dhruv Pai、Timor Averbuch、Ashley Zhang、Ben Keigwin 和 Alec Dewulf —— Tilde Research