你需要更好的注意力先验

本文通过熵最优传输的视角对注意力机制进行了泛化，揭示了标准注意力对应于一个由隐式均匀先验正则化的传输问题。我们引入了一种新的注意力机制——具有可训练先验的广义最优传输注意力（GOAT），用可学习的连续先验替代了这一朴素假设。该先验与优化后的内核（如FlashAttention）完全兼容。GOAT还提供了基于熵最优传输的注意力陷阱解释，并提出了解决方案，避免了标准注意力的表示权衡。最后，通过将空间信息融入核心注意力计算中，GOAT学习到了一种可外推的先验，结合了学习位置嵌入的灵活性和固定编码的长度泛化能力。