大模型原理之 Softmax

摘要：Softmax 是将任意实数 logits 转化为概率分布的归一化指数函数，是大模型中出场频率最高的数学函数。本文以"为什么→是什么→在哪→怎么稳→怎么训"为主线：它被概率分布的三项硬约束（各项为正、总和为一、保序）加上光滑可导要求逼成唯一自然解；在 Transformer 中出现于输出层（决定"生成什么"）与注意力层（决定"关注什么"），配套 d k \sqrt{d_k} dk 缩放与因果掩码；工程上利用平移不变性减 max 防溢出，online softmax 支撑 Flash Attention 将显存从 O ( n 2 ) O(n^2) O(n2) 降到 O ( n ) O(n) O(n)；训练时与交叉熵复合，梯度简化为"预测减真实"。最后按"控尖锐度、稀疏化、省算力"梳理主要变体，指出其共性是以牺牲标准 softmax 的某条性质换取特定收益。

关键词：Softmax；大语言模型；注意力机制；交叉熵；温度参数；Flash Attention

一、为什么需要 Softmax

模型最后一层输出的 logits 是任意实数（可正可负、无界），但下游需要的是一个概率分布。这座"打分 → 概率"的桥必须满足四个条件：

每项 > 0 ：指数 e z e^z ez 恒为正
总和 = 1：归一化，除以总和
保持相对强弱：指数天然单调
可以训练：指数处处光滑可导

其他候选全被条件逼出局：线性归一化 除完还有负数；argmax 没有梯度；sigmoid 各项独立、加总不为 1。而指数函数有一个独特性质------差分变比值 ： z i − z j z_i - z_j zi−zj 直接变成概率比 e z i / e z j e^{z_i}/e^{z_j} ezi/ezj，logit 的加法对应概率的乘法，这是其他函数不具备的。由此还自然导出平移不变性 （ softmax ( z + c ) = softmax ( z ) \text{softmax}(z+c) = \text{softmax}(z) softmax(z+c)=softmax(z)），为第四节的数值稳定性埋下伏笔。Softmax 不是被挑选出来的，而是被这些条件逼出来的唯一自然解 ------顺带还有指数放大差距的福利，使它成为 argmax 的可导软化版 $1$ 。

二、公式与计算

softmax ( z ) i = e z i ∑ j = 1 K e z j \text{softmax}(z)i = \frac{e^{z_i}}{\sum{j=1}^{K} e^{z_j}} softmax(z)i=∑j=1Kezjezi

以 logits [2.0, 1.0, 0.1, 3.0] 为例，三步算完：

步骤	操作	结果
① 取指数	e z i e^{z_i} ezi	`[7.39, 2.72, 1.11, 20.09]`
② 求和	归一化分母	`31.31`
③ 相除	逐项除以总和	`[0.236, 0.087, 0.035, 0.642]`

输出全正、总和为 1，头部概率被显著放大------这就是"软化的 argmax"。

三、在大模型中的两个位置

