Softmax

Softmax는 로짓이라고 하는 원시 점수 벡터를 확률 분포로 변환하는 함수입니다. 모든 출력값은 0과 1 사이에 있고, 출력값의 합은 정확히 1입니다. 벡터 z에 대해 softmax(z)ᵢ = e^(zᵢ) / Σⱼ e^(zⱼ)를 계산합니다. 각 점수를 지수화한 다음 모든 지수값의 합으로 나눕니다. Softmax는 여러 클래스 중 하나를 선택하는 신경망의 표준 출력 계층입니다. 또한 트랜스포머의 어텐션 메커니즘에서 유사도 점수를 가중치로 바꾸고, 언어 모델의 어휘 점수를 다음 토큰 확률로 변환합니다. 이 이름은 가장 큰 점수 하나에 모든 확률을 할당하는 argmax를 부드럽고 미분 가능한 형태로 만든 함수라는 점을 나타냅니다.

Softmax 작동 방식#

지수 함수는 모든 출력을 양수로 만들고 점수 간 차이를 확대하므로 가장 큰 로짓이 확률에서 불균형적으로 큰 비중을 차지합니다. Softmax는 로짓의 순서를 보존하므로 점수가 가장 높은 항목의 확률이 항상 가장 높습니다. 로짓 간 차이만 중요하므로 모든 로짓에 같은 상수를 더해도 출력은 바뀌지 않습니다. 라이브러리는 이 특성을 수치 안정성에 활용하기 위해 지수 함수를 계산하기 전에 최대 로짓을 빼는데, e¹⁰⁰⁰은 32비트 부동소수점에서 오버플로되기 때문입니다.

온도 T는 분포의 첨예도를 조절합니다. softmax(z / T)에서 T가 1보다 작으면 최상위 클래스 쪽으로 분포가 뾰족해지고, 1보다 크면 균등 분포에 가까워집니다. 통계역학에서는 같은 수식이 볼츠만 분포로 나타나며, 이때 T는 실제로 온도를 의미합니다. 아래 PyTorch 코드 조각은 두 가지 효과를 모두 보여줍니다.

import torch

logits = torch.tensor([2.0, 1.0, 0.1])

# Divide logits by a temperature before softmax: low sharpens, high flattens
for temperature in [0.5, 1.0, 2.0]:
    probs = torch.softmax(logits / temperature, dim=0)
    print(f"T={temperature}: {[round(p, 3) for p in probs.tolist()]}")

# Adding a constant leaves softmax unchanged, so libraries subtract the max to avoid overflow
big = logits + 1000
naive = big.exp() / big.exp().sum()
print(f"Naive: {naive.tolist()}, torch.softmax: {[round(p, 3) for p in torch.softmax(big, dim=0).tolist()]}")

T = 1일 때 로짓 2.0, 1.0, 0.1은 [0.659, 0.242, 0.099]이 됩니다. 온도를 절반으로 낮추면 분포는 [0.864, 0.117, 0.019]처럼 더 뾰족해지고, 두 배로 높이면 [0.502, 0.304, 0.194]처럼 더 평평해집니다. 모든 로짓에 1,000을 더하면 단순한 수식은 오버플로를 일으켜 [nan, nan, nan]을 출력하지만, torch.softmax는 여전히 [0.659, 0.242, 0.099]를 반환합니다.

학습에서의 Softmax#

분류기는 softmax와 교차 엔트로피 손실을 함께 사용해 학습합니다. 손실은 정답 클래스 y에 할당된 확률의 음의 로그인 −log(pᵧ)이며, 로짓에 대한 그래디언트는 예측 확률에서 원-핫 레이블을 뺀 p − y입니다. 따라서 프레임워크는 두 연산을 결합합니다. PyTorch의 CrossEntropyLoss는 원시 로짓을 받아 내부적으로 log-softmax를 계산하며, 이는 확률을 각각 계산한 다음 로그를 취하는 것보다 안정적입니다. 이렇게 학습한 모델은 로짓을 출력하며, 추론 시 확률이 필요할 때만 softmax를 적용합니다.

Softmax 사용 사례#

• 이미지 분류: 분류기의 마지막 계층은 클래스마다 로짓 하나를 출력하고, softmax는 이를 사용자가 확인하는 신뢰도 점수로 변환합니다.
• 어텐션: 트랜스포머는 쿼리와 모든 키 사이의 유사도 점수에 softmax(QKᵀ / √dₖ)를 계산하므로 각 토큰의 어텐션 가중치 합은 1입니다. √dₖ로 나누면 차원이 커져도 점수가 증가해 softmax가 포화되는 현상을 방지할 수 있습니다.
• 텍스트 생성: 대규모 언어 모델(LLM)은 매 단계마다 수만 개에서 수십만 개에 이르는 어휘 토큰에 softmax를 적용하며, 샘플링 온도가 앞서 설명한 T입니다. 그런 다음 핵 샘플링은 확률 합이 임계값에 도달하는 최소 토큰 집합에서 토큰을 추출합니다.
• 강화 학습: 행동 값에 softmax를 적용하면 대개 최선의 행동을 선택하면서 다른 행동도 탐색하는 정책을 얻을 수 있습니다. 온도에 따라 탐색과 활용의 균형이 달라집니다.

YOLO26 분류 모델은 추론 시 분류 헤드에 softmax를 적용하며, 각 결과의 probs 속성에는 1,000개 ImageNet 클래스에 대한 전체 분포가 포함됩니다.

from ultralytics import YOLO

# Load a YOLO26 classification model trained on ImageNet
model = YOLO("yolo26n-cls.pt")
results = model("https://ultralytics.com/images/bus.jpg")

# probs holds the softmax output of the classification head
probs = results[0].probs
for index, conf in zip(probs.top5, probs.top5conf.tolist()):
    print(f"{results[0].names[index]}: {conf:.3f}")
print(f"Sum over all 1,000 classes: {probs.data.sum().item():.3f}")

가장 가능성이 높은 클래스 5개와 해당 확률을 출력하며, 헤드가 softmax 분포를 반환하므로 1,000개 클래스 전체의 확률 합은 1.000으로 출력됩니다.

Softmax와 관련 함수 비교#

• 시그모이드: 각 점수를 독립적으로 0과 1 사이로 압축하므로 출력값의 합이 1일 필요는 없습니다. 시그모이드는 사람과 백팩이 모두 포함된 이미지처럼 이진 및 다중 레이블 문제에 적합하며, 객체 탐지기에서 각 클래스의 점수를 계산할 때도 일반적으로 사용됩니다. 클래스가 두 개일 때 softmax는 두 로짓의 차이에 대한 시그모이드와 정확히 같습니다.
• Argmax: 가장 큰 점수의 인덱스를 반환하는 미분 불가능한 하드 결정입니다. Softmax는 같은 순위를 유지하면서 미분 가능하므로 학습에 사용할 수 있습니다. 또한 softmax 확률에 argmax를 적용한 결과는 로짓에 argmax를 적용한 결과와 같습니다.
• Log-softmax: 안정성을 위해 직접 계산하는 softmax의 로그이며, 교차 엔트로피 손실에서 사용합니다.
• 온도 스케일링: 학습 후 검증 데이터에 단일 온도를 맞춰 softmax 확률이 모델의 실제 정답률과 일치하도록 하는 보정 방법입니다. 양수 T로 모든 로짓을 나누어도 순서는 유지되므로 예측 클래스는 바뀌지 않습니다.