Softmax 的来历与 LLM 应用

0. 先看结论

Softmax 做的事很直接：把一组分数变成概率分布。它的思想可追溯到统计物理里的 Boltzmann/Gibbs 分布，后来进入选择理论、统计学习和神经网络。到了 LLM 里，最常见的是两个位置：

1. 输出层：把模型对所有 token 的打分 logits 变成“下一个 token 的概率”。
2. 注意力层：把 query 和 key 的相似度分数变成“该看哪些上下文 token 的权重”。

在 MoE 模型里，softmax 还可能出现在 router 中，给不同 expert 分配路由权重。

flowchart LR
    A[原始分数 logits] --> B[指数放大差异 exp]
    B --> C[归一化: 除以总和]
    C --> D[概率分布: 每项 > 0, 总和 = 1]

    D --> E[输出层: 预测下一个 token]
    D --> F[注意力层: 给上下文 token 分配权重]

0.1 名词速查

刚接触 LLM 的话，可以先看这张表。后面遇到术语，再回来查。

基础数学词：

名词	通俗解释	在本文中的作用
分数	模型给某个选项的原始评价	softmax 的输入
logit	softmax 前的原始分数	token 或位置的未归一化偏好
概率	0 到 1 之间的数，总和为 1	softmax 的输出
概率分布	一组候选项各自的概率	表示“每个选项有多可能”
归一化	把一组数缩放到总和为 1	softmax 分母的作用
指数函数 exp	把数变成 $e^x$	放大高分和低分的差距
log-sum-exp	$\log \sum_j e^{z_j}$	稳定计算 log softmax 的核心技巧
sigmoid	把一个分数压到 0 到 1	常用于二分类或多标签，不要求总和为 1
temperature	温度参数	控制概率分布尖锐或平滑
cross entropy	交叉熵损失	训练时惩罚“真实 token 概率低”
perplexity	困惑度	语言模型常见评估指标，来自交叉熵

LLM 结构词：

名词	通俗解释	在本文中的作用
token	模型处理文本的基本单位	LLM 每次预测下一个 token
vocabulary	词表，所有可选 token 的集合	输出层要给每个 token 打分
hidden state	模型内部的上下文表示	用来生成 vocab logits
embedding	token 的向量表示	把离散 token 变成模型可计算的向量
output matrix	输出矩阵	把 hidden state 映射成词表 logits
attention	注意力机制	决定当前位置该看哪些上下文位置
Q/K/V	query、key、value 三组向量	计算注意力分数并聚合信息
attention score	注意力分数	softmax 前的位置相关性分数
mask	屏蔽规则	禁止模型看不该看的位置
causal mask	因果遮罩	自回归生成时屏蔽未来 token

解码与工程词：

名词	通俗解释	在本文中的作用
greedy	每次选最大概率 token	最确定的解码方式
top-k	只保留概率最高的 k 个候选	限制低质量长尾 token
top-p	保留累计概率达到 p 的候选集合	根据不确定性动态调候选数
logit bias	人为调整某些 token 的 logit	鼓励或禁止特定输出
repetition penalty	重复惩罚	降低已出现 token 的再次出现概率
full softmax	对整个词表完整计算 softmax	准确但大词表成本高
softmax bottleneck	softmax 输出表达能力限制	解释语言模型输出层的理论瓶颈
label smoothing	标签平滑	训练时不把正确答案当成 100% one-hot
MoE router	专家路由器	在 MoE 模型中为 token 选择专家
expert	专家网络	MoE 中被路由选中的子网络
FlashAttention	高效 attention 实现	不改结果，减少显存读写
HBM	GPU 高带宽显存	长上下文 attention 的主要瓶颈之一

---

1. Softmax 是什么

给定一组实数分数：

$$z = [z_1, z_2, ..., z_K]$$

softmax 输出第 $i$ 个类别的概率：

$$\mathrm{softmax}(z_i)=\frac{e^{z_i}}{\sum_{j=1}^{K}e^{z_j}}$$

带温度参数的版本：

$$\mathrm{softmax}(z_i/T)=\frac{e^{z_i/T}}{\sum_{j=1}^{K}e^{z_j/T}}$$

