技术博客

缩放点积的 1/√d:方差视角的推导

第 46 / 147 篇 · 2026-07-20 · ai · 注意力机制数值稳定性概率基础

注意力公式里的 $1/\sqrt d$ 常被当成"经验调出来的缩放"。它不是。只要对初始化时的 $q, k$ 做一个最朴素的方差假设,这个常数就会自己掉出来,而且你能同时看到不除会发生什么、除错了会发生什么。

推导:点积的方差随 d 线性增长

设 $q, k \in \mathbb{R}^d$ 的分量独立同分布,均值为 0、方差为 $\sigma^2$(Xavier 类初始化下这近似成立,且模型在训练中大部分时间的尺度也在这附近)。点积

$$ s = q^\top k = \sum_{a=1}^{d} q_a k_a $$

是 $d$ 个独立乘积项的和。每项的期望 $\mathbb{E}[q_a k_a]=0$,方差

$$ \mathrm{Var}(q_a k_a) = \mathbb{E}[q_a^2]\mathbb{E}[k_a^2] = \sigma^4 $$

(独立性让协方差项消失)。独立和的方差相加:

$$ \mathrm{Var}(s) = \sum_{a=1}^d \sigma^4 = d\,\sigma^4 \quad\Longrightarrow\quad \mathrm{std}(s) = \sqrt d\,\sigma^2. $$

注意方差恰好是 $d$ 的一次函数,标准差恰好是 $\sqrt d$ 的倍数——分母 $\sqrt d$ 不是在调某个神秘超参,它是在做量纲归一:$s/\sqrt d$ 的方差回到 $\sigma^4$,与维度 $d$ 无关。

不除会发生什么:softmax 饱和

为什么 logits 尺度是生死问题?softmax 的梯度性质由 logit 间差距决定。两个 logits 差 $\Delta$,softmax 输出的比值是 $e^{\Delta}$。$\mathrm{std}(s)=\sqrt d\,\sigma^2$ 意味着典型 logit 差距是 $\sqrt d$ 量级:$d=1024$ 时 logits 散布在几十的标准差上,最大的那个 logit 会赢出 $e^{30}$ 级别的倍率——softmax 输出接近 one-hot,这就是饱和。

饱和有两个直接后果。其一,输出对 value 的混合退化成"只取 argmax 那一个",注意力的"软检索"能力名存实亡。其二,梯度消失:softmax 的雅可比是 $\mathrm{diag}(p)-pp^\top$,当 $p$ 接近 one-hot 时它接近零矩阵,$\partial \ell/\partial s \approx 0$,$W^Q, W^K$ 收不到有效梯度。训练初期恰好是表征最混乱、最需要梯度的阶段,初始化尺度下不除 $1/\sqrt d$ 等于开局自废。

为什么不是 1/d

一个自然的疑问:既然 std 是 $\sqrt d$ 倍,那把方差归一(除 $d$)不是更彻底?不行,因为要控制的是 logits 的尺度(标准差),不是方差。softmax 只看得见 logits 的差值,它的敏感区间由 std 刻画:std 太小(除 $d$),所有 logit 挤在一起,softmax 输出趋于均匀分布 $p_j \approx 1/m$,每个 value 被平均混合,注意力同样失去选择性——只是换成了"太钝"的失败模式。除 $\sqrt d$ 把 std 钉回 $\sigma^2$ 量级,是"太尖"和"太钝"中间唯一与 $d$ 无关的稳态点。

这个推理也解释了为什么温度参数(temperature)的存在是合理的:$\mathrm{softmax}(s/\tau)$ 里的 $\tau$ 就是留给模型后期再调节这个 std 的自由度,训练充分后最优尺度未必还是初始化假设下的 $\sqrt d$。

一个补充视角:玻尔兹曼分布的量纲

softmax 源自统计物理的玻尔兹曼分布 $p_i \propto e^{-E_i/kT}$,其中 $E_i/kT$ 必须无量纲。logits 扮演的角色是"负能量",$1/\sqrt d$ 相当于给能量选了一个与系统尺寸(维度 $d$)无关的单位制,保证"温度"(训练中学到的 logit 尺度)有稳定的语义。这和概率篇里 temperature 的讨论是同一枚硬币的两面。

s = q·k 的分布随 d 展宽:std = √d·σ²,除 √d 后回到与 d 无关的尺度 d 小:std 窄 d 大:std = √d·σ² 除以 √d 后两者重合 Var(q·k) = d·σ⁴ Var(s/√d) = σ⁴(与 d 无关) logits 差 Δ → 输出比 e^Δ

小结

$1/\sqrt d$ 的来源一句话:初始化假设下 $\mathrm{std}(q^\top k)=\sqrt d\,\sigma^2$,除以 $\sqrt d$ 是把 logit 标准差归一到与维度无关的量级,躲开 softmax 饱和导致的软检索退化和梯度消失。它归一的是标准差而非方差,所以分母是 $\sqrt d$ 而不是 $d$。这个推导还附带一个判断标准:任何对注意力分数的改动(相对位置偏置、ALiBi、可学习温度),都可以问一句"它把 logit 的 std 改成了多少"——方差账是这条线路上最可靠的审计工具。