Softmax Attention

分类: 深度学习基础

Softmax Attention

定义

Transformer 中标准的注意力机制，使用 softmax 函数对 query-key 点积进行归一化，生成注意力权重

数学形式

通过核方法可重写为对偶形式：

核心要点

softmax 对应一个无穷维 RKHS 中的核函数 $\kappa(\mathbf{x}, \mathbf{y}) = \exp(\mathbf{x}\mathbf{y}^\top / \sqrt{d})$

对偶形式揭示每个 token 对”隐式权重矩阵”的 rank-1 贡献

计算复杂度 $O(N^2 d)$，是长序列处理的主要瓶颈

代表工作

IWP: 利用 softmax 的核展开推导 token pruning 的理论指标

FlashAttention: IO-aware 的高效 softmax attention 实现

相关概念

Linear Attention

Multi-Head Attention

核方法

RKHS

FlashAttention