Softmax Cross-Entropy (logits, stable) [medium]

v1 교차 엔트로피 는 확률 분포 두 개 입력. 딥러닝에서는 softmax 출력을 따로 계산해서 clip 하면 두 번의 수치 오차 누적:

1. logits → softmax 변환 (큰 logit 에서 overflow 위험).
2. log(softmax) 계산 (작은 확률 → log 폭주).

Standard idiom: logits 에서 바로 한 번에:

$\text{CE}(z, y) = -z_y + \log \sum_j e^{z_j}$

이걸 log-sum-exp trick 으로 안정화:

$\log \sum_j e^{z_j} = m + \log \sum_j e^{z_j - m}, \quad m = \max_j z_j$

이러면 exp(음수) 만 써서 overflow 없음. PyTorch nn.CrossEntropyLoss, TF softmax_cross_entropy_with_logits 가 모두 이 방식.

입출력

• logits shape (N, K): 각 샘플의 raw 스코어 (softmax 이전).
• labels shape (N,): 정수 클래스 인덱스 ∈ [0, K).
• 반환: shape (N,) 각 샘플 loss, nats 단위 (자연로그).

왜 per-sample?

sklearn, PyTorch 가 per-sample 배열 반환 → reduction ('mean', 'sum') 은 호출부 결정. 배치별 weighting/masking 을 다르게 줄 수 있음.

과제

함수 softmax_ce_with_logits(logits, labels) 를 완성하세요.

• 반환: shape (N,) 배열, dtype float, nats 단위.
• 극단 logit ($|z| \gtrsim 1000$) 에서도 NaN, inf 없이 유한.
• 힌트: m = logits.max(axis=1, keepdims=True), lse = m + log(sum(exp(logits - m), axis=1)).

테스트 케이스

#	이름	검증
1	반환 shape (N,)
2	정답에 큰 logit → loss ≈ 0
3	균등 logit → loss = log(K)
4	큰 logit (z=1000) 에서 finite	안정성
5	naive 구현 (softmax then log) 과 수치적 일치 (저 범위)
6	단조 그래디언트 검증: 오답 logit 증가 → loss 증가
7	PyTorch F.cross_entropy(reduction='none') 와 일치 (있으면)