用高斯核替代注意力中的线性投影,更简单高效。
Projection-Free Transformers via Gaussian Kernel Attention

- 用高斯核直接计算令牌相似度,仅需学习带宽参数。
- 参数减少58%,训练浮点运算量降为49%,仍稳定收敛。
- 适合追求轻量、可解释性的模型设计场景。
Transformer 中的自注意力通常通过 $\mathrm{softmax}(QK^\top/\sqrt{d})V$ 实现,其中 $Q=XW_Q$、$K=XW_K$、$V=XW_V$ 是输入 $X$ 的学习线性投影。本文探究这些投影是否必要,提出 extbf{高斯核注意力}(GKA),一种无需学习投影的点积注意力替代方案。它通过径向基函数(RBF)核直接计算每头令牌特征间的亲和度,每个头仅需学习带宽参数 $σ_h$,输出投影 $W_O$ 保持与标准 Transformer 接口兼容。GKA 可被理解为对令牌的归一化核回归,连接现代 Transformer 与经典非局部滤波及核平滑方法。在视觉与语言建模任务中评估表明:在 exttt{nanochat} 框架下,采用因果掩码与滑动窗口约束的 GKA 模型,在深度 20 时参数量仅为标准注意力的 0.42 倍,训练总浮点运算量为 0.49 倍,训练验证差距接近零,基准测试表现良好,但在此算力规模下比特每字节(BPB)略高。整体上,GKA 提供了一种极简、可解释的注意力机制,具备显式局部性尺度,为 Transformer 设计开辟了新的精度-效率权衡维度。
原文摘要 · Abstract (English)
Self-attention in Transformers is typically implemented as $\mathrm{softmax}(QK^\top/\sqrt{d})V$, where $Q=XW_Q$, $K=XW_K$, and $V=XW_V$ are learned linear projections of the input $X$. We ask whether these learned projections are necessary, or whether they can be replaced by a simpler similarity-based diffusion operator. We introduce \textbf{Gaussian Kernel Attention} (GKA), a drop-in replacement for dot-product attention that computes token affinities directly using a Gaussian radial basis function (RBF) kernel applied to per-head token features. Each head learns only a bandwidth parameter $σ_h$, while a single output projection $W_O$ preserves compatibility with the standard Transformer interface. GKA can be interpreted as normalized kernel regression over tokens, linking modern Transformer architectures to classical non-local filtering and kernel smoothing methods. We evaluate GKA in both vision and language modeling settings. For autoregressive language modeling within the \texttt{nanochat} framework, we implement causal masking and sliding-window constraints by masking and renormalizing the Gaussian kernel. At depth 20, a GKA model with $0.42\times$ the parameters and $0.49\times$ the total training FLOPs of a standard attention baseline trains stably, exhibits a near-zero train-validation gap, and demonstrates competitive behavior on standard benchmarks, albeit with higher bits-per-byte (BPB) at this compute scale. Overall, GKA provides a minimal, interpretable attention mechanism with an explicit locality scale, offering a dimension in the accuracy-efficiency trade-off for Transformer design.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。