본 논문은 현대 딥러닝 프레임워크에서 사용되는 softmax, log-softmax, cross-entropy 구현에 내재된 두 가지 수치적 오류를 정의하고, 이를 해결하기 위한 안정적인 알고리즘을 제안합니다.

1. 핵심 수치적 오류

• Softmax Collapse (소프트맥스 붕괴): 모델이 정답 클래스에 대해 과도한 확신을 가질 경우, logit margin $m = z_{t} - \log\sum_{j\neq t} e^{z_{j}}$가 증가합니다. 표준 구현에서는 $1 + e^{-m}$ 형태의 분모를 계산할 때, $m$이 정밀도 임계값(precision cliff, $c_{d}$)을 넘어서면 $e^{-m}$이 1에 비해 너무 작아져 반올림 오차로 인해 사라집니다. 이로 인해 정답 클래스의 그래디언트 $\tilde{g}_{t}$가 0으로 수렴하는 반면, 비정답 클래스의 그래디언트는 0이 되지 않아 그래디언트 합이 0이 되지 않는 비대칭적 오류가 발생합니다. 이는 Adam 등 적응형 최적화 알고리즘에서 Slingshot 현상을 유발합니다.
• Softmax Cancellation (소프트맥스 상쇄): 지식 증류(Knowledge Distillation) 과정에서 학생 로짓 $\mathbf{z}_{\mathrm{S}}$와 교사 로짓 $\mathbf{z}_{\mathrm{T}}$가 매우 가까울 때, softmax 확률값 간의 뺄셈 연산(q - p) 과정에서 정밀도 손실이 발생합니다. 표준 구현은 미세한 차이를 노이즈로 처리하여, 로짓이 동일함에도 불구하고 0이 아닌 그래디언트가 남는 현상을 보입니다.

2. 제안 방법론

저자들은 logit을 명시적으로 이동시키고 log-sum-exp 및 softplus 연산을 재구성하여 위 문제를 해결합니다.

• Softmax Collapse 해결: $i_{\max} = \arg\max_{i} z_{i}$를 기준으로 로짓을 이동시켜 $\tilde{z}_{i_{\max}} = 0$이 되게 합니다. 이후 log-softmax를 $\tilde{z}_{k} - \operatorname{softplus}(-m)$ 형태로 구현하여, $1$을 직접 더하지 않고도 그래디언트가 $-e^{-m}/(1+e^{-m})$ 형태를 유지하게 함으로써 정밀도를 보존합니다.
• Softmax Cancellation 해결: 확률값의 뺄셈 대신 로짓 차이 $\boldsymbol{\delta} = \mathbf{z}_{\mathrm{S}} - \mathbf{z}_{\mathrm{T}}$를 사용하여 log-space에서 계산합니다. $\operatorname{expm1}(x) = e^x - 1$과 $\operatorname{log1p}(x) = \log(1+x)$ 함수를 활용하여 $\mathbf{q} - \mathbf{p}$를 직접 유도함으로써, $\mathbf{z}_{\mathrm{S}} \to \mathbf{z}_{\mathrm{T}}$ 상황에서도 그래디언트가 0으로 정밀하게 수렴하게 합니다.

3. 연구 결과 및 기여

• Slingshot 제거: 제안된 기법은 학습 과정에서 발생하는 손실 함수의 급격한 스파이크를 제거하여 grokking 현상을 안정화하고 학습 속도를 높입니다.
• Margin Equilibrium: 가중치 감쇠(weight decay)와 Adam의 $\epsilon$ 하이퍼파라미터가 로짓 마진의 상한선 $m^{\star}(\epsilon) \approx \log(1/\epsilon)$을 결정하는 평형 상태를 형성함을 증명했습니다.
• 범용성: 제안된 알고리즘은 bf16, fp32, fp64 등 모든 부동소수점 정밀도 환경에서 기존의 정밀도 임계값 문제를 근본적으로 해결함을 확인했습니다.