Softmax和Softmax-Loss函数及梯度计算

1. 结合Logistic Regression 分析 Softmax

训练集由 $m$ 个已标记的样本构成：$\{(x^{(1)}, y^{(1)}), ... , (x^{(m)}, y^{(m)})\}$，其中输入的第 $i$ 个样例 $x^{(i)} \in \Re^{n+1}$，即特征向量 $x$ 的维度为 $n + 1$ ，其中 $x_0 = 1$ 对应截距项。由于 logistic 回归是针对二分类问题的，因此类标记 $y^{(i)} \in \{0,1\}$。假设函数(hypothesis function) 如下：
$h_\theta(x) = \frac{1}{1+e^{(-\theta^Tx)}}$
我们将训练模型参数 $\theta$，使其能够最小化代价函数 ：
$J(\theta) = -\frac{1}{m} \left[ \sum_{i=1}^m y^{(i)} \log h_\theta(x^{(i)}) + (1-y^{(i)}) \log (1-h_\theta(x^{(i)})) \right]$
而在 Softmax回归中，我们解决的是多分类问题（相对于 logistic 回归解决的二分类问题），类标 $\textstyle y$ 可以取 $\textstyle k$ 个不同的值（而不是 2 个）。因此，对于训练集 $\{ (x^{(1)}, y^{(1)}), \ldots, (x^{(m)}, y^{(m)}) \}$，我们有 $y^{(i)} \in \{1, 2, \ldots, k\}$。（注意此处的类别下标从 1 开始，而不是 0）。例如，在 MNIST 数字识别任务中，我们有 $\textstyle k=10$ 个不同的类别。

对于输入的每一个样例 $x^{(i)} \in \Re^{n+1}$ ，输出其属于每一类别 $j$ 的概率值 $\textstyle p(y=j | x)$，即输出一个 $k$ 维向量（向量元素的和为1）来表示这 $k$ 个估计的概率值。即 Softmax 的假设函数为 $h_{\theta}(x^{(i)})$
$h_\theta(x^{(i)}) = \begin{bmatrix} p(y^{(i)} = 1 | x^{(i)}; \theta) \\ p(y^{(i)} = 2 | x^{(i)}; \theta) \\ \vdots \\ p(y^{(i)} = k | x^{(i)}; \theta) \end{bmatrix} =\frac{1}{ \sum_{j=1}^{k}{e^{ \theta_j^T x^{(i)} }} } \begin{bmatrix} e^{ \theta_1^T x^{(i)} } \\ e^{ \theta_2^T x^{(i)} } \\ \vdots \\ e^{ \theta_k^T x^{(i)} } \\ \end{bmatrix}$
其中 $\theta_1, \theta_2, \ldots, \theta_k \in \Re^{n+1}$ 是模型的参数。请注意 $\frac{1}{ \sum_{j=1}^{k}{e^{ \theta_j^T x^{(i)} }} }$ 这一项对概率分布进行归一化，使得所有概率之和为 1

其代价函数为：
$J(\theta) = - \frac{1}{m} \left[ \sum_{i=1}^{m} \sum_{j=1}^{k} 1\left\{y^{(i)} = j\right\} \log \frac{e^{\theta_j^T x^{(i)}}}{\sum_{l=1}^k e^{ \theta_l^T x^{(i)} }}\right]$
其中 $1\{\cdot\}$ 为指示函数：
$1\{\cdot\} = \begin{cases} 1\{表达式值为真\} = 1 \\ 1\{表达式值为假\} = 0 \\ \end{cases}$
梯度下降公式：
$\nabla_{\theta_j} J(\theta) = - \frac{1}{m} \sum_{i=1}^{m}{ \left[ x^{(i)} \left( 1\{ y^{(i)} = j\} - p(y^{(i)} = j | x^{(i)}; \theta) \right) \right] }$

2. Softmax

