Forgetting Transformer: Softmax Attention with a Forget Gate

现代循环序列模型的一个重要组成部分是遗忘门。虽然Transformer没有显式的循环形式，但我们将遗忘门自然地融入Transformer中，通过以数据依赖的方式降低未归一化的注意力分数的权重。我们称这种注意力机制为遗忘注意力，并将由此产生的模型称为遗忘Transformer（FoX）。我们证明了FoX在长上下文语言建模、长度外推和短上下文下游任务上优于Transformer，而在长上下文下游任务上的表现与Transformer相当。此外，它与FlashAttention算法兼容，且不需要任何位置嵌入。包括“大海捞针”测试在内的多项分析表明，FoX保留了Transformer相对于Mamba-2、HGRN2和DeltaNet等循环序列模型的优越长上下文能力。我们还引入了一种“Pro”块设计，融合了一些常见的循环序列模型架构组件，并发现它显著提升了FoX和Transformer的性能。我们的代码可在https://github.com/zhixuan-lin/forgetting-transformer 获取。

该论文试图解决长上下文语言建模、长度外推以及短上下文下游任务中的性能问题。它探索了如何在非显式循环形式的Transformer模型中引入类似循环神经网络中的遗忘门机制，以提高处理长时间依赖性的能力。这是一个新问题，因为它尝试将传统RNN中的概念与最新的Transformer架构相结合。

关键思路是通过降低未归一化的注意力分数来模拟遗忘门的效果，从而提出了一种新的注意力机制——Forgetting Attention，并基于此构建了Forgetting Transformer (FoX) 模型。这种做法不同于传统的Transformer，后者没有明确的遗忘机制。此外，作者还引入了一个'Pro'块设计，整合了一些常见的递归序列模型组件，进一步提升了FoX和标准Transformer的表现。

实验设计包括对长上下文语言建模、长度外推及短上下文下游任务的测试，结果表明FoX优于普通Transformer。值得注意的是，FoX保留了Transformer对于长时间依赖性处理的能力，甚至超过了如Mamba-2、HGRN2和DeltaNet等递归序列模型。代码已开源，可在GitHub上找到（https://github.com/zhixuan-lin/forgetting-transformer）。未来值得深入研究的方向包括但不限于：优化'Pro'块设计、探索更多应用场景以及与其他高效算法（例如FlashAttention）结合的可能性。

最近在这个领域内的相关研究包括《Universal Transformers》, 《Recurrent Memory with Flow Control for Neural Program Interpretation》, 和《Long Short-Term Memory-Networks for Machine Reading》。这些研究都在探索改进现有模型处理长时间依赖性的方法。