1. softmax：为什么是指数，为什么先减最大值

模型最后一层吐出 \(V\) 个数 \(z_1,\dots,z_V\)，叫 logits。它们可以是任意实数，加起来也不等于 1，不能当概率。softmax 的任务是把它们变成一个合法的概率分布。

1.1 先确定形式：分子非负，分母归一

最自然的形式是“每个分数过一个非负函数，再除以总和”：

这样自动满足 \(p_i>0\)、\(\sum_i p_i=1\)。问题是 \(f\) 该取什么。如果希望“两个 token 的相对概率只取决于它们分数之差”，就应当存在某个函数 \(h\)，使得

令 \(b=0\)，可知 \(f(a)=f(0)h(a)\)，所以 \(g(a)=f(a)/f(0)=h(a)\)。再利用比值只依赖差值的条件，就得到 \(g(a+b)=g(a)g(b)\)。下面解这个函数方程。

1. 令 \(g(u)=f(u)/f(0)\)，则 \(g(0)=1\)，且 \(g(a+b)=g(a)g(b)\)。
2. 对整数 \(n\)：\(g(n)=g(1)^n\)（归纳：\(g(n)=g(n-1)g(1)\)）。对负整数用 \(g(-n)g(n)=g(0)=1\)。
3. 对有理数 \(u=m/n\)：\(g(u)^n=g(nu)=g(m)=g(1)^m\)，所以 \(g(m/n)=g(1)^{m/n}\)。
4. 若额外假设 \(g\) 连续（或可测、局部有界等正则条件），由有理数稠密并取极限得 \(g(u)=g(1)^u\)。

所以 \(f(z)=f(0)\,c^z\)，其中 \(c=g(1)>0\)。常数 \(f(0)\) 会在分子分母里约掉。取 \(c=e\) 最方便，因为 \(\frac{d}{dz}e^z=e^z\)，求导时不用额外带 \(\ln c\)。若取别的底 \(c\)，只是把 \(z\) 重新缩放成 \(z\ln c\)，分布完全一样。

1.2 手算一个三分类

取 \(z=(2,1,0)\)。先算指数：\(e^2=7.389\)，\(e^1=2.718\)，\(e^0=1\)，总和 \(S=11.107\)。于是

验证：三项相加为 1；分数每降 1，概率大约除以 \(e\)（0.665/0.245≈2.72）。这就是“logit 差 1 ≈ 概率比 e 倍”。

1.3 减去最大值：一个恒等式

直接算 \(e^{z_i}\) 会溢出：float32 的最大值约 \(3.4\times10^{38}\)，而 \(e^{89}\approx4.5\times10^{38}\) 就爆了。解决办法是分子分母同乘 \(e^{-M}\)：

这是恒等变形，不是近似。取 \(M=\max_j z_j\) 后，最大的指数是 \(e^0=1\)，其余都 \(\le 1\)，分母最多 \(V\)，不可能溢出；同时最小的指数可能下溢到 0，但那只代表“概率小到可以忽略”，不影响结果。

1.4 导数：softmax 不是逐元素的

令 \(S=\sum_k e^{z_k}\)，则 \(p_i=e^{z_i}/S\)。对 \(z_j\) 求偏导，要分 \(i=j\) 和 \(i\ne j\)：

1. 当 \(i=j\)：分子导数是 \(e^{z_i}\)，分母导数也是 \(e^{z_j}=e^{z_i}\)。由商法则 \[ \frac{\partial p_i}{\partial z_i}=\frac{e^{z_i}S-e^{z_i}e^{z_i}}{S^2}=p_i(1-p_i). \]
2. 当 \(i\ne j\)：分子对 \(z_j\) 的导数为 0，分母导数为 \(e^{z_j}\)，所以 \[ \frac{\partial p_i}{\partial z_j}=\frac{0-e^{z_i}e^{z_j}}{S^2}=-p_ip_j . \]

合并成一行（\(\delta_{ij}=1\) 当 \(i=j\)，否则 0）：

这个矩阵叫 Jacobian。它有一个重要性质：每一行加起来为 0，因为 \(\sum_i p_i(\delta_{ij}-p_j)=p_j-p_j\sum_i p_i=0\)。直觉是“把所有 logit 同时加一个常数，概率不变”，所以导数在“全 1 方向”上必须为 0，和 1.3 的平移不变性一致。

1.5 log-sum-exp：把 softmax 写在对数域

语言模型经常要“某个 token 的对数概率”。直接算 \(\ln p_i=z_i-\ln\sum_j e^{z_j}\)，但 \(\sum e^{z_j}\) 会溢出。用同样的平移技巧：

于是

这个量永远 \(\le 0\)（因为 \(p_i\le1\)），数值上也不会溢出。仓库里的 shiftAndLogSum 就是在算 \(z_i-M\) 和 \(\ln\sum e^{z_i-M}\)，LogSoftmax32 再把两者相减。