基于高斯核注意力的投影自由 Transformer

Transformer中的自注意力通常实现为 $\mathrm{softmax}(QK^\top/\sqrt{d})V$，其中 $Q=XW_Q$、$K=XW_K$、$V=XW_V$ 是输入 $X$ 的学习线性投影。我们询问这些学习投影是否必要，或者是否可以被更简单的数据相似性扩散算子取代。我们提出高斯核注意力(Gaussian Kernel Attention, GKA)，这是一种用于替换点积注意力的方案，通过对每个头的 token 特征应用高斯径向基函数(RBF)核来直接计算 token 之间的相似度。每个头仅学习一个带宽参数 $\sigma_h$，而单一的输出投影 $W_O$ 保持与标准 Transformer 接口的兼容性。GKA可解释为对 token 进行归一化核回归，链接现代 Transformer 架构与经典的非局部滤波和核平滑方法。我们在视觉和语言建模任务中评估了GKA。对于基于 \texttt{nanochat} 框架的自回归语言建模，我们通过屏蔽和重新归一化高斯核来实现因果掩码和滑动窗口约束。在深度为 20 时，一个参数量为标准注意力基线的 $0.42\times$ 且总训练 FLOPs 为 $0.49\times$ 的 GKA 模型能够稳定训练，呈现接近零的训练-验证间隙，并在标准基准上表现出竞争力，尽管在此计算规模下表现出更高的比特/字节(Bits-per-byte, BPB)。总体而言，GKA 提供了一种最小且可解释的注意力机制，具有明确的局部性尺度，为 Transformer 设计提供了准确率与效率之间一个维度的权衡。