其中：

符号	含义	说明
$z_i$	第 $i$ 个选项的原始分数，也叫 logit	分数越大，概率越大
$e^{z_i}$	指数变换	把差距放大，并保证为正数
$\sum_j e^{z_j}$	归一化项，也类似物理里的 partition function	让所有概率加起来等于 1
$T$	temperature，温度	低温更“贪心”，高温更“发散”

1.1 logits 是什么

logits 指 模型在 softmax 之前给每个候选项打的原始分数。

它不是概率，也不需要满足概率的约束：

对比项	logits	概率
取值范围	可以是任意实数，例如 -12.3、0、27.8	必须在 0 到 1 之间
是否要相加为 1	不需要	必须总和为 1
含义	相对偏好、相对证据、未归一化分数	可解释为抽样概率
进入 softmax 前后	softmax 之前	softmax 之后

在 LLM 里，模型每次预测下一个 token 时，会给词表里的每个 token 都算一个 logit：

上下文: "今天天气很"

token 候选      logit      softmax 后概率
---------------------------------------
"好"            8.2        0.51
"热"            7.5        0.25
"冷"            6.4        0.08
"苹果"         -1.3        很小

也可以把 logits 当成“还没换算成概率的投票分”：

flowchart LR
    A[模型内部表示 hidden state] --> B[线性层打分]
    B --> C[logits: 原始分数]
    C --> D[softmax]
    D --> E[概率分布]

几个容易混淆的点：

问题	答案
logit 越大是否概率一定越大？	是。softmax 保序，最大 logit 对应最大概率。
logit 是不是百分比？	不是。8.2 不代表 82%。
logit 可以是负数吗？	可以。负数只表示相对分数低，不代表概率为负。
为什么不直接输出概率？	任意概率分布不方便直接训练；模型先输出自由实数，再由 softmax 转成概率。
采样时一定要显式算 softmax 吗？	不一定。greedy 只要找最大 logit；但随机采样、top-p 等通常需要概率或等价处理。

1.2 一个极小例子

假设三个 token 的 logits 是：

token	logit
猫	3.0
狗	2.0
车	0.0

softmax 后大致为：

token	softmax 概率	说明
猫	0.705	最高分，概率最大
狗	0.259	次高分，仍有机会
车	0.035	分数低，但不是 0

token	logit 条形	概率条形
猫	████████████ 3.0	██████████████ 0.705
狗	████████ 2.0	█████ 0.259
车	0.0	█ 0.035

logit 只差 1，softmax 概率可能差很多。

---

2. 为什么叫“soft max”

max 的意思是“只选最大的那个”。例如：

[3, 2, 0] -> max 选择第 1 项

神经网络训练需要可导函数，不能直接用“硬选择”。softmax 给出的是“软选择”：

[3, 2, 0] -> [0.705, 0.259, 0.035]

它偏向最大值，同时保留其他选项的非零概率，所以能参与梯度下降。

flowchart LR
    A[硬 max] --> B[只选最大项]
    B --> C[不适合反向传播]

    D[softmax] --> E[最大项权重大]
    D --> F[其他项仍有概率]
    F --> G[平滑可导]

---

3. 历史由来：从物理到 LLM

Softmax 不是专门为神经网络凭空发明的。更贴近历史的说法是：归一化指数分布在多个学科里反复出现；“softmax”这个名字是在神经网络语境里流行起来的。

3.1 时间线

时间	关键事件	意义
1868	Boltzmann 分布	指数形式描述状态概率
1902	Gibbs 统计力学	partition function 等概念系统化
1959	Luce 选择公理	用相对偏好解释选择概率
1989	Bridle NIPS 论文	softmax 术语在神经网络语境中扩散
2001/2003	Bengio 神经概率语言模型	softmax 输出整个词表的下一个词概率
2005	hierarchical softmax	缓解大词表 full softmax 成本
2013	word2vec 负采样	用近似目标绕开大规模 softmax
2017	Transformer	attention 中使用 scaled dot-product softmax
2017/2018	adaptive softmax / softmax bottleneck	系统讨论大词表效率与表达能力
2019/2020	nucleus sampling	基于 softmax 概率做动态截断采样
2022	FlashAttention	不改变结果，用 IO-aware 算法高效计算 attention softmax