结合上面理解此时的 Softmax函数：
$h_{\theta}(x) = \sigma(\theta^T x) = \sigma(z) = \left({\color{Red}{\sigma{1}(z)}}, \ldots, \sigma_{m}(z)\right)$
此时的 $x$ 为一个样例输入 $\in \Re^{n+1}$，等价与上面的一个 $x^{(i)}$，其中共有 $m$ 个类别，$z_i$ 表示输入样例 $x$ 是第 $i$ 个类别的线性预测结果，等价于 $z_i = \theta^T_i x$

Softmax 函数 $\sigma(z)=\left({\color{Red}{\sigma_{1}(z)}}, \ldots, \sigma_{m}(z)\right)$ 定义如下：
${\color{Green}{o_i}} = {\color{Red}{\sigma_{i}(z)}} =\frac{\exp \left(z_{i}\right)}{\sum_{j=1}^{m} \exp \left(z_{j}\right)}, \quad i=1, \ldots, m \\ {\color{Green}{【观察到的数据 \ x \ (或z)\ 属于类别\ i\ 的概率，或者称作似然 (Likelihood)】}}$
它在 Logistic Regression 里其到的作用是讲线性预测值转化为类别概率：$m$ 代表类别数，假设 $z_i = w_i^Tx + b_i$ 是第 $i$ 个类别的线性预测结果，带入 $Softmax$ 的结果其实就是先对每一个 $z_i$ 取 exponential 变成非负，然后除以所有项之和进行归一化，现在每个 $o_{i}=\sigma_{i}(z)$ 就可以解释成：观察到的数据 $x$ 属于类别 $i$ 的概率，或者称作似然 (Likelihood)。

然后 Logistic Regression 的目标函数是根据最大似然原则来建立的，假设数据 $x$ 所对应的类别为 $y$，则根据我们刚才的计算最大似然就是要最大化 $o_y$ 的值 (通常是使用 negative log-likelihood 而不是 likelihood，也就是说最小化 $-log(o_y)$ 的值，这两者结果在数学上是等价的)。后面这个操作就是 caffe 文档里说的 Multinomial Logistic Loss，具体写出来是这个样子：
$\ell(y, o)=-\log \left(o_{y}\right)$

3. Softmax-Loss = Softmax + Multinomial Logistic Loss

而 Softmax-Loss 其实就是把两者结合到一起，只要把 $o_y$ 的定义展开即可：
$\tilde{\ell}(y, z)=-\log \left(\frac{e^{z_{y}}}{\sum_{j=1}^{m} e^{z_{j}}}\right)=\log \left(\sum_{j=1}^{m} e^{z_{j}}\right)-z_{y} \\ {\color{Green}{\small{【将观察到的数据 \ x \ (或z)\ 属于类别 \ y \ 的概率做最大似然估计，或最小\ negative\ log\ 似然】}}} \\$

比如如果我们要写一个 Logistic Regression 的 solver，那么因为要处理的就是这个东西，比较自然地就可以将整个东西合在一起来考虑，或者甚至将 $z_i = w_i^Tx + b_i$ 的定义直接一起带进去然后对 $w$ 和 $b$ 进行求导来得到 Gradient Descent 的 update rule.

反过来，如果是在设计 Deep Neural Networks 的库，则可能会倾向于将两者分开来看待：因为 Deep Learning 的模型都是一层一层叠起来的结构，一个计算库的主要工作是提供各种各样的 layer，然后让用户可以选择通过不同的方式来对各种 layer 组合得到一个网络层级结构就可以了。比如用户可能最终目的就是得到各个类别的概率似然值，这个时候就只需要一个 Softmax Layer，而不一定要进行 Multinomial Logistic Loss 操作；或者是用户有通过其他什么方式已经得到了某种概率似然值，然后要做最大似然估计，此时则只需要后面的 Multinomial Logistic Loss 而不需要前面的 Softmax 操作。因此提供两个不同的 Layer 结构比只提供一个合在一起的 Softmax-Loss Layer 要灵活许多。从代码的角度来说也显得更加模块化。但是这里自然地就出现了一个问题：numerical stability。

