给Transformer加了个遗忘门，让长文本处理更高效。
Forgetting Transformer: Softmax Attention with a Forget Gate
- 用数据依赖的权重衰减机制，将遗忘门融入注意力计算
- 在长序列建模和长度外推上超越原版Transformer
- 无需位置编码，兼容FlashAttention，适合长文本任务
现代循环序列模型的核心组件是遗忘门。尽管Transformer没有显式的循环结构，我们发现可通过数据依赖方式对未归一化的注意力分数进行降权，自然地引入遗忘门。我们称该注意力机制为遗忘注意力（Forgetting Attention），对应的模型为遗忘Transformer（FoX）。实验表明，FoX在长上下文语言建模、长度外推及短上下文下游任务中表现优于Transformer，而在长上下文下游任务中性能相当。同时，它兼容FlashAttention算法且无需位置嵌入。多项分析（包括针堆测试）显示，FoX在长上下文能力上仍显著优于Mamba-2、HGRN2和DeltaNet等循环模型。我们还提出一种“Pro”模块设计，融合了常见循环架构组件，显著提升了FoX与Transformer的性能。代码已开源。
原文摘要 · Abstract (English)
An essential component of modern recurrent sequence models is the forget gate. While Transformers do not have an explicit recurrent form, we show that a forget gate can be naturally incorporated into Transformers by down-weighting the unnormalized attention scores in a data-dependent way. We name this attention mechanism Forgetting Attention and the resulting model the Forgetting Transformer (FoX). We show that FoX outperforms the Transformer on long-context language modeling, length extrapolation, and short-context downstream tasks, while performing on par with the Transformer on long-context downstream tasks. Moreover, it is compatible with the FlashAttention algorithm and does not require any positional embeddings. Several analyses, including the needle-in-the-haystack test, show that FoX also retains the Transformer's superior long-context capabilities over recurrent sequence models such as Mamba-2, HGRN2, and DeltaNet. We also introduce a "Pro" block design that incorporates some common architectural components in recurrent sequence models and find it significantly improves the performance of both FoX and the Transformer. Our code is available at https://github.com/zhixuan-lin/forgetting-transformer.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。