일반적으로 Softmax 함수는 일반적으로 모델의 최종 출력층(final layer) 에서 사용된다. Softmax를 통해 출력된 값은 각 클래스에 대한 확률 분포로 해석할 수 있으며, 이를 기반으로 Cross Entropy Loss를 계산하여 ground truth와 예측 결과 간의 차이를 측정한다. 이때, Backpropagation 과정에서 Softmax와 Cross-Entropy에 대한 Chain Rule을 적용하여, Loss를 입력(logit)에 대해 미분함으로써 역전파를 시작한다.
우리는 Softmax와 Cross Entropy를 각각 따로 처리하는 대신 결합하여 하나의 연산처럼 다루는 것이 일반적이다. 이렇게 하면 수치적으로 더 안정적이며, 계산이 간편해진다. 특히, Softmax의 출력이 0에 가까운 경우 \( \log{p_i} \)를 계산할 때 발생할 수 있는 언더플로우 문제를 방지할 수 있다.
이는 최종적으로 다음과 같이 표현된다.
$$ \frac{\partial \mathcal{L}}{\partial z_i} = \tilde{y}_t^{(i)} - \hat{y}_t^{(i)} $$
$$ \frac{\partial \mathcal{L}}{\partial \mathbf{h}_t^{\langle \text{logit} \rangle (i)}} = -\frac{1}{I} \sum_{i=0}^{I-1} \left( \tilde{y}_t^{(i)} - \hat{y}_t^{(i)} \right) $$
우선 여기서 사용되는 Softmax와 Cross Entropy에 대해 기본적인 정보만 알고 넘어가자.
Softmax
Softmax 함수는 다중 클래스 분류 문제에서 각 클래스에 대한 확률 분포를 구하는 데 사용된다. 주어진 로짓 벡터 \( \mathbf{z} = (z_1, z_2, \dots, z_C) \)에 대해, Softmax는 각 요소를 다음과 같이 변환한다.
$$ \text{softmax}(z_i) = \frac{\exp(z_i)}{\sum_{j=1}^{C} \exp(z_j)} $$
여기서 \( C \)는 클래스의 수를 의미한다. Softmax 함수는 입력값을 0과 1 사이의 값으로 정규화하며, 전체 합이 1이 되도록 한다.
Cross Entropy Loss
Cross Entropy Loss는 예측 확률 분포와 실제 정답 분포 간의 차이를 측정하는 함수이다.
모델의 출력 확률 벡터 \( p \)와 정답 분포 \( y \)에 대해 Cross Entropy Loss는 다음과 같이 정의된다.
$$ L = - \sum^{C-1}_{i=0} y_i \log{\hat{y_i}} $$
여기서 \( y_i \) 는 정답 라벨의 one-hot 인코딩 값이며, \( \hat{y_i} \)는 Softmax를 거친 예측 확률이다.
증명
이제 다음 수식이 도출되는 과정을 순서대로 전개해보자.
$$ \frac{\partial \mathcal{L}}{\partial z_i} = \tilde{y}_t^{(i)} - \hat{y}_t^{(i)} $$
1. Softmax 정의
Softmax 함수는 다음과 같이 정의할 수 있다.
$$ \hat{y}_t^{(i)} = \frac{e^{h_t^{(i)}}}{\sum_j e^{h_t^{(j)}}} $$
여기서 \( h_t^{(i)} \)는 logit이다.
2. Cross Entropy Loss 정의
Cross-Entropy Loss \(\mathcal{L}\)는 다음과 같다.
$$ \mathcal{L} = -\sum_{i} \tilde{y}_t^{(i)} \log \hat{y}_t^{(i)} $$
여기서 \(\tilde{y}_t^{(i)}\)는 ground truth 분포를 나타낸다.
3. Loss를 logit \( h_t^{(i)} \)에 대해 미분
Chain Rule을 적용하면 다음과 같다.
\[
\frac{\partial \mathcal{L}}{\partial h_t^{(i)}} = \sum_k \frac{\partial \mathcal{L}}{\partial \hat{y}_t^{(k)}} \frac{\partial \hat{y}_t^{(k)}}{\partial h_t^{(i)}}
\]
먼저, Loss를 softmax 출력 \(\hat{y}_t^{(k)}\)에 대해 미분하면 다음과 같다.
\[
\frac{\partial \mathcal{L}}{\partial \hat{y}_t^{(k)}} = -\frac{\tilde{y}_t^{(k)}}{\hat{y}_t^{(k)}}
\]
Softmax를 logit에 대해 미분하면,
\[
\frac{\partial \hat{y}_t^{(k)}}{\partial h_t^{(i)}} =
\begin{cases}
\hat{y}_t^{(i)} (1 - \hat{y}_t^{(i)}) & \text{if } i = k \\
-\hat{y}_t^{(i)} \hat{y}_t^{(k)} & \text{if } i \neq k
\end{cases}
\]
4. Chain Rule 적용
따라서,
\[
\frac{\partial \mathcal{L}}{\partial h_t^{(i)}} = \sum_k \left( -\frac{\tilde{y}_t^{(k)}}{\hat{y}_t^{(k)}} \right) \frac{\partial \hat{y}_t^{(k)}}{\partial h_t^{(i)}}
\]
이를 두 부분으로 나누어 계산한다.
Case 1. \(k = i\) 인 경우
\[
-\frac{\tilde{y}_t^{(i)}}{\hat{y}_t^{(i)}} \cdot \hat{y}_t^{(i)} (1 - \hat{y}_t^{(i)}) = -\tilde{y}_t^{(i)} (1 - \hat{y}_t^{(i)})
\]
Case 2. \(k \neq i\) 인 경우
\[
\sum_{k \neq i} \left( -\frac{\tilde{y}_t^{(k)}}{\hat{y}_t^{(k)}} \right) \left( -\hat{y}_t^{(i)} \hat{y}_t^{(k)} \right) = \hat{y}_t^{(i)} \sum_{k \neq i} \tilde{y}_t^{(k)}
\]
5. 합치기
과정 4에서 도출된 두 식을 합치면 다음과 같다.
$$ \frac{\partial \mathcal{L}}{\partial h_t^{(i)}} = -\tilde{y}_t^{(i)} (1 - \hat{y}_t^{(i)}) + \hat{y}_t^{(i)} \sum_{k \neq i} \tilde{y}_t^{(k)} $$
이를 풀어서 전개하면 다음과 같다.
\[
\begin{align*}
\frac{\partial \mathcal{L}}{\partial h_t^{(i)}}
&= \left( -\tilde{y}_t^{(i)} + \tilde{y}_t^{(i)} \hat{y}_t^{(i)} \right) + \left( \hat{y}_t^{(i)} \sum_{k \neq i} \tilde{y}_t^{(k)} \right) \\
&= -\tilde{y}_t^{(i)} + \hat{y}_t^{(i)} \left( \tilde{y}_t^{(i)} + \sum_{k \neq i} \tilde{y}_t^{(k)} \right) \\
&= -\tilde{y}_t^{(i)} + \hat{y}_t^{(i)} \times 1 \\
&= \hat{y}_t^{(i)} - \tilde{y}_t^{(i)}
\end{align*}
\]
이므로 다시 간단하게 정리하면 다음과 같다.
$$ \frac{\partial \mathcal{L}}{\partial h_t^{(i)}} = \hat{y}_t^{(i)} - \tilde{y}_t^{(i)} $$
참고로 \( \tilde{y} \)가 확률 분포(probability distribution)이기 때문에 다음이 성립한다.
$$ \sum_{k \neq i} \tilde{y}_t^{(k)} + \tilde{y}_t^{(i)} = 1 $$
6. 최종 결론
그러므로 Softmax와 Cross-Entropy를 함께 사용할 때, logit에 대한 Loss의 gradient는 다음과 같다.
$$ \boxed{\frac{\partial \mathcal{L}}{\partial h_t^{(i)}} = \hat{y}_t^{(i)} - \tilde{y}_t^{(i)}} $$
복잡해 보이는 미분과정이지만 결과적으론 아주 명료하고 깔끔하게 정리가 된다. 이를 토대로 Backpropagation 과정에서 다음 가중치의 update도 chain rule을 이용해 전개할 수 있다.
관련 포스팅
2024.10.27 - [Data Science/ML & DL] - Softmax에 대한 고찰
2025.04.16 - [Data Science/Data Analysis] - Cross Entropy와 KL(Kullback–Leibler) Divergence
참고 자료
https://levelup.gitconnected.com/killer-combo-softmax-and-cross-entropy-5907442f60ba