3.2 各阶段对照

阶段	代表工作	今天的对应关系
统计物理	Boltzmann, Gibbs	“指数打分 + 归一化”的原型
选择理论	R. Duncan Luce	相对偏好转选择概率
神经网络	John S. Bridle	多分类输出层概率化
神经语言模型	Bengio 等	LLM 输出层的前身
Transformer	Vaswani 等	LLM 注意力权重的核心公式
长上下文工程	FlashAttention	exact softmax attention 的高效实现

几个常见公式形态：

场景	公式
统计物理	$p_i \propto e^{-E_i/kT}$
选择理论	$P(i)=u_i/\sum_j u_j$
神经网络/LLM	$P(i)=e^{z_i}/\sum_j e^{z_j}$
Transformer attention	$\mathrm{softmax}(QK^T/\sqrt{d_k})V$

---

4. 三条来源线索

4.1 物理线：Boltzmann/Gibbs 分布

统计物理里，一个系统处在某个能量状态 $E_i$ 的概率常写成：

$$p_i = \frac{e^{-E_i/kT}}{\sum_j e^{-E_j/kT}}$$

这和 softmax 几乎同形：

统计物理	机器学习
能量 $E_i$ 越低，概率越高	logit $z_i$ 越高，概率越高
温度 $T$ 控制分布平滑程度	temperature 控制采样随机性
partition function 负责归一化	softmax 分母负责归一化

把 $z_i=-E_i/kT$，Boltzmann 分布就可以看成 softmax。

flowchart LR
    A[物理: 能量 E] --> B[取负并除以温度]
    B --> C[指数 e^score]
    C --> D[除以总和]
    D --> E[状态概率]

    F[ML: logit z] --> C

4.2 选择理论线：Luce 选择公理

Luce 在 1959 年的《Individual Choice Behavior》里研究“人在多个选项中如何随机选择”。如果每个选项有正的吸引力 $u_i$，可以写成：

$$P(i)=\frac{u_i}{\sum_j u_j}$$

如果令 $u_i=e^{z_i}$，得到：

$$P(i)=\frac{e^{z_i}}{\sum_j e^{z_j}}$$

也就是 softmax。
所以 softmax 既是数学技巧，也可以看作“相对偏好转概率”的模型。

4.3 神经网络线：Bridle 与 softmax 命名

在 1989 年 NIPS 论文中，John S. Bridle 将随机识别模型与神经网络训练联系起来，使用了归一化指数形式。随后 1990 年相关工作明确把神经网络分类输出解释为概率，使用 normalized exponential，也就是后来常说的 softmax。

从这以后，softmax 成为多分类神经网络的标准输出层之一。

---

5. Softmax 的几个性质

5.1 输出是合法概率分布

性质	说明
非负	$e^{z_i}>0$，所以每个概率都大于 0
归一化	分母是所有指数值之和，所以总和为 1
保序	logit 越大，softmax 概率越大
可导	能参与反向传播

5.2 平移不变性

对所有 logits 同时加一个常数，结果不变：

$$\mathrm{softmax}(z)=\mathrm{softmax}(z+c)$$

例如：

softmax([3, 2, 0]) == softmax([103, 102, 100])

实际计算时会用这个性质做稳定化：

$$\mathrm{softmax}(z_i)=\frac{e^{z_i-\max(z)}}{\sum_j e^{z_j-\max(z)}}$$

这样可以避免 $e^{z_i}$ 太大导致数值溢出。

5.3 温度控制“尖锐度”

温度 $T$	效果	生成文本中的表现
$T \to 0$	分布非常尖锐	更像 greedy，保守、重复风险更高
$T = 1$	原始分布	常规随机采样
$T > 1$	分布更平	更发散、多样，但错误风险上升
$T \to \infty$	接近均匀分布	几乎乱选