先看全景，softmax 在 Transformer 中一内一外只出现两次 $2$ ：
#mermaid-svg-77McbRJz3zIz19Ha{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-77McbRJz3zIz19Ha .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-77McbRJz3zIz19Ha .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-77McbRJz3zIz19Ha .error-icon{fill:#552222;}#mermaid-svg-77McbRJz3zIz19Ha .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-77McbRJz3zIz19Ha .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-77McbRJz3zIz19Ha .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-77McbRJz3zIz19Ha .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-77McbRJz3zIz19Ha .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-77McbRJz3zIz19Ha .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-77McbRJz3zIz19Ha .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-77McbRJz3zIz19Ha .marker{fill:#333333;stroke:#333333;}#mermaid-svg-77McbRJz3zIz19Ha .marker.cross{stroke:#333333;}#mermaid-svg-77McbRJz3zIz19Ha svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-77McbRJz3zIz19Ha p{margin:0;}#mermaid-svg-77McbRJz3zIz19Ha .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-77McbRJz3zIz19Ha .cluster-label text{fill:#333;}#mermaid-svg-77McbRJz3zIz19Ha .cluster-label span{color:#333;}#mermaid-svg-77McbRJz3zIz19Ha .cluster-label span p{background-color:transparent;}#mermaid-svg-77McbRJz3zIz19Ha .label text,#mermaid-svg-77McbRJz3zIz19Ha span{fill:#333;color:#333;}#mermaid-svg-77McbRJz3zIz19Ha .node rect,#mermaid-svg-77McbRJz3zIz19Ha .node circle,#mermaid-svg-77McbRJz3zIz19Ha .node ellipse,#mermaid-svg-77McbRJz3zIz19Ha .node polygon,#mermaid-svg-77McbRJz3zIz19Ha .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-77McbRJz3zIz19Ha .rough-node .label text,#mermaid-svg-77McbRJz3zIz19Ha .node .label text,#mermaid-svg-77McbRJz3zIz19Ha .image-shape .label,#mermaid-svg-77McbRJz3zIz19Ha .icon-shape .label{text-anchor:middle;}#mermaid-svg-77McbRJz3zIz19Ha .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-77McbRJz3zIz19Ha .rough-node .label,#mermaid-svg-77McbRJz3zIz19Ha .node .label,#mermaid-svg-77McbRJz3zIz19Ha .image-shape .label,#mermaid-svg-77McbRJz3zIz19Ha .icon-shape .label{text-align:center;}#mermaid-svg-77McbRJz3zIz19Ha .node.clickable{cursor:pointer;}#mermaid-svg-77McbRJz3zIz19Ha .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-77McbRJz3zIz19Ha .arrowheadPath{fill:#333333;}#mermaid-svg-77McbRJz3zIz19Ha .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-77McbRJz3zIz19Ha .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-77McbRJz3zIz19Ha .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-77McbRJz3zIz19Ha .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-77McbRJz3zIz19Ha .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-77McbRJz3zIz19Ha .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-77McbRJz3zIz19Ha .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-77McbRJz3zIz19Ha .cluster text{fill:#333;}#mermaid-svg-77McbRJz3zIz19Ha .cluster span{color:#333;}#mermaid-svg-77McbRJz3zIz19Ha div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-77McbRJz3zIz19Ha .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-77McbRJz3zIz19Ha rect.text{fill:none;stroke-width:0;}#mermaid-svg-77McbRJz3zIz19Ha .icon-shape,#mermaid-svg-77McbRJz3zIz19Ha .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-77McbRJz3zIz19Ha .icon-shape p,#mermaid-svg-77McbRJz3zIz19Ha .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-77McbRJz3zIz19Ha .icon-shape .label rect,#mermaid-svg-77McbRJz3zIz19Ha .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-77McbRJz3zIz19Ha .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-77McbRJz3zIz19Ha .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-77McbRJz3zIz19Ha :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 堆叠 × N 层
tokens
Embedding
注意力层

softmax(QKᵀ/√d_k)·V

位置②：决定关注什么
前馈层 (FFN)
输出层 (LM Head)

logits → softmax → 概率分布

位置①：决定生成什么
采样 (温度 / top-k / top-p)

图 1 Softmax 在 Transformer 中的两个位置

3.1 输出层决定"生成什么"

自回归生成时，每一步对整个词表做一次 softmax，得到下一个 token 的概率分布，再从中采样。所有解码参数都作用在它的输出之后：

温度 T：softmax(z/T)。T<1 分布尖锐（保守确定），T>1 平坦（随机有创造性）；T→0 退化为 argmax，T→∞ 退化为均匀分布------一个旋钮控制全部谱系。
Top-k / Top-p：只保留概率最高的 k 个、或累计概率达 p 的最小集合，再归一化采样，防止选中低质量词。温度 + top-p 是生产环境最常见组合。

3.2 注意力层决定"关注什么"

Attention ( Q , K , V ) = softmax ⁣ ( Q K ⊤ d k ) V \text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}}\right)V Attention(Q,K,V)=softmax(dk QK⊤)V

softmax 把 query 与 key 的相似度分数归一化为和为 1 的权重，再对 V 加权求和。两个配套设计都源于 softmax 的特性：

除以 d k \sqrt{d_k} dk ：维度大时点积分差过大，会把 softmax 推入饱和区（趋近 one-hot、梯度消失）；缩放把方差拉回 1，保持梯度健康 $2$ 。
Causal mask ：屏蔽位置加 − ∞ -\infty −∞，softmax 后概率严格为 0------"看不到未来"的数学表达。

四、数值稳定性

logit 较大时 e z i e^{z_i} ezi 会溢出（float32 上限约 e 88 e^{88} e88）。利用 softmax 的平移不变性 （整体加减常数结果不变），实现时先减去最大值 $3$ ：

softmax ( z ) i = e z i − max ⁡ ( z ) ∑ j e z j − max ⁡ ( z ) \text{softmax}(z)_i = \frac{e^{z_i - \max(z)}}{\sum_j e^{z_j - \max(z)}} softmax(z)i=∑jezj−max(z)ezi−max(z)

