소프트맥스 함수(softmax function)는 소프트아그맥스(softargmax) 또는 정규화된 지수 함수(normalized exponential function)라고도 불리며, 실수 벡터를 확률 분포로 변환하는 수학적 함수이다. 이는 로지스틱 함수를 다차원으로 일반화한 것이며, 특히 다중 클래스 분류 문제에서 머신 러닝 및 딥 러닝에 널리 사용된다.
정의
입력 벡터 \( \mathbf{z} = (z_1, \dots, z_K) \in \mathbb{R}^K \) (여기서 \( K > 1 \))가 주어졌을 때, 소프트맥스 함수 \( \sigma: \mathbb{R}^K \to (0,1)^K \)는 각 성분 \( i \)에 대해 다음과 같이 정의된다:
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} \]
출력은 합이 1이 되는 양의 값들의 벡터이므로, 범주형 확률 분포로 해석할 수 있다. 지수 연산은 입력 값의 차이를 증폭시키므로, 더 큰 입력 값은 더 큰 확률을 얻게 된다. 예를 들어, 벡터 \( (1, 2, 8) \)에 소프트맥스를 적용하면 대략 \( (0.001, 0.002, 0.997) \)이 되어, 거의 모든 확률 질량이 가장 큰 원소에 할당된다.
속성
소프트맥스 함수는 다음과 같은 중요한 속성을 가진다:
- 출력 범위: 각 성분은 엄격히 0과 1 사이의 값을 가진다.
- 합이 1: 모든 출력 성분의 합은 정확히 1이다.
- 순서 보존: 만약 \( z_i > z_j \)이면, \( \sigma(\mathbf{z})_i > \sigma(\mathbf{z})_j \)이다.
- 상수 이동 불변성: 모든 입력에 상수 \( c \)를 더해도 출력은 변하지 않는다. 즉, \( e^{z_i + c} / \sum_j e^{z_j + c} = e^{z_i} / \sum_j e^{z_j} \)이다. 이 속성은 계산 시 입력에서 최댓값을 빼서 수치적 안정성을 확보하는 데 자주 사용된다.
온도 매개변수
소프트맥스 함수의 변형으로 온도 매개변수 \( \beta \) (또는 그 역수 \( T = 1/\beta \))를 도입할 수 있다:
\[ \sigma(\mathbf{z})_i = \frac{e^{\beta z_i}}{\sum_{j=1}^{K} e^{\beta z_j}} \]
\( \beta > 1 \)이면 출력 분포는 최댓값 주변에 더 집중되어 "날카로워지고", \( 0 < \beta < 1 \)이면 분포는 더 균일해진다. 이 매개변수는 top-k 샘플링 및 온도 스케일링에서 출력의 무작위성을 제어하는 데 흔히 사용되며, 특히 생성형 AI 모델에서 활용된다.
응용
소프트맥스는 많은 신경망 구조에서 핵심 구성 요소이다:
- 분류: 출력 계층의 활성화 함수로 사용되어 원시 로짓(logits)을 클래스 확률로 변환하며, 교차 엔트로피 손실과 함께 학습된다.
- 어텐션 메커니즘: 트랜스포머 모델에서 시퀀스 위치 간의 어텐션 가중치를 계산하는 데 사용된다. 예를 들어, 멀티헤드 어텐션에서 쿼리와 키의 내적에 소프트맥스를 적용하여 가중치를 정규화한다.
- 강화 학습: 행동 선호도(선호 점수)를 확률적 정책으로 변환하여 행동 선택에 사용된다.
- 혼합 모델: 가우시안 혼합 모델과 같은 모델에서 혼합 가중치를 계산하는 데 사용된다.
수치적 안정성
지수 함수는 입력 값이 클 때 오버플로우(overflow)를 일으킬 수 있다. 일반적인 해결책은 지수 연산 전에 입력에서 최댓값을 빼는 것이다:
\[ \sigma(\mathbf{z})_i = \frac{e^{z_i - \max(\mathbf{z})}}{\sum_{j=1}^{K} e^{z_j - \max(\mathbf{z})}} \]
이 변환은 상수 이동 불변성 덕분에 결과를 변경하지 않으면서, 가장 큰 지수 값이 0이 되도록 하여 오버플로우를 방지한다.
관련 개념
소프트맥스 함수는 \( K = 2 \)인 경우 로지스틱 함수의 특수한 형태이다. 또한 손실 함수인 교차 엔트로피와 밀접하게 관련되어 있으며, 시퀀스 생성에서 빔 서치 디코딩과 함께 사용된다. 대규모 언어 모델 추론에서는 로짓에 소프트맥스를 적용하여 토큰 확률을 얻고, 온도 조정을 통해 창의성과 일관성 사이의 균형을 맞춘다.