- 简介传统的自注意力机制依赖于softmax操作符，需要使用类似RoPE的位置嵌入或位置偏差来考虑令牌顺序。但当前的方法在处理长度泛化时仍然面临挑战。我们提出了一种基于棒折过程的替代注意力机制：对于当前令牌之前的每个令牌，我们确定一个断裂点$\beta_{i,j}$，该点表示剩余部分分配给当前令牌的比例。我们重复这一过程，直到棒子完全分配完毕，从而生成一系列注意力权重。这个过程自然地融入了最近性偏见，这在语法解析中具有语言学上的动机（Shen等人，2017）。我们研究了用棒折注意力机制取代传统基于softmax的注意力机制的影响。然后，我们讨论了数值稳定的棒折注意力机制的实现，并将Flash Attention调整以适应这种机制。当作为当前softmax+RoPE注意力系统的直接替代品使用时，我们发现棒折注意力在长度泛化和下游任务中的表现与当前方法相当。棒折注意力在长度泛化方面也表现出色，使得一个使用$2^{11}$上下文窗口训练的模型能够在$2^{14}$上下文窗口下表现出色，并且困惑度有所改善。
- 图表
- 解决问题该论文试图解决自注意力机制在处理长序列时的长度泛化问题。这是一个长期存在的挑战，特别是在自然语言处理任务中，模型需要处理比训练时更长的输入。
- 关键思路论文提出了一种基于断裂过程的新型注意力机制，称为“stick-breaking attention”。与传统的基于softmax的自注意力机制不同，这种新方法通过为每个前序token确定一个断裂点来分配剩余的注意力权重，从而自然地引入了时间顺序的偏置。这种方法不仅简化了模型结构，还提高了对长序列的泛化能力。
- 其它亮点1. 该方法在多个下游任务上表现出色，尤其是在长度泛化方面，允许模型在训练时使用较短的上下文窗口（如$2^{11}$），但在测试时仍能处理更长的序列（如$2^{14}$），并取得显著的困惑度改进。 2. 论文详细讨论了如何实现数值稳定的stick-breaking attention，并将其与Flash Attention结合，以提高计算效率。 3. 实验部分使用了多种数据集，包括但不限于语言建模和语法解析任务，验证了该方法的有效性。 4. 作者提供了开源代码，便于其他研究者复现和进一步研究。
- 1. "Attention Is All You Need" (Vaswani et al., 2017) - 提出了Transformer模型和自注意力机制，是当前许多NLP模型的基础。 2. "RoFormer: Enhanced Transformer with Rotary Position Embedding" (Su et al., 2021) - 引入了旋转位置编码（RoPE），以改善Transformer在处理长序列时的表现。 3. "Long-Range Arena: A Benchmark for Efficient Transformers" (Tay et al., 2020) - 提出了一系列基准任务，用于评估不同模型在处理长序列时的性能。 4. "Performer: Generalized Attention at Scale" (Choromanski et al., 2020) - 提出了Performer模型，通过随机特征映射来近似注意力机制，以提高计算效率。
提问交流