Softmax와 Cross Entropy

softmax 함수와 그 도함수

softmax 함수는 다음과 같이 정의된다:

\[\operatorname{softmax}(\mathbf{z})_i = \sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{k=1}^{K} e^{z_k}} \qquad \text{for } i = 1, 2, \ldots, K\]

$z_j$에 대한 도함수를 구하려면 몫의 미분법을 적용한다. 이하에서 $\sum_k$는 $\sum_{k=1}^{K}$를 뜻하고, $\delta_{ij}$는 크로네커 델타($i = j$이면 $1$, 아니면 $0$)이다. 앞으로 쓸 유용한 성질은 sifting 이다: $\sum_k a_k \delta_{kj} = a_j$.

\[\begin{aligned} \frac{\partial \sigma(\mathbf{z})_i}{\partial z_j} &= \frac{\left(\dfrac{\partial}{\partial z_j}e^{z_i}\right)\sum_k e^{z_k} - e^{z_i}\left(\dfrac{\partial}{\partial z_j}\sum_k e^{z_k}\right)} {\left(\sum_k e^{z_k}\right)^2} \\[8pt] &= \frac{e^{z_i}\,\delta_{ij}\sum_k e^{z_k} - e^{z_i}e^{z_j}} {\left(\sum_k e^{z_k}\right)^2} \\[8pt] &= \frac{e^{z_i}}{\sum_k e^{z_k}} \left(\delta_{ij} - \frac{e^{z_j}}{\sum_k e^{z_k}}\right) \\[8pt] &= \sigma(\mathbf{z})_i \left(\delta_{ij} - \sigma(\mathbf{z})_j\right) \\[8pt] &= \begin{cases} \sigma(\mathbf{z})_i \left(1 - \sigma(\mathbf{z})_i\right) & \text{if } i = j \\ -\,\sigma(\mathbf{z})_i \,\sigma(\mathbf{z})_j & \text{if } i \neq j \end{cases} \end{aligned}\]

이제 softmax 함수의 전체 Jacobian 행렬을 얻었다:

\[J_{\sigma}(\mathbf{z}) = \begin{bmatrix} \sigma(\mathbf{z})_1 (1 - \sigma(\mathbf{z})_1) & -\sigma(\mathbf{z})_1 \sigma(\mathbf{z})_2 & \cdots & -\sigma(\mathbf{z})_1 \sigma(\mathbf{z})_K \\ -\sigma(\mathbf{z})_2 \sigma(\mathbf{z})_1 & \sigma(\mathbf{z})_2 (1 - \sigma(\mathbf{z})_2) & \cdots & -\sigma(\mathbf{z})_2 \sigma(\mathbf{z})_K \\ \vdots & \vdots & \ddots & \vdots \\ -\sigma(\mathbf{z})_K \sigma(\mathbf{z})_1 & -\sigma(\mathbf{z})_K \sigma(\mathbf{z})_2 & \cdots & \sigma(\mathbf{z})_K (1 - \sigma(\mathbf{z})_K) \end{bmatrix}\]

더 간결하게 쓰면,

\[J_{\sigma}(\mathbf{z}) = \operatorname{diag}\left(\sigma(\mathbf{z})\right) - \sigma(\mathbf{z})\,\sigma(\mathbf{z})^{\top},\]

여기서 $\delta_{ij}$ 항이 대각 성분에, $-\sigma_i\sigma_j$ 항이 외적에 대응한다.

그런데 전체 Jacobian 행렬이 정말 필요할까? 실제로 필요한 것은 softmax의 입력에 대한 손실의 도함수뿐이다. 즉 모든 $i$, $j$에 대한 $\frac{\partial \sigma(\mathbf{z})_i}{\partial z_j}$가 아니라 다음이 관심 대상이다.

\[\frac{\partial L}{\partial z_j} = \sum_{i=1}^{K} \frac{\partial L}{\partial \sigma(\mathbf{z})_i} \frac{\partial \sigma(\mathbf{z})_i}{\partial z_j}.\]

Cross Entropy 손실과 그 도함수

여기서 예측 확률은 softmax의 출력이다: $\hat{\mathbf{y}} = \sigma(\mathbf{z})$. 정답 $\mathbf{y}$는 확률 분포, 즉 $\sum_{i} y_i = 1$이라고 가정한다(one-hot 레이블이 가장 흔한 경우다). 예제 하나에 대한 Cross Entropy 손실은 이렇게 정의된다:

\[L(\mathbf{y}, \hat{\mathbf{y}}) = -\sum_{i=1}^{K} y_i \log(\hat{y}_i)\]