• Softmax-Loss 单层损失层的梯度计算

假设我们直接使用一层 Softmax-Loss 层，计算输入数据 $z_k$ 属于类别 $y$ 的概率的极大似然估计。由于 Softmax-Loss 层是最顶层的输出层，则可以直接用最终输出 (loss): $\tilde{\ell}(y, z)$ 求对输入 $z_k$ 的偏导数：
$\begin{aligned} \frac{\partial \tilde{\ell}(y, z)}{\partial z_{k}} & = \frac{\partial}{\partial z_{k}} \left(\log \left(\sum_{j=1}^{m} e^{z_{j}}\right)-z_{y} \right)\\ & = \frac{\exp \left(z_{k}\right)}{\sum_{j=1}^{m} \exp \left(z_{j}\right)}-\delta_{k y}=\sigma_{k}(z)-\delta_{k y} \end{aligned}$
其中 $\sigma_k(z)$ 是 Softmax-Loss 的中间步骤 Softmax 在 Forward Pass 的计算结果，而
$\delta_{k y}=\left\{\begin{array}{ll}{1} & {k=y} \\ {0} & {k \neq y}\end{array}\right.$

即 Softmax-Loss 层的梯度为：
$\begin{aligned} \frac{\partial \tilde{\ell}(y, z)}{\partial z_{k}} = \begin{cases} \sigma_{k}(z) - 1 , & k = y \\ \sigma_{k}(z) , &k \ne y \end{cases} \end{aligned}$

• Softmax + Multinomial Logistic Loss 两层分开叠加构造的损失层梯度的计算

接下来看，如果是 Softmax 层和 Multinomial Logistic Loss 层分成两层会是什么样的情况呢？继续回忆刚才的记号：我们把 Softmax 层的输出，也就是 Loss 层的输入记为 $o_{i}=\sigma_{i}(z)$，因此我们首先要计算顶层的 Multinomial Logistic Loss 层输出，对 Softmax 层输入的梯度：
$\begin{aligned} \frac{\partial \ell(y, o)}{\partial o_{i}} & = \frac{\partial}{\partial o_{i}} \left(-\log \left(o_{y}\right)\right) \\ & = -\frac{\delta_{i y}}{o_{y}} \end{aligned}$
即 Multinomial Logistic Loss 层梯度为：
$\begin{aligned} \frac{\partial \ell(y, o)}{\partial o_{i}} \begin{cases} -\frac{1}{o_{y}}, &i = y \\ 0 , &i \ne y \end{cases} \end{aligned}$
然后我们把这个导数向下传递，现在到达 Softmax 层，在 apply chain rule 之前，首先计算层内的导数
$\begin{aligned} \frac{\partial o_{i}}{\partial z_{k}} & = \frac{\partial}{\partial z_{k}} \frac{\exp \left(z_{i}\right)}{\sum_{j=1}^{m} \exp \left(z_{j}\right)} \\ & =\frac{\delta_{i k} e^{z_{i}}\left(\sum_{j=1}^{m} e^{z_{j}}\right)-e^{z_{i}} e^{z_{k}}}{\left(\sum_{j=1}^{m} e^{z_{j}}\right)^{2}} \\ & =\delta_{i k} o_{i}-o_{i} o_{k} \end{aligned}$
即 Softmax 层的梯度为：
$\begin{aligned} \frac{\partial o_{i}}{\partial z_{k}} \begin{cases} o_i(1 - o_k), & k = i \\ -o_i o_k, & k \ne i \end{cases} \end{aligned}$
如果用 Chain Rule 带进去验算一下的话：
$\begin{aligned} \sum_{i=1}^{m} \frac{\partial o_{i}}{\partial z_{k}} \cdot \frac{\partial \ell(y, o)}{\partial o_{i}} & = (\delta_{i k} o_{i}-o_{i} o_{k}) \cdot (-\frac{\delta_{i y}}{o_{y}}) \\ & {\color{Red}{\small{\Downarrow 根据链式法则，当 i = y \ 才能往前传递，即把上面的\ i\ 都用\ y \ 替换即可}}} \\ & = (\delta_{y k} o_{y}-o_{y} o_{k}) \cdot (-\frac{1}{o_{y}}) \\ & = o_{k}-\delta_{y k} \end{aligned}$

和刚才的结果一样的，看来我们求导没有求错。虽然最终结果是一样的，但是我们可以看出，如果分成两层计算的话，要多算好多步骤，除了计算量增大了一点，我们更关心的是数值上的稳定性。由于浮点数是有精度限制的，每多一次运算就会多累积一定的误差，注意到分成两步计算的时候我们需要计算 $\delta_{iy}/o_y$ 这个量，如果碰巧这次预测非常不准，$o_y$ 的值，也就是正确的类别所得到的概率非常小（接近零）的话，这里会有 overflow 的危险。下面我们来实际试验一下，首先定义好两种不同的计算函数：

function softmax(z)
  #z = z - maximum(z)
  o = exp(z)
  return o / sum(o)
end

function gradient_together(z, y)
  o = softmax(z)
  o[y] -= 1.0
  return o
end

function gradient_separated(z, y)
  o = softmax(z)
  ∂o_∂z = diagm(o) - o*o'
  ∂f_∂o = zeros(size(o))
  ∂f_∂o[y] = -1.0 / o[y]
  return ∂o_∂z * ∂f_∂o
end

然后由于 float (Float32) 比 double (Float64) 的精度要小很多，我们就以 double 的计算结果为近似的“正确值”，然后来比较两种情况下通过 float 来计算得到的结果和正确值之差。绘图代码如下：

using DataFrames
using Gadfly

M = 100
y = 1
zy = vec(10f0 .^ (-38:5:38)) # float range ~ [1.2*10^-38, 3.4*10^38]
zy = [-reverse(zy);zy]
srand(12345)
n_rep = 50

discrepancy_together = zeros(length(zy), n_rep)
discrepancy_separated = zeros(length(zy), n_rep)

for i = 1:n_rep
  z = rand(Float32, M)  # use float instead of double
  
  discrepancy_together[:,i] = [begin
    z[y] = x
    true_grad = gradient_together(convert(Array{Float64},z), y)
    got_grad = gradient_together(z, y)
    abs(true_grad[y] - got_grad[y])
  end for x in zy]
  discrepancy_separated[:,i] = [begin
    z[y] = x
    true_grad = gradient_together(convert(Array{Float64},z), y)
    got_grad = gradient_separated(z, y)
    abs(true_grad[y] - got_grad[y])
  end for x in zy]
end

df1 = DataFrame(x=zy, y=vec(mean(discrepancy_together,2)),
                label="together")
df2 = DataFrame(x=zy, y=vec(mean(discrepancy_separated,2)),
                label="separated")
df = vcat(df1, df2)

format_func(x) = @sprintf("%s10<sup>%d</sup>", x<0?"-":"",int(log10(abs(x))))
the_plot = plot(df, x="x", y="y", color="label",
                Geom.point, Geom.line, Geom.errorbar,
                Guide.xticks(ticks=int(linspace(1, length(zy), 10))),
                Scale.x_discrete(labels=format_func),
                Guide.xlabel("z[y]"), Guide.ylabel("discrepancy"))

这里我们做的事情是保持 $z$ 的其他座标不变，而改变 $z_y$ 也就是对应于真是 label 的那个座标的数值大小，我们刚才的推测是当 $o_y$ 很接近零的时候会有 overflow 的危险，而 $o_y = \sigma_y(z)$，忽略掉 normalization 的话，正比于 $exp(z_y)$，所以我们需要把 $z_y$ 那个座标设成绝对值很大的负数。在得到的图中我们可以看到以整个数值范围内的情况对比。 图中横座标是 $z_y$ 的大小，纵座标是分别用两种方法计算出来的结果和“真实值”之间的差距大小。

首先可以看到的是单层直接计算确实比分成两层算要好一点，不过从纵座标上也可以看到两者差距其实非常小。往左边看的话，会发现黄色的点没有了，那是因为结果得到了 NaN了，比如 $o_y$ 由于求一个绝对值非常大的负数的 exponential，导致下溢超出 float 可以表示的小数点精度范围，直接变成 0 了，此时 $1/o_y$ 就是 Inf，当要乘以 $o_y$ 进行 cancel 的时候得到 $0 \times \infty$，对于浮点数这个操作会直接得到 NaN，也就是 Not a Number。反过来看蓝线的话，好像有点奇怪的是越往左边好像反而变得更加精确了，其实是因为我们的“真实值”也 underflow 了，因为 double 虽然比 float 精度高很多，但是也是有限制的。根据 Wikipedia，float 的精度范围大致是 $10^{-38} \sim 10^{38}$，而 double 的精度范围大致是 $10^{-308} \sim 10^{308}$，大了很多，但是我们不妨来看一下图中的 $-10^2$ 这个座标点，注意到
$e^{x}=10^{x / \log 10}$
所以 $\exp \left(-10^{2}\right) \approx 10^{-44}​$，对于 float 来说已经下溢了，对于 double 来说还是可以表示的范围，但是和 0 的差别也已经如此小，在图上已经看不出区别来了。指数再移一格的话，$\exp \left(-10^{3}\right) \approx 10^{434}​$，会直接导致 double 也 underflow，结果我们的“真实值”也会是零，所以“误差”直接变成零了。

比较有趣的是往右边的正数半轴看，发现到了 $10^2$ 之后蓝线和黄线都没有了，说明他们都得到了 NaN，不过这里是另一个问题：对一个比较大的数求 exponential 非常容易发生 overflow。还是用刚才的式子可以看到 ，已经超过了 float 可以表达的最大上限，所以会变成 Inf，然后在 normalize 的一步会出现 Inf/Inf 这样的情况，于是就得到 NaN 了。

这个问题其实也是有解决办法的，我们刚才贴的代码里的 softmax 函数第一行有一行被注释掉的代码，就是 在求 exponential 之前将 $z$ 的每一个元素减去 $z_i$ 的最大值。这样求 exponential 的时候会碰到的最大的数就是 0 了，不会发生 overflow 的问题，但是如果其他数原本是正常范围，现在全部被减去了一个非常大的数，于是都变成了绝对值非常大的负数，所以全部都会发生 underflow，但是 underflow 的时候得到的是 0，这其实是非常 meaningful 的近似值，而且后续的计算也不会出现奇怪的 NaN。

证明：将输入 $z$ 的每一个元素都减去 $z_i$ 元素中的最大值，然后求 Softmax 函数结果相等

$\begin{aligned} \sigma_i(z) & = \frac{e^{z_i}}{\sum_{j = 1}^{m}e^{z_j}} \\ & = \frac{e^{z_i - z_{max}}}{\sum_{j = 1}^{m}e^{z_j - z_{max}}} \\ & = \frac{e^{z_i }/e^{z_{max}}}{\sum_{j = 1}^{m}e^{z_j}/e^{z_{max}}} \\ & = \frac{e^{z_i}}{\sum_{j = 1}^{m}e^{z_j}} \end{aligned}$

4. Reference

• Softmax vs. Softmax-Loss: Numerical Stability
• Softmax UFLDL
• https://blog.csdn.net/u014313009/article/details/51045303
• https://blog.csdn.net/u014380165/article/details/77284921
• https://eli.thegreenplace.net/2016/the-softmax-function-and-its-derivative/