【每日论文】Forgetting Transformer: Softmax Attention with a Forget Gate

下载PDF或查看论文，请点击：LlamaFactory - huggingface daily paper - 每日论文解读 | LlamaFactory | LlamaFactory

摘要

现代循环序列模型的一个关键组件是遗忘门。虽然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”（FoX）的新型Transformer变体，通过引入遗忘门机制，在长上下文语言建模、长度外推和短上下文下游任务中优于Transformer，同时保持了Transformer的长上下文检索能力。

问题1：这篇论文想要解决什么具体问题？

• 问题背景：尽管Transformer在处理长上下文信息方面表现出色，但它缺乏一种显式的机制来以数据依赖的方式忘记过去的信息。这种机制在循环序列模型中很常见，并且在短上下文任务中已被证明是至关重要的。

• 现有方案不足：现有的循环序列模型虽然具有忘记过去信息的机制，但在长上下文能力方面表现不佳，而Transformer虽然在长上下文信息处理上出色，但缺乏忘记过去信息的机制。

• 研究目标：提出一种将遗忘门机制自然地集成到Transformer中的方法，从而在保持Transformer长上下文能力的同时，提高其在短上下文任务上的性能。

问题2：论文的核心创新点是什么？

• 技术创新：提出了一种名为“遗忘注意力”（Forgetting Attention）的新机制，通过下采样未归一化的注意力分数来实现数据依赖的遗忘。

• 方法改进：将遗忘门机制应用于softmax注意力，并通过下采样未归一化的注意力分数来实现。

• 优势：与Transformer相比，FoX在长上下文语言建模、长度外推和短上下文下游任务上表现更优，同时保持了Transformer的长上下文检索能力。

问题3：实验结果如何验证了方法的有效性？

• 关键实验：在长上下文语言建模、长度外推和短上下文下游任务上进行了实验，与Transformer和多种循环序列模型进行了比较。

• 性能提升：FoX在长上下文语言建模、长度外推和短上下文下游任务上优于Transformer，同时在长上下文下游任务上与Transformer表现相当。

• 对比结果：FoX在所有测试的长上下文任务中均优于Mamba-2、HGRN2和DeltaNet等循环序列模型。

问题4：这个研究的实际应用价值是什么？

• 应用场景：FoX可以应用于需要长上下文信息处理的任务，如自然语言处理、机器翻译、文本摘要等。

• 实施建议：在构建大型语言模型时，可以考虑将FoX作为Transformer的替代方案，以在保持长上下文能力的同时提高短上下文任务的性能。

• 局限与展望：目前的研究仅限于因果序列建模，未来可以扩展到非因果情况。此外，可以根据遗忘门值自适应地剪枝计算，以降低训练和推理成本。