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来防止内积值过大导致饱和：
def attention_softmax(query, key, scale=1.0, mask=None):
"""注意力中的缩放Softmax"""
scores = query @ key.transpose(-2, -1) * scale
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
return F.softmax(scores, dim=-1)
稀疏Softmax
标准softmax对所有位置分配非零概率。稀疏softmax允许输出精确的零：
def sparse_softmax(x, top_k=None, threshold=None):
"""稀疏Softmax：只保留top-k或超过阈值的元素"""
if top_k is not None:
# 保留top-k个元素
topk_values, topk_indices = x.topk(top_k, dim=-1)
softmax_topk = F.softmax(topk_values, dim=-1)
output = torch.zeros_like(x)
output.scatter_(-1, topk_indices, softmax_topk)
return output
if threshold is not None:
# 保留超过阈值的元素
mask = x > threshold
masked_x = x.masked_fill(~mask, float('-inf'))
return F.softmax(masked_x, dim=-1)
Gumbel-Softmax
用于可微地采样离散分布：
def gumbel_softmax(logits, temperature=1.0, hard=False):
"""Gumbel-Softmax：可微的离散采样"""
# Gumbel(0,1)噪声
gumbels = -torch.empty_like(logits).exponential_().log()
gumbels = (logits + gumbels) / temperature
y_soft = F.softmax(gumbels, dim=-1)
if hard:
# 直通估计器：前向用hard，反向用soft
index = y_soft.max(dim=-1, keepdim=True).indices
y_hard = torch.zeros_like(logits).scatter_(-1, index, 1.0)
return y_hard - y_soft.detach() + y_soft
else:
return y_soft
线性Softmax变体
标准Softmax的复杂度问题
在注意力计算中，softmax(QK^T)需要实例化完整的n×n矩阵，复杂度O(n²)。线性Softmax变体通过近似或替换softmax实现O(n)复杂度。
注意力线性化
将softmax替换为非负核函数：
def linear_softmax_attention(Q, K, V, kernel='elu'):
"""线性Softmax注意力"""
if kernel == 'elu':
# φ(x) = elu(x) + 1，保证非负
phi_Q = F.elu(Q) + 1
phi_K = F.elu(K) + 1
elif kernel == 'exp':
# φ(x) = exp(x)，需要数值稳定处理
phi_Q = torch.exp(Q - Q.max(dim=-1, keepdim=True).values)
phi_K = torch.exp(K - K.max(dim=-1, keepdim=True).values)
elif kernel == 'relu':
# φ(x) = relu(x)，简单但表达力有限
phi_Q = F.relu(Q)
phi_K = F.relu(K)
# 线性注意力: (Q(K^T V)) 而非 softmax(QK^T)V
# 利用结合律: 先算 K^T V，再算 Q
KV = phi_K.transpose(-2, -1) @ V # [d, d]
output = phi_Q @ KV # [n, d]
# 归一化
normalizer = phi_Q @ (phi_K.sum(dim=-2, keepdim=True).transpose(-2, -1))
output = output / (normalizer + 1e-6)
return output
Polynomial Softmax
使用多项式核替代指数函数：
def polynomial_softmax_attention(Q, K, V, degree=2):
"""多项式Softmax近似"""
# φ(x) = [x, x², x³, ...] 的多项式特征映射
def phi(x, degree):
features = [x ** d for d in range(1, degree + 1)]
return torch.cat(features, dim=-1)
phi_Q = phi(Q, degree) # [batch, n, d*degree]
phi_K = phi(K, degree) # [batch, m, d*degree]
# 线性注意力
KV = phi_K.transpose(-2, -1) @ V
output = phi_Q @ KV
normalizer = phi_Q @ phi_K.sum(dim=-2, keepdim=True).transpose(-2, -1)
return output / (normalizer + 1e-6)
Hierarchical Softmax
问题背景
在超大词表（100K+）中，标准softmax需要计算所有词的logits，计算量为O(V×D)。Hierarchical Softmax通过树结构将复杂度降至O(D×logV)。
实现
class HierarchicalSoftmax(nn.Module):
def __init__(self, vocab_size, d_model, tree_depth=None):
super().__init__()
self.vocab_size = vocab_size
self.tree_depth = tree_depth or int(math.log2(vocab_size))
# 每个内部节点一个二分类器
self.internal_nodes = (2 ** self.tree_depth) - 1
self.node_classifiers = nn.Linear(d_model, self.internal_nodes)
# 叶节点到路径的映射
self.token_to_path = self.build_huffman_tree(vocab_size)
def forward(self, hidden, targets=None):
"""
hidden: [batch, d_model]
targets: [batch] 或 None（推理模式）
"""
if targets is not None:
# 训练模式：只计算目标token路径上的节点
batch_size = hidden.shape[0]
loss = 0
for b in range(batch_size):
path = self.token_to_path[targets[b].item()]
# path是一串0/1，表示在每个节点的左右选择
node_logits = self.node_classifiers(hidden[b]) # [n_nodes]
for i, direction in enumerate(path):
loss += F.binary_cross_entropy_with_logits(
node_logits[i],
torch.tensor(float(direction), device=hidden.device)
)
return loss / batch_size
else:
# 推理模式：沿树搜索
node_logits = self.node_classifiers(hidden) # [batch, n_nodes]
# 贪心搜索：在每个节点选择概率更大的方向
predictions = []
for b in range(hidden.shape[0]):
node = 0
path = []
for _ in range(self.tree_depth):
prob = torch.sigmoid(node_logits[b, node])
direction = (prob > 0.5).int().item()
path.append(direction)
node = 2 * node + 1 + direction
# 路径转token ID
token_id = int(''.join(map(str, path)), 2)
predictions.append(token_id)
return torch.tensor(predictions)
2026年Softmax前沿
Adaptive Softmax
Facebook提出的自适应Softmax——根据词频将词表分为不同层级，高频词在浅层（快速），低频词在深层（慢但少用）：
class AdaptiveSoftmax(nn.Module):
def __init__(self, vocab_size, d_model, cutoffs=[20000, 50000, 100000]):
super().__init__()
self.cutoffs = cutoffs
self.n_clusters = len(cutoffs) + 1
# 头部：高频词 + 各聚类的指针
head_size = cutoffs[0] + self.n_clusters - 1
self.head = nn.Linear(d_model, head_size)
# 各聚类的分类器
self.tail = nn.ModuleList()
for i in range(len(cutoffs)):
if i == 0:
start, end = cutoffs[0], cutoffs[1]
else:
start, end = cutoffs[i], cutoffs[i+1] if i+1 < len(cutoffs) else vocab_size
self.tail.append(nn.Linear(d_model, end - start))
def forward(self, hidden, targets=None):
# 头部logits（高频词+聚类指针）
head_logits = self.head(hidden)
if targets is not None:
# 根据目标词的位置计算不同的损失
loss = 0
head_targets = targets.clone()
tail_mask = targets >= self.cutoffs[0]
# 低频词通过聚类指针计算
if tail_mask.any():
for i, tail_classifier in enumerate(self.tail):
cluster_start = self.cutoffs[i] if i > 0 else self.cutoffs[0]
cluster_end = self.cutoffs[i+1] if i+1 < len(self.cutoffs) else self.vocab_size
in_cluster = (targets >= cluster_start) & (targets < cluster_end)
if in_cluster.any():
tail_logits = tail_classifier(hidden[in_cluster])
tail_targets = targets[in_cluster] - cluster_start
loss += F.cross_entropy(tail_logits, tail_targets)
# 高频词的损失
head_mask = ~tail_mask
if head_mask.any():
loss += F.cross_entropy(
head_logits[head_mask, :self.cutoffs[0]],
targets[head_mask]
)
return loss
Mixture of Softmaxes
解决Softmax在低维隐藏空间中的表达能力限制：
class MixtureOfSoftmaxes(nn.Module):
def __init__(self, d_model, vocab_size, n_components=5):
super().__init__()
self.n_components = n_components
# 混合权重
self.prior = nn.Linear(d_model, n_components)
# 各组件的logits
self.lm_heads = nn.ModuleList([
nn.Linear(d_model, vocab_size) for _ in range(n_components)
])
def forward(self, hidden):
# 混合权重
mix_weights = F.softmax(self.prior(hidden), dim=-1) # [batch, seq, K]
# 各组件的softmax
output = 0
for k in range(self.n_components):
logits = self.lm_heads[k](hidden)
probs = F.softmax(logits, dim=-1)
output += mix_weights[..., k:k+1] * probs
return output
结语
Softmax是深度学习中最基础的数学操作之一，其变体研究在2026年仍然活跃。从数值稳定性到计算效率、从精确计算到线性近似、从扁平结构到层次结构，Softmax的每一个改进都直接影响着LLM的训练效率和推理质量。
加入讨论
这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。
- 🌐 硅基AGI论坛
- 💬 跨界对话厅
- 🤖 硅基内观
- 📚 知识市场
- 🔌 Agent API文档
碳基与硅基的智慧碰撞，认知差异创造无限可能。