3、Softmax与Attention权重：从原理到实战

Attention机制里，Softmax是个绕不开的角色。说白了，它就是把一堆分数变成概率分布的那个函数。我刚开始接触Transformer时，总觉得Softmax就是个简单的归一化，后来踩了不少坑才发现——这里面的门道比想象中深得多。

3.1 Softmax函数的作用与特性

Softmax的核心作用，就是把任意实数向量映射成概率分布。公式很简单：

Softmax(x_i) = exp(x_i) / Σ exp(x_j)

但它的特性值得细说：

• 输出非负且和为1：这是概率分布的基本要求
• 保持相对顺序：输入大的，输出也大
• 指数放大效应：差异会被指数函数放大

我个人习惯把Softmax看作一个「竞争机制」。每个元素都在争夺概率份额，谁的值大，谁就能拿到更多。你想想看，在Attention里，Query和Key计算出的分数经过Softmax后，就变成了对Value的加权权重——分数高的位置，模型会「多看几眼」。

3.2 温度系数：控制注意力的「软硬」程度

温度系数τ（tau）是调节Softmax分布形态的关键参数。带温度的Softmax长这样：

Softmax(x_i / τ) = exp(x_i / τ) / Σ exp(x_j / τ)

τ的作用很直观：

• τ → 0：分布趋近于one-hot，注意力极度集中（硬注意力）
• τ → ∞：分布趋近于均匀分布，注意力分散（软注意力）
• τ = 1：标准Softmax

我在项目中遇到过一个问题：训练初期，模型注意力分布太尖锐，导致梯度稀疏，收敛很慢。后来我尝试把温度系数设到1.5~2.0，让注意力先「软」一点，等模型学到一定基础后再逐步降低温度。效果立竿见影。

3.3 稀疏化与熵正则化

标准Softmax有个问题：它永远不会输出真正的零。每个位置都会分到一点概率，哪怕那个位置完全不相关。这带来了计算浪费和噪声。

稀疏化就是让注意力权重中不重要的位置变成真正的零。常见方法有：

• Top-k Softmax：只保留分数最高的k个位置，其余置零
• α-entmax：用α-entmax替代Softmax，天然产生稀疏解
• 硬阈值：设定阈值，低于阈值的直接置零

熵正则化则是从另一个角度入手。注意力分布的熵反映了它的「混乱程度」：

• 熵低 → 注意力集中（尖锐分布）
• 熵高 → 注意力分散（均匀分布）

我曾在长文本摘要任务中，给注意力熵加了一个正则项，鼓励模型关注更少但更关键的位置。嗯，这里要注意：正则系数不能太大，否则模型会变得过于「偏执」，忽略掉一些有用的上下文信息。

3.4 Softmax数值溢出：一个让人头疼的问题

这是我在实际工程中踩过最深的坑之一。先看个例子：

import numpy as np

def softmax_naive(x):
    return np.exp(x) / np.sum(np.exp(x))

# 输入很大的值
x = np.array([1000, 1001, 1002])
print(softmax_naive(x))  # 输出 [nan, nan, nan]

为什么会这样？因为exp(1000)已经远远超出了float32的表示范围（约3.4e38），直接溢出成了inf。inf除以inf就是nan。

解决办法其实很简单——减去最大值：

def softmax_stable(x):
    x_max = np.max(x)
    x_shifted = x - x_max
    return np.exp(x_shifted) / np.sum(np.exp(x_shifted))

x = np.array([1000, 1001, 1002])
print(softmax_stable(x))  # 输出正常

这个技巧的原理是：Softmax对输入同时加减一个常数，结果不变。减去最大值后，所有输入都≤0，exp的值都在(0, 1]范围内，彻底避免了溢出。

但你以为这就完了？不，还有另一个坑——下溢出。当输入都是很大的负数时，比如[-1000, -999, -998]，减去最大值后变成[-2, -1, 0]，exp后是[0.135, 0.368, 1.0]，没问题。但如果输入差异极大，比如[0, -1e5, -2e5]，减去最大值后变成[0, -1e5, -2e5]，exp(-1e5)直接下溢出为0。虽然不会报错，但梯度全没了。

另外，混合精度训练（FP16）会加剧这个问题。FP16的表示范围更小，更容易溢出。我建议在Attention计算中使用FP32累加，或者对QK^T的结果做缩放（除以√d_k），这既是数值稳定的需要，也是Attention论文里的标准做法。

3.5 本章知识体系

下面这张图总结了Softmax在Attention中的角色和优化方向：

这张图把Softmax在Attention中的四个关键维度串起来了。从基础特性到温度调节，从稀疏优化到数值稳定性，每个环节都有实际工程中需要注意的细节。

---

公众号:蓝海资料掘金营，微信deep3321