token	T=0.5	T=1.0	T=2.0
A	█████████████████ 0.86	█████████████ 0.64	█████████ 0.46
B	██ 0.12	█████ 0.24	██████ 0.28
C	█ 0.02	██ 0.09	███ 0.17
D	0.00	█ 0.03	██ 0.10

5.4 softmax 和 sigmoid 的区别

Softmax 和 sigmoid 都能把分数压到概率语义附近，但用法不同。

对比	softmax	sigmoid
输入	一组互相竞争的分数	一个分数，或多个互不排斥的分数
输出约束	所有项相加为 1	每一项独立在 0 到 1 之间
适合场景	多分类、下一个 token、attention 权重	二分类、多标签判断
判断方式	“这些选项里选一个”	“这个标签是否成立”

例子：

任务	更常见选择	原因
下一个 token 是 猫/狗/车 哪一个	softmax	token 之间互相竞争
一张图是否同时有 猫 和 草地	sigmoid	多个标签可以同时成立

---

6. Softmax 在 LLM 中的第一处核心应用：输出层

LLM 的基本任务是：给定前文，预测下一个 token。

flowchart LR
    A[输入上下文 tokens] --> B[Transformer layers]
    B --> C[最后一个位置的 hidden state]
    C --> D[线性层: hidden -> vocab logits]
    D --> E[softmax]
    E --> F[下一个 token 概率分布]
    F --> G[采样或选择 token]

6.1 从 hidden state 到词表概率

假设词表大小是 $V$，模型最后得到一个隐藏向量 $h$。输出层一般写成：

$$z = hW^T + b$$

其中 $z$ 是长度为 $V$ 的 logits。$z_i$ 表示：在当前上下文下，第 $i$ 个 token 作为下一个 token 的未归一化分数。
$W$ 可以看作“每个 token 的输出方向表”。hidden state $h$ 和某个 token 的输出方向越匹配，这个 token 的 logit 越高。

随后：

$$P(\mathrm{token}=i|\mathrm{context})=\frac{e^{z_i}}{\sum_{j=1}^{V}e^{z_j}}$$

6.2 训练时：softmax + cross entropy

训练数据中，真实下一个 token 是 $y$。模型希望提高 $P(y)$。常用损失：

$$\mathcal{L}=-\log P(y)$$

也就是负对数似然。概率越低，损失越大：

模型给真实 token 的概率	损失 $-\log P(y)$	含义
0.90	0.105	很好
0.50	0.693	一般
0.10	2.303	差
0.01	4.605	很差

flowchart LR
    A[模型 logits] --> B[softmax]
    B --> C[P(token)]
    D[真实 token y] --> E[cross entropy]
    C --> E
    E --> F[反向传播]

6.2.1 训练细节：为什么损失函数直接吃 logits

在 PyTorch、TensorFlow 这类框架里，交叉熵损失通常接收 raw logits，而不是手动 softmax 后的概率。

原因有两个：

原因	解释
数值更稳定	框架会把 log_softmax 和负对数似然融合计算，避免 exp、log 带来的溢出或精度损失
梯度更简洁	对“softmax + cross entropy”组合，梯度就是 预测概率 - 真实标签

对 one-hot 标签来说，可以直接记成：

梯度 = p - y

其中：

符号	含义
$p$	softmax 后的预测概率
$y$	真实标签的 one-hot 分布

如果真实 token 是 猫，但模型只给了 0.20 概率，那么 猫 对应的梯度会推动它的 logit 上升；其他 token 的概率越高，被压低的力度也越大。

6.2.2 label smoothing 和 perplexity

训练语言模型时，经常还会遇到两个和 softmax 输出层相关的词：

名词	作用	说明
label smoothing	把 100% 正确答案改成“多数给正确 token，少量分给其他 token”	避免模型过度自信
perplexity	$\exp(\mathrm{cross\ entropy})$	模型平均有多“困惑”

一个简化例子：

标签形式	猫	狗	车
one-hot	1.00	0.00	0.00
label smoothing	0.90	0.05	0.05

perplexity 越低，模型越能预测真实文本。但它不能单独代表生成质量；采样策略、任务目标和偏好对齐也会影响最终输出。

