Softmax-CE Loss + Gradient [medium]

v1 Softmax CE 손실 은 forward 만. 학습에는 gradient 가 필수. 놀라운 사실: softmax 와 cross-entropy 를 연쇄적으로 미분 하면 극단적으로 간결한 식이 나옴:

$\frac{\partial \mathcal{L}}{\partial z_{i,k}} = \frac{1}{N} \left( \text{softmax}(z_i)_k - \mathbb{1}[k = y_i] \right)$

즉 grad = (softmax(z) - one_hot(y)) / N — 둘을 따로 미분해서 곱하는 것보다 훨씬 수치적으로 안정적 이고 빠름. PyTorch CrossEntropyLoss.backward(), TF softmax_cross_entropy_with_logits_v2 내부가 바로 이 식.

왜 이 형태?

$\frac{\partial}{\partial z_k} \log p_k = \frac{1}{p_k} \cdot p_k(1-p_k) = 1 - p_k \quad (k=y)$ $\frac{\partial}{\partial z_j} \log p_y = -p_j \quad (j \ne y)$

정답 위치는 $p_y - 1$, 나머지는 $p_j$. 정리:

$\nabla_z \mathcal{L}_i = p_i - e_{y_i}$

$e_{y_i}$ = one-hot of $y_i$. 배치 평균이면 $1/N$.

과제

함수 softmax_ce_loss_grad(logits, y_true) 를 완성하세요.

• logits shape (N, K), y_true shape (N,) 정수.
• 반환: (loss, grad)
  • loss: Python float — 평균 CE (nats).
  • grad: shape (N, K) — ∂L/∂z.
• log-sum-exp 로 loss 안정화, grad 는 softmax - one-hot.

테스트 케이스