예측 확률에 대한 도함수는 다음과 같다:

\[\frac{\partial L}{\partial \hat{y}_i} = -\frac{y_i}{\hat{y}_i}\]

softmax와 Cross Entropy 합치기

두 도함수를 연쇄 법칙에 대입하면 놀라운 일이 벌어진다:

\[\begin{aligned} \frac{\partial L}{\partial z_j} &= \sum_{i=1}^{K} \frac{\partial L}{\partial \hat{y}_i} \frac{\partial \hat{y}_i}{\partial z_j} \\[4pt] &= \sum_{i=1}^{K} \left(-\frac{y_i}{\hat{y}_i}\right) \hat{y}_i \left(\delta_{ij} - \hat{y}_j\right) && (\hat{y}_i = \sigma(\mathbf{z})_i) \\[4pt] &= \sum_{i=1}^{K} (-y_i)\left(\delta_{ij} - \hat{y}_j\right) && \left(\tfrac{1}{\hat{y}_i} \cdot \hat{y}_i = 1\right) \\[4pt] &= -\sum_{i=1}^{K} y_i\,\delta_{ij} + \hat{y}_j \sum_{i=1}^{K} y_i \\[4pt] &= -y_j + \hat{y}_j \underbrace{\sum_{i=1}^{K} y_i}_{=\,1} && (\text{sifting}) \\[4pt] &= \hat{y}_j - y_j \end{aligned}\]

$K \times K$ Jacobian 행렬이 단순한 뺄셈으로 무너졌다. 예측값 빼기 정답. $\mathbf{y}$가 one-hot일 때 이것은 밑바닥부터 구현한 코드마다 등장하는 probs - onehot 계산 그 자체이고, PyTorch 같은 프레임워크가 두 연산을 F.cross_entropy 하나로 융합해두는 이유이기도 하다. 합쳐진 기울기가 이렇게 단순한데 전체 Jacobian 행렬을 계산하는 건 낭비이기 때문이다. 유도 과정에서 $\sum_i y_i = 1$만 사용했으므로, 이 결과는 label smoothing을 포함해 임의의 분포형 정답에 대해 성립한다.

다른 길: one-hot 정답을 위한 지름길

위의 일반적인 유도는 임의의 분포형 정답에 대해 작동하지만, 처음부터 one-hot $\mathbf{y}$로 못박는다면 Jacobian 행렬을 아예 피해 가는 더 짧은 길이 있다.

참 클래스를 $c$라 하면 $y_i = \delta_{ic}$이다. Cross Entropy 손실은 항 하나로 줄어든다:

\[L = -\sum_{i=1}^{K} y_i \log(\hat{y}_i) = -\log(\hat{y}_c)\]

이제 softmax를 직접 대입하고 로그를 쪼갠다:

\[L = -\log\left(\frac{e^{z_c}}{\sum_k e^{z_k}}\right) = -z_c + \log\sum_k e^{z_k}\]

로그가 분자의 지수를 지워버려서, $L$이 로짓에 대한 명시적인 함수로 남았다. 중간에 $\hat{\mathbf{y}}$를 거쳐 연쇄시킬 필요가 없다. 항별로 미분하면:

\[\frac{\partial L}{\partial z_j} = \underbrace{-\,\delta_{jc}}_{\partial(-z_c)/\partial z_j} + \underbrace{\frac{e^{z_j}}{\sum_k e^{z_k}}}_{\partial \log\sum_k e^{z_k} / \partial z_j} = \hat{y}_j - \delta_{jc}\]

첫 항은 $j = c$일 때 $-1$이고 아니면 $0$이다. 둘째 항은 잠시 멈춰 볼 만하다. log-sum-exp의 도함수가 바로 softmax 자신이다. 돌이켜보면 이것이 앞 절의 소거가 일어날 수밖에 없었던 진짜 이유다. softmax는 처음부터 $\log\sum_k e^{z_k}$의 기울기였고, 거기에 로그를 합성하면 그 구조가 그대로 드러날 뿐이다.

$\delta_{jc}$는 정확히 one-hot 정답의 $j$번째 성분이므로, $\hat{y}j - \delta{jc}$는 당연하게도 일반 결과 $\hat{y}_j - y_j$와 일치한다.

그렇다면 왜 굳이 긴 유도를 했을까? 지름길이 더 빠르긴 하지만, one-hot 가정이 첫 줄부터 박혀 있다. 일반적인 경로는 Jacobian 행렬과 몇 번의 첨자 조작을 대가로 치르는 대신, $\sum_i y_i = 1$을 만족하는 어떤 정답에도 성립하는 결과를 준다. label smoothing까지 포함해서.