6.3 推理时：softmax 后如何选 token

解码方式	是否用 softmax 概率	做法	风格
greedy	间接使用	选最大 logit/probability	稳定但可能单调
temperature sampling	使用	logits 除以 $T$ 后 softmax，再采样	控制随机性
top-k	使用	只保留概率最高的 k 个 token	去掉长尾
top-p / nucleus	使用	保留累计概率达到 p 的最小 token 集	动态控制候选集
beam search	使用	保留若干高概率路径	机器翻译常见，开放生成中可能乏味

flowchart LR
    A[logits] --> B[可选: repetition penalty / logit bias]
    B --> C[temperature 缩放]
    C --> D[softmax 概率]
    D --> E[top-k / top-p 截断]
    E --> F[重新归一化]
    F --> G[采样 token]

---

7. Softmax 在 LLM 中的第二处核心应用：注意力机制

Transformer 的注意力公式：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$

实际实现里还会加 mask 和 attention bias：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M + B\right)V$$

这里 softmax 处理的是 attention scores，不是输出层的 vocab logits。二者都是“softmax 前的分数”，但对象不同：输出层 logits 对词表 token打分，attention scores 对上下文位置打分。

符号	含义
$Q$	query：当前位置“想找什么”
$K$	key：每个上下文位置“提供什么索引”
$V$	value：每个上下文位置携带的信息
$QK^T$	query 和 key 的相似度
$\sqrt{d_k}$	缩放项，防止点积过大让 softmax 饱和
$M$	mask，例如 causal mask
$B$	位置偏置，例如相对位置偏置、ALiBi 等

7.1 注意力的软选择过程

flowchart LR
    A[Q 和 K] --> B[相似度分数]
    B --> C[scale + mask]
    C --> D[softmax]
    D --> E[注意力权重]
    E --> F[加权 V]

7.2 为什么 attention 需要 softmax

需求	softmax 的作用
在多个上下文 token 中分配权重	把相似度变成概率权重
更关注相关 token	高相似度经过指数后权重大
保持可导	能端到端训练
支持 mask	被禁止的位置加 $-\infty$ 后概率约为 0
支持位置偏置	在 softmax 前加 bias 就能改变注意力分布

7.3 causal mask 怎么起作用

自回归 LLM 生成第 $t$ 个 token 时不能看到未来 token。实现上会把未来位置的 attention logit 设成 $-\infty$：

softmax([2.1, 1.3, -inf, -inf]) = [0.69, 0.31, 0.00, 0.00]

flowchart LR
    A[注意力分数] --> B{位置是否合法?}
    B -- 历史和当前 --> C[保留原分数]
    B -- 未来位置 --> D[设为 -inf]
    C --> E[softmax]
    D --> E
    E --> F[未来 token 权重为 0]

---

8. 两个 softmax 的区别：输出层 vs 注意力层

维度	输出层 softmax	注意力 softmax
作用对象	词表中的所有 token	当前序列中的上下文位置
输入	vocab logits	attention scores
输出	下一个 token 的概率	每个上下文位置的权重
规模瓶颈	词表大小 $V$，常见数万到数十万	序列长度 $L$，attention 矩阵 $L \times L$
训练损失直接使用	是，cross entropy	不直接作为最终损失，但影响 hidden states
推理阶段是否可采样	是，从 token 分布采样	否，用于加权聚合 value

位置	流向
输出层 softmax	hidden state -> vocab logits -> token 概率 -> 采样下一个 token
注意力 softmax	QK 相似度 -> 位置权重 -> 加权求和 V -> 更新 hidden state

8.1 扩展应用：MoE 路由里的 softmax

如果模型采用 Mixture-of-Experts（MoE）结构，softmax 还可能出现在 router 里。router 不选下一个 token，也不选上下文位置；它给当前 token 选择要经过哪些 expert。

flowchart LR
    A[token hidden state] --> B[router logits]
    B --> C[softmax]
    C --> D[top-k experts]
    D --> E[加权合并 expert 输出]

