自由能混合器

标准注意力机制在存储键/值方面保留所有信息，但通过每头的凸组合读取，限制了通道级选择。我们提出一种自由能混合器 (FEM)：一种基于自由能 (log-sum-exp) 的读取机制，对快速先验 (例如来自查询/键的标准注意力) 施加值驱动的、每通道的对数线性倾斜。不同于试图改进和丰富 $(q,k)$ 评分分布的方法，FEM 将其视为先验，并在保持不变复杂度的同时，提供一种值感知的后验读取，平滑地从平均化过渡到每通道选择，learnable 逆温度增加时仍保持并行性和原始渐近复杂度 ($O(T^2)$ 用于 softmax；$O(T)$ 用于可线性化变体)。我们实现了一种两级门控 FEM，可无缝集成到标准注意力、线性注意力、线性 RNN 和 SSMs 中。在 NLP、视觉和时间序列任务中，FEM 在匹配参数预算下一致性地优于强基准。