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/