位置	softmax 对象	输出含义
输出层	词表 token	下一个 token 的概率
注意力层	上下文位置	每个位置的信息权重
MoE router	expert 列表	每个 expert 的路由权重

MoE 路由后面常接 top-k 选择，例如每个 token 只激活 1 个或 2 个 expert。这样可以让模型拥有更多参数，同时让每次推理只使用其中一小部分。普通 dense Transformer 没有这一层 router softmax。

---

9. 工程问题：为什么 softmax 不是“算一下”这么简单

9.1 数值稳定性

直接算 $e^{z_i}$ 可能溢出。例如 $e^{1000}$ 超出常规浮点范围。

稳定写法：

$$\mathrm{softmax}(z_i)=\frac{e^{z_i-m}}{\sum_j e^{z_j-m}}, \quad m=\max_j z_j$$

训练中一般不会显式先 softmax 再 log，而是用 fused 的 log_softmax + negative log likelihood 或 cross_entropy，减少溢出和精度损失。

这背后常用到 log-sum-exp 技巧：

$$\log \sum_j e^{z_j}=m+\log \sum_j e^{z_j-m}, \quad m=\max_j z_j$$

它和“先减最大值再 softmax”是同一个稳定性思路：数学结果不变，但指数项不会过大。深度学习框架里的 log_softmax、cross_entropy 一般会在内部做类似处理。

9.2 大词表成本

输出层 full softmax 的复杂度与词表大小 $V$ 线性相关：

每个位置都要给 V 个 token 打分并归一化

如果：

参数	示例
batch size	8
sequence length	2048
vocab size	100,000

那么仅 logits 就有：

8 * 2048 * 100000 = 1,638,400,000 个数

这会占用大量显存和计算。

常用缓解思路：

方法	核心思想	常见场景
hierarchical softmax	用树结构把 $O(V)$ 降到约 $O(\log V)$	早期神经语言模型、word2vec
sampled softmax / NCE / negative sampling	训练时只比较少量负样本	大词表表示学习
adaptive softmax	高频词精细算，低频词分簇算	大词表语言模型
vocabulary pruning	推理时只考虑部分候选	特定任务或约束解码

9.3 注意力的 $L^2$ 成本

注意力 softmax 的输入是 $L \times L$ 的分数矩阵。序列越长，成本越高：

上下文长度 $L$	attention score 数量 $L^2$
2,048	4,194,304
8,192	67,108,864
32,768	1,073,741,824
128,000	16,384,000,000

FlashAttention 不改变 softmax 的数学结果。它通过分块和 online softmax，避免把完整 attention 矩阵长时间写到显存里，从而减少 HBM 读写。

flowchart LR
    A[标准 attention] --> B[生成完整 LxL 分数矩阵]
    B --> C[写入/读取显存]
    C --> D[softmax]
    D --> E[乘 V]

    F[FlashAttention] --> G[分块读 Q/K/V]
    G --> H[块内 online softmax]
    H --> I[累积输出]
    I --> J[减少显存 IO]

9.4 softmax 饱和与梯度问题

如果某个 logit 远大于其他 logit，softmax 会非常接近 one-hot：

softmax([20, 0, 0]) ≈ [1.0, 0.0, 0.0]

这会导致：

现象	影响
分布过尖	模型过度自信
梯度变小	学习信号变弱
生成更确定	多样性下降
attention 过集中	模型可能忽视长距离信息

Transformer 里的 $1/\sqrt{d_k}$ 缩放，就是为了缓解点积随维度增大而使 softmax 过尖的问题。

---

10. Softmax 与 LLM 解码参数

不少 LLM API 的生成参数，都是在 softmax 前后改概率分布。

参数/技术	发生位置	数学含义	效果
temperature	softmax 前	$z_i \leftarrow z_i/T$	控制随机性
top-k	softmax 后或截断 logits 后	只保留最高 k 个候选	去掉低概率长尾
top-p / nucleus	softmax 后	保留累计概率达到 p 的最小集合	动态候选集
repetition penalty	softmax 前	降低已出现 token 的 logit	减少重复
logit bias	softmax 前	手动提高/降低某些 token 的 logit	控制词汇偏好
greedy	softmax 可省略	取最大 logit	确定性输出

