Softmax

Softmax：将分数变为概率 Softmax函数将任意实数向量转换为概率分布——所有元素为正且和为1。它是分类、注意力、语言模型等核心组件的数学基础。 Softmax(x_i) = exp(x_i) / Σ_j exp(x_j) 看似简单的公式背后，隐藏着数值稳定性、计算效率和梯度行为的深刻问题。 数值稳定性 溢出问题 当输入值很大时，exp(x)会溢出。例如exp(1000)远超FP32的范围（~3.4e38）。解决方案是减去最大值： def softmax_stable(x, dim=-1): """数值稳定的Softmax""" x_max = x.max(dim=dim, keepdim=True).values exp_x = torch.exp(x - x_max) return exp_x / exp_x.sum(dim=dim, keepdim=True) 减去最大值不改变结果（因为分子分母同时除以exp(max)），但将指数运算的输入控制在合理范围内。 Log-Softmax 在许多场景中（如交叉熵损失），我们需要的是log概率而非概率本身。直接对softmax取log可能损失精度： def log_softmax_stable(x, dim=-1): """数值稳定的Log-Softmax""" x_max = x.max(dim=dim, keepdim=True).values shifted = x - x_max log_sum_exp = torch.log(torch.exp(shifted).sum(dim=dim, keepdim=True)) return shifted - log_sum_exp PyTorch的F.log_softmax和F.cross_entropy内部都使用了这种稳定的实现。 Softmax的梯度特性 饱和问题 当某个输入远大于其他输入时，Softmax的输出接近one-hot——一个接近1，其余接近0。此时梯度几乎为零，导致学习停滞。 # 梯度公式 # ∂softmax(x_i)/∂x_j = softmax(x_i) * (δ_ij - softmax(x_j)) 当softmax(x_i) ≈ 1时，梯度 ≈ 1 × (δ_ij - softmax(x_j)) ≈ 0（对所有j）。 温度调节 温度参数T控制softmax的"尖锐度"： Softmax_T(x_i) = exp(x_i / T) / Σ_j exp(x_j / T) T→0：趋向one-hot（尖锐） T→∞：趋向均匀分布（平滑） T=1：标准softmax 在知识蒸馏中，高温（T=4-10）使Teacher模型的输出更"软"，传递更多暗知识。 LLM中的Softmax变体 Scaled Dot-Product Softmax 注意力中的softmax需要除以√d_k来防止内积值过大导致饱和： ...