论文

MoSAR:学习可离散化的注意力距离机制

MoSAR: Mixture of Semantic Attention Regimes for Learning Adaptive and Approximable Attention Geometries

模型架构Transformer

摘要

密集自注意力的二次复杂度仍然是长上下文语言建模的中心瓶颈。许多有效的替代方案通过提前决定注意力应该稀疏或局部来解决这一成本。我们认为,注意力近似应该被视为一个几何问题,从数据中学习相关的交互几何:自然语言依赖性是依赖于输入的并且难以提前规定,因此模型应该学习位置相关性在哪里可以衰减以及在哪里必须保留更广泛的交互。我们引入了语义注意力机制混合(MoSAR),它通过查询-键交互来学习这种自适应的、受控衰减的几何结构。输入条件查询和键路由器在位置编码后应用,在短、中和全局范围内选择混合,从而诱导连续的距离相关的注意力场,而不是固定的稀疏模式。该几何形状是在训练期间学习的,随后可以通过 top-1 路由进行离散化。在使用匹配的 500M 参数模型的受控预训练实验中,MoSAR 在不降低语言建模质量的情况下学习了低范围注意力几何,从而在训练上下文长度上改善了密集 RoPE 的困惑度。在长度外推下,MoSAR 在所有评估的变体中实现了最佳的困惑度,包括 ALiBi 等强基线。此外,学习到的几何在确定性 top-1 离散化下保持稳定,这表明它不仅具有自适应性,而且适合推理时的低成本近似。

MoSAR不同查询—键机制组合的距离衰减曲线示例。
图1(图注摘要):MoSAR不同查询—键机制组合的距离衰减曲线示例。