Softmax 구현의 비밀: 왜 최댓값을 빼는가?

• •Softmax 공식은 exp(xi) / sum(exp(xj)) 형태이다
• •입력값이 커질수록 지수 함수의 결과값이 급격히 증가하여 오버플로우 위험이 있다
• •Transformer 구조 내부에는 수많은 Softmax가 존재하여 관리가 필요하다

Softmax는 다중 클래스 분류 문제에서 모델의 출력을 확률로 해석하기 위해 마지막 레이어에 주로 사용된다.

• •입력 벡터의 최댓값 C를 모든 요소에서 빼는 방식을 사용한다
• •수학적으로는 원래의 Softmax 공식과 완전히 동일한 결과를 보장한다
• •지수 함수의 결과값을 1 이하로 억제하여 수치적 안정성을 확보한다

지수 법칙 e^(a-b) = e^a / e^b를 이용한 수학적 트릭이다.

• •최댓값 뺄셈을 적용하지 않으면 저정밀도 연산에서 nan 또는 inf가 발생한다
• •이 기법은 딥러닝 프레임워크 내부(PyTorch 등)에 기본적으로 구현되어 있다
• •Flash Attention 이해를 위한 필수적인 기초 지식이다

float16은 표현 가능한 수의 범위가 좁아 오버플로우에 매우 취약하다.