减 max 后指数最大仅为 e 0 = 1 e^0=1 e0=1，彻底避免溢出。在此基础上，online softmax 支持分块流式计算：
#mermaid-svg-uIxToOCJODvB9RGU{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;fill:#333;}@keyframes edge-animation-frame{from{stroke-dashoffset:0;}}@keyframes dash{to{stroke-dashoffset:0;}}#mermaid-svg-uIxToOCJODvB9RGU .edge-animation-slow{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 50s linear infinite;stroke-linecap:round;}#mermaid-svg-uIxToOCJODvB9RGU .edge-animation-fast{stroke-dasharray:9,5!important;stroke-dashoffset:900;animation:dash 20s linear infinite;stroke-linecap:round;}#mermaid-svg-uIxToOCJODvB9RGU .error-icon{fill:#552222;}#mermaid-svg-uIxToOCJODvB9RGU .error-text{fill:#552222;stroke:#552222;}#mermaid-svg-uIxToOCJODvB9RGU .edge-thickness-normal{stroke-width:1px;}#mermaid-svg-uIxToOCJODvB9RGU .edge-thickness-thick{stroke-width:3.5px;}#mermaid-svg-uIxToOCJODvB9RGU .edge-pattern-solid{stroke-dasharray:0;}#mermaid-svg-uIxToOCJODvB9RGU .edge-thickness-invisible{stroke-width:0;fill:none;}#mermaid-svg-uIxToOCJODvB9RGU .edge-pattern-dashed{stroke-dasharray:3;}#mermaid-svg-uIxToOCJODvB9RGU .edge-pattern-dotted{stroke-dasharray:2;}#mermaid-svg-uIxToOCJODvB9RGU .marker{fill:#333333;stroke:#333333;}#mermaid-svg-uIxToOCJODvB9RGU .marker.cross{stroke:#333333;}#mermaid-svg-uIxToOCJODvB9RGU svg{font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:16px;}#mermaid-svg-uIxToOCJODvB9RGU p{margin:0;}#mermaid-svg-uIxToOCJODvB9RGU .label{font-family:"trebuchet ms",verdana,arial,sans-serif;color:#333;}#mermaid-svg-uIxToOCJODvB9RGU .cluster-label text{fill:#333;}#mermaid-svg-uIxToOCJODvB9RGU .cluster-label span{color:#333;}#mermaid-svg-uIxToOCJODvB9RGU .cluster-label span p{background-color:transparent;}#mermaid-svg-uIxToOCJODvB9RGU .label text,#mermaid-svg-uIxToOCJODvB9RGU span{fill:#333;color:#333;}#mermaid-svg-uIxToOCJODvB9RGU .node rect,#mermaid-svg-uIxToOCJODvB9RGU .node circle,#mermaid-svg-uIxToOCJODvB9RGU .node ellipse,#mermaid-svg-uIxToOCJODvB9RGU .node polygon,#mermaid-svg-uIxToOCJODvB9RGU .node path{fill:#ECECFF;stroke:#9370DB;stroke-width:1px;}#mermaid-svg-uIxToOCJODvB9RGU .rough-node .label text,#mermaid-svg-uIxToOCJODvB9RGU .node .label text,#mermaid-svg-uIxToOCJODvB9RGU .image-shape .label,#mermaid-svg-uIxToOCJODvB9RGU .icon-shape .label{text-anchor:middle;}#mermaid-svg-uIxToOCJODvB9RGU .node .katex path{fill:#000;stroke:#000;stroke-width:1px;}#mermaid-svg-uIxToOCJODvB9RGU .rough-node .label,#mermaid-svg-uIxToOCJODvB9RGU .node .label,#mermaid-svg-uIxToOCJODvB9RGU .image-shape .label,#mermaid-svg-uIxToOCJODvB9RGU .icon-shape .label{text-align:center;}#mermaid-svg-uIxToOCJODvB9RGU .node.clickable{cursor:pointer;}#mermaid-svg-uIxToOCJODvB9RGU .root .anchor path{fill:#333333!important;stroke-width:0;stroke:#333333;}#mermaid-svg-uIxToOCJODvB9RGU .arrowheadPath{fill:#333333;}#mermaid-svg-uIxToOCJODvB9RGU .edgePath .path{stroke:#333333;stroke-width:2.0px;}#mermaid-svg-uIxToOCJODvB9RGU .flowchart-link{stroke:#333333;fill:none;}#mermaid-svg-uIxToOCJODvB9RGU .edgeLabel{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-uIxToOCJODvB9RGU .edgeLabel p{background-color:rgba(232,232,232, 0.8);}#mermaid-svg-uIxToOCJODvB9RGU .edgeLabel rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-uIxToOCJODvB9RGU .labelBkg{background-color:rgba(232, 232, 232, 0.5);}#mermaid-svg-uIxToOCJODvB9RGU .cluster rect{fill:#ffffde;stroke:#aaaa33;stroke-width:1px;}#mermaid-svg-uIxToOCJODvB9RGU .cluster text{fill:#333;}#mermaid-svg-uIxToOCJODvB9RGU .cluster span{color:#333;}#mermaid-svg-uIxToOCJODvB9RGU div.mermaidTooltip{position:absolute;text-align:center;max-width:200px;padding:2px;font-family:"trebuchet ms",verdana,arial,sans-serif;font-size:12px;background:hsl(80, 100%, 96.2745098039%);border:1px solid #aaaa33;border-radius:2px;pointer-events:none;z-index:100;}#mermaid-svg-uIxToOCJODvB9RGU .flowchartTitleText{text-anchor:middle;font-size:18px;fill:#333;}#mermaid-svg-uIxToOCJODvB9RGU rect.text{fill:none;stroke-width:0;}#mermaid-svg-uIxToOCJODvB9RGU .icon-shape,#mermaid-svg-uIxToOCJODvB9RGU .image-shape{background-color:rgba(232,232,232, 0.8);text-align:center;}#mermaid-svg-uIxToOCJODvB9RGU .icon-shape p,#mermaid-svg-uIxToOCJODvB9RGU .image-shape p{background-color:rgba(232,232,232, 0.8);padding:2px;}#mermaid-svg-uIxToOCJODvB9RGU .icon-shape .label rect,#mermaid-svg-uIxToOCJODvB9RGU .image-shape .label rect{opacity:0.5;background-color:rgba(232,232,232, 0.8);fill:rgba(232,232,232, 0.8);}#mermaid-svg-uIxToOCJODvB9RGU .label-icon{display:inline-block;height:1em;overflow:visible;vertical-align:-0.125em;}#mermaid-svg-uIxToOCJODvB9RGU .node .label-icon path{fill:currentColor;stroke:revert;stroke-width:revert;}#mermaid-svg-uIxToOCJODvB9RGU :root{--mermaid-font-family:"trebuchet ms",verdana,arial,sans-serif;} 用新 max 重缩放旧累积值
块 1