10.1 top-k 和 top-p 的区别

flowchart LR
    A[概率分布] --> B[top-k: 固定 k 个]
    A --> C[top-p: 累计概率 p]
    B --> D[候选数固定]
    C --> E[候选数动态]

示例：

token	概率	top-k=3	top-p=0.9
A	0.50	保留	保留
B	0.20	保留	保留
C	0.12	保留	保留
D	0.08	丢弃	保留
E	0.04	丢弃	丢弃
其他	0.06	丢弃	丢弃

top-p=0.9 会保留 A+B+C+D，因为累计概率 0.90。

---

11. Softmax 的表达能力问题：softmax bottleneck

在语言模型中，输出层常可写成：

$$z = hW^T$$

其中：

符号	含义
$h$	当前上下文的 hidden state
$W$	token embedding/output matrix
$z$	vocab logits

这意味着 logits 会受到隐藏维度和输出矩阵秩的限制。Yang 等人在 “Breaking the Softmax Bottleneck” 中指出，许多基于 softmax 的语言模型存在表达能力瓶颈：自然语言的条件分布很复杂，而标准 softmax 输出结构可能限制模型能表示的分布族。

相关思路：

方向	代表思想	作用
Mixture of Softmaxes	多个 softmax 混合	用多个低秩分布组合出更高秩分布
更强 hidden state	增大模型深度/宽度	提高上下文表示能力
untied/tied embeddings 权衡	输入输出 embedding 是否共享	参数效率 vs 表达自由度
alternative output distributions	sparsemax、entmax 等	让概率分布可稀疏或形状不同

现代 LLM 仍大量使用标准 softmax。softmax bottleneck 提供的是理解输出层限制的理论视角，不表示 softmax 已经被主流 LLM 淘汰。

---

12. 替代方案与变体

主流路线：

方法	用途	主要收益	主要代价
softmax	标准输出层/注意力	稳定、简单、生态成熟	大词表和长序列成本高
FlashAttention	exact attention 实现	更快、更省显存	依赖内核和硬件实现
hierarchical softmax	大词表近似	降低 $O(V)$ 成本	树结构会影响效果
sampled softmax / NCE	训练近似	减少负类计算	有偏差/方差权衡

研究或特定场景：

方法	特点	适用场景
sparsemax	可产生 0 概率	希望分布稀疏
entmax	softmax 与 sparsemax 折中	希望可控稀疏
adaptive softmax	高频词细算、低频词分簇	大词表语言模型
linear attention	不一定使用 softmax	追求低于 $L^2$ 的长上下文复杂度
Gumbel-Softmax	可微近似离散采样	训练含离散选择的模型

---

13. Softmax 在 LLM 中的位置

这里用总览表和短流程图放在一起看，避免一张大图在 Obsidian 里拉得太长。

场景	softmax 前的分数	softmax 后的结果	直接用途
训练输出层	vocab logits	token 概率或 log 概率	计算 cross entropy，更新参数
每层 attention	attention scores	上下文位置权重	加权聚合 value，更新 hidden state
推理解码	处理后的 vocab logits	候选 token 概率	greedy、temperature、top-k、top-p 选 token
MoE router	expert logits	expert 路由权重	选择并加权少量 expert

13.1 训练输出层

flowchart LR
    A[文本 tokens] --> B[Transformer]
    B --> C[vocab logits]
    C --> D[log_softmax]
    D --> E[cross entropy]

13.2 每层 attention

flowchart LR
    A[QK^T] --> B[scale/mask/bias]
    B --> C[softmax]
    C --> D[attention 权重]
    D --> E[加权 V]

13.3 推理解码

flowchart LR
    A[最后位置 logits] --> B[temperature/penalty]
    B --> C[softmax]
    C --> D[top-k/top-p]
    D --> E[下一个 token]

13.4 MoE 路由

flowchart LR
    A[hidden state] --> B[router logits]
    B --> C[softmax]
    C --> D[top-k experts]
    D --> E[expert 输出]

---

14. 常见误解

