用可学习先验改进注意力机制，解决传统方法的表示缺陷。
You Need Better Attention Priors
- 将注意力视为熵正则最优传输问题，引入可训练连续先验。
- 在多个数据集上实现更好性能，且支持FlashAttention等优化核。
- 兼具学习位置嵌入的灵活性与固定编码的长度泛化能力。
我们通过熵正则最优传输（Entropic Optimal Transport, EOT）视角重新定义注意力机制，发现标准注意力对应于隐含均匀先验的运输问题。为此，提出可学习先验的广义最优传输注意力（GOAT），取代原始假设。该机制与FlashAttention等优化内核完全兼容，同时为注意力衰减现象提供理论解释并给出解决方案，避免了标准注意力中的表示权衡。此外，通过将空间信息融入核心注意力计算，GOAT学习到可外推的先验，结合了学习型位置嵌入的灵活性与固定编码的长度泛化优势。
原文摘要 · Abstract (English)
We generalize the attention mechanism by viewing it through the lens of Entropic Optimal Transport, revealing that standard attention corresponds to a transport problem regularized by an implicit uniform prior. We introduce Generalized Optimal transport Attention with Trainable priors (GOAT), a new attention mechanism that replaces this naive assumption with a learnable, continuous prior. This prior maintains full compatibility with optimized kernels such as FlashAttention. GOAT also provides an EOT-based explanation of attention sinks and materializes a solution for them, avoiding the representational trade-offs of standard attention. Finally, by absorbing spatial information into the core attention computation, GOAT learns an extrapolatable prior that combines the flexibility of learned positional embeddings with the length generalization of fixed encodings.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。