Transformer模型越来越多地被用作认知和神经表征的计算模型,因此自注意力所实现的机制不仅关乎工程性能,也引起了关注。认知科学中有一个互补的传统,通过振荡器同步等动态交互来建模协调、绑定和记忆;我们将这一机制引入自注意力,提出了Kuramoto注意力层,其值更新是一个同步步骤。每个token携带一组相位振荡器,因此其隐藏状态存在于高维环面上。注意力权重形成自适应耦合图,使用原始相位状态作为值,使得在固定注意力权重下,值更新恰好是Kuramoto耦合方向。softmax选择哪些振荡器耦合,而值路径将每个token移向其所选token的注意力加权圆均值。我们在enwiki8和CodeParrot数据集上训练Kuramoto注意力,与参数匹配的RoPE和SwiGLU Transformer进行比较。在CodeParrot上5M参数时,它在中位数和均值上都优于Transformer,均值差距为0.012验证和0.010测试比特每字节。在enwiki8上5M时,所有六次运行的验证/测试中位数均低于Transformer,所有种子的均值在0.01 BPC以内;其中五次运行还形成了一个紧密的低均值簇。在1M参数时,它在enwiki8上落后约0.02 BPC,在CodeParrot上落后0.013-0.015比特每字节。消融实验和相位诊断显示了该层的同步和几何动机组件如何影响模型性能。结果是一种自注意力机制,其学习到的计算可以直接解读为相位状态上的自适应同步。
核心观点
- 提出了Kuramoto注意力层,将振荡器同步机制引入自注意力,值更新对应于同步步骤
- 每个token的隐藏状态位于高维环面上,注意力权重形成自适应耦合图
- 在enwiki8和CodeParrot上,5M参数模型在验证和测试指标上优于参数匹配的Transformer基线
- 1M参数时性能略逊于基线,但差距较小
- 消融实验和相位诊断揭示了同步和几何组分对模型性能的影响
技术要点
该方法利用相位振荡器的同步动力学,将自注意力的值更新建模为Kuramoto耦合方向,通过softmax选择耦合的振荡器,并使用圆均值进行聚合。训练中使用与RoPE和SwiGLU Transformer相同的参数设置,进行对比实验。
应用场景
适用于需要建模序列中长程依赖和动态交互的任务,如语言建模、代码生成等,尤其适合作为认知和神经科学的计算模型。