用棒棒分段机制替代softmax,提升长序列建模能力。
Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study

- 基于棒棒分段过程生成注意力权重,天然带近邻偏好。
- 训练时2048上下文,推理可扩展至16384,困惑度下降。
- 兼容Flash Attention,可直接替换现有模型结构。
传统自注意力依赖softmax算子,需位置编码(如RoPE)来捕捉词序,但面临长序列泛化挑战。本文研究一种基于棒棒分段过程的新型注意力机制,在大规模设置下进行深入分析。该机制为每个历史词分配一个分割点,表示其占当前“棒棒”的比例,即注意力权重;在剩余部分重复此过程,直至所有词被分配权重,形成自然递减的注意力分布。该过程天然包含近邻偏好,具有语言学合理性。我们探讨了用棒棒分段注意力替代传统softmax+RoPE的可行性,并提出数值稳定实现方案,适配Flash Attention。作为现有softmax+RoPE系统的即插即用替代品,棒棒分段注意力在长序列泛化和下游任务上表现媲美主流方法。尤其在长度泛化方面表现突出:在2^11=2048上下文窗口训练的模型,可在2^14=16384长度上保持良好性能,困惑度显著降低。
原文摘要 · Abstract (English)
The self-attention mechanism traditionally relies on the softmax operator, necessitating positional embeddings like RoPE, or position biases to account for token order. But current methods using still face length generalisation challenges. We investigate an alternative attention mechanism based on the stick-breaking process in larger scale settings. The method works as follows: For each token before the current, we determine a break point, which represents the proportion of the stick, the weight of the attention, to allocate to the current token. We repeat this on the remaining stick, until all tokens are allocated a weight, resulting in a sequence of attention weights. This process naturally incorporates recency bias, which has linguistic motivations for grammar parsing. We study the implications of replacing the conventional softmax-based attention mechanism with stick-breaking attention. We then discuss implementation of numerically stable stick-breaking attention and adapt Flash Attention to accommodate this mechanism. When used as a drop-in replacement for current softmax+RoPE attention systems, we find that stick-breaking attention performs competitively with current methods on length generalisation and downstream tasks. Stick-breaking also performs well at length generalisation, allowing a model trained with $2^{11}$ context window to perform well at $2^{14}$ with perplexity improvements.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。