Log-Softmax Backward [medium]

v1 log-softmax 의 forward 는 $y_i = x_i - \text{logsumexp}(x)$. 이제 역전파 공식:

$\frac{\partial y_i}{\partial x_j} = \delta_{ij} - s_j, \quad s = \text{softmax}(x)$

upstream gradient $g = \partial L / \partial y$ 가 주어지면:

$\frac{\partial L}{\partial x_j} = \sum_i g_i \left( \delta_{ij} - s_j \right) = g_j - s_j \cdot \sum_i g_i$

벡터화 구현

axis 로 지정된 축을 기준으로 softmax 를 구한 뒤:

s = softmax(x, axis=axis)           # (..., K)
grad_in = g - s * g.sum(axis=axis, keepdims=True)

왜 이게 깔끔한가

• Jacobian 을 명시적으로 만들 필요 없음 — $O(K^2)$ 가 아니라 $O(K)$.
• NLL loss 의 g = -e_y (one-hot) 를 대입하면 $\partial L/\partial x = s - e_y$ (softmax-CE gradient) 로 정확히 환원됨.

불변식

임의의 g 에 대해 $\sum_j \partial L/\partial x_j = \sum g_j - \sum g_i = 0$. 로짓 전체를 상수만큼 올려도 log-softmax 가 shift-invariant 라는 사실의 쌍대.

과제

함수 log_softmax_backward(x, grad_output, axis=-1) 를 완성하세요.

• x, grad_output: 같은 shape.
• 반환: grad_input, 같은 shape.

테스트 케이스