误解	更准确的说法
softmax 会“理解”概率	softmax 只做数值归一化，概率质量来自训练目标和数据
最大 logit 的 token 一定会被选中	采样时只是概率最大，不一定被抽中
temperature 是在 softmax 后调概率	通常是 softmax 前缩放 logits
attention softmax 输出的是词概率	attention softmax 输出的是上下文位置权重
mask 是 softmax 后把概率改成 0	常见实现是在 softmax 前把非法 logit 加上 $-\infty$
FlashAttention 是近似 attention	FlashAttention 原始核心是 exact attention，只改变计算组织方式
softmax 概率就是模型“真实置信度”	神经网络概率可能未校准，不能直接等同真实可信度
训练时应该先手动 softmax 再算 cross entropy	多数框架的交叉熵损失应直接输入 raw logits
MoE router 的 softmax 是 token 概率	router softmax 输出的是 expert 权重，不是词表概率

---

15. 面向工程实践的检查清单

15.1 写 softmax 代码时

检查项	建议
是否减去 max	必须，避免 overflow
是否在训练中手写 softmax + log	尽量用框架的 fused cross entropy
mask 是否在 softmax 前加	推荐 pre-softmax mask
mask 值是否足够小	常用 -inf 或 dtype 可承受的极小值
fp16/bf16 下是否稳定	注意溢出、下溢和累加精度
损失函数是否期待 logits	先看框架文档，避免重复 softmax

15.2 调 LLM 生成时

想要的效果	可调参数
更确定、更保守	降低 temperature，降低 top-p
更有创意	提高 temperature，适度提高 top-p
减少胡言乱语	降低 temperature，限制 top-p/top-k
减少重复	使用 repetition penalty 或 presence/frequency penalty
禁止/鼓励特定词	使用 logit bias

---

16. 收束

Softmax 把“相对分数”变成“可训练、可采样、可解释的概率分布”。
它的历史大致经历了三次迁移：

flowchart LR
    A[统计物理: 状态能量 -> 概率] --> B[选择理论: 偏好强度 -> 选择概率]
    B --> C[神经网络: logits -> 分类概率]
    C --> D[语言模型: 上下文 -> token 概率]
    C --> E[Transformer: token 相似度 -> 注意力权重]

在 LLM 中，softmax 同时参与两件事：决定下一个 token 的概率，以及决定上下文信息怎么流动。在 MoE 模型里，它还可能参与 expert 路由。公式本身不长，但围绕它会碰到大词表训练、长上下文注意力、采样策略、数值稳定、概率校准和表达能力瓶颈等问题。

---

17. 参考资料

17.1 历史与基础

• Boltzmann/Gibbs 分布：Boltzmann distribution
• Luce 选择公理：Scholarpedia: Luce’s choice axiom
• Luce 1959 书目：Individual Choice Behavior
• Bridle 1989：Training Stochastic Model Recognition Algorithms
• Bridle 1990：Probabilistic Interpretation of Feedforward Classification Network Outputs

17.2 语言模型与 LLM

• 神经概率语言模型：Bengio et al., A Neural Probabilistic Language Model
• hierarchical softmax：Morin & Bengio, Hierarchical Probabilistic Neural Network Language Model
• word2vec 与负采样：Mikolov et al., Distributed Representations of Words and Phrases
• adaptive softmax：Grave et al., Efficient Softmax Approximation for GPUs
• softmax bottleneck：Yang et al., Breaking the Softmax Bottleneck

17.3 Transformer、注意力与生成

• Transformer 注意力：Vaswani et al., Attention Is All You Need
• nucleus sampling / top-p：Holtzman et al., The Curious Case of Neural Text Degeneration
• FlashAttention：Dao et al., FlashAttention
• FlashAttention-2：Dao, FlashAttention-2

17.4 训练技巧、MoE 与变体

• label smoothing：Szegedy et al., Rethinking the Inception Architecture for Computer Vision
• Gumbel-Softmax：Jang et al., Categorical Reparameterization with Gumbel-Softmax
• Switch Transformer：Fedus et al., Switch Transformers
• Mixtral MoE：Jiang et al., Mixtral of Experts