max=5

分母=20
块 2

max=8

重缩放÷e³
块 3

max=9

重缩放÷e¹
最终结果

图 2 Online softmax 分块流式计算

这是 Flash Attention 的核心技巧------显存从 O ( n 2 ) O(n^2) O(n2) 降到 O ( n ) O(n) O(n)，让长上下文成为可能 $3$ $4$ 。

五、训练与交叉熵的组合

预训练最小化交叉熵 L = − log ⁡ softmax ( z ) y true \mathcal{L} = -\log\,\text{softmax}(z){y{\text{true}}} L=−logsoftmax(z)ytrue。softmax 单独求导较繁琐，但与交叉熵复合后梯度简化为：

∂ L ∂ z = p − y \frac{\partial \mathcal{L}}{\partial z} = p - y ∂z∂L=p−y

即"预测分布减真实分布"，形式简洁且无饱和问题------这是两者成为标配组合的数学原因。常用指标困惑度即交叉熵的指数： PPL = e L \text{PPL} = e^{\mathcal{L}} PPL=eL。

六、主要变体

控尖锐度 ：温度；知识蒸馏用高温软化 teacher 输出，传递"暗知识" $5$ ；
稀疏化 ：Sparsemax 输出可含精确 0 $6$ ；Gumbel-Softmax 实现可微的近似 one-hot 采样 $7$ ；
省算力 ：层次化 softmax（ O ( V ) O(V) O(V) → O ( log ⁡ V ) O(\log V) O(logV)）、负采样、Flash Attention $4$ 。

变体本质都是在牺牲标准 softmax 的某条性质，换取特定收益------反衬出标准版的地位。

总结

Softmax 的本质是把任意实数 logits 转化为概率分布的归一化指数函数------大模型训练时用交叉熵拟合它、推理时从它采样、每层注意力靠它分配权重，是模型"做选择"与"分注意力"的共同出口。

参考文献

$1$ J. S. Bridle. "Probabilistic Interpretation of Feedforward Classification Network Outputs, with Relationships to Statistical Pattern Recognition." Neurocomputing , 1990.

$2$ A. Vaswani et al. "Attention Is All You Need." NeurIPS , 2017.

$3$ M. Milakov, N. Gimelshein. "Online Normalizer Calculation for Softmax." arXiv , 2018.

$4$ T. Dao et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." NeurIPS , 2022.

$5$ G. Hinton, O. Vinyals, J. Dean. "Distilling the Knowledge in a Neural Network." arXiv , 2015.

$6$ A. Martins, R. Astudillo. "From Softmax to Sparsemax: A Sparse Model of Attention and Multi-Label Classification." ICML , 2016.

$7$ E. Jang, S. Gu, B. Poole. "Categorical Reparameterization with Gumbel-Softmax." ICLR, 2017.