softmax의 온도: 로짓을 nat으로 바꾸는 환율

이 장을 연 물음으로 돌아가자. LLM에서 답을 뽑을 때 조절하는 temperature는 물리의 온도와 같은 역할을 하고 있을까? 물리에서 kT는 "에너지 몇 J(줄)이 1 nat인가"를 정하는 환율이었다. nat은 자연로그 ln으로 잰 값에 붙이는 단위로, 경우의 수가 e배가 되면 ln W가 1 nat 늘어난다.

숫자로 보기

먼저 가장 익숙한 softmax에 온도를 넣어 보자. 온도 τ(타우)가 있는 softmax는 로짓을 τ로 나눈 뒤 pᵢ ∝ e^(zᵢ/τ)로 확률을 매긴다. ML 문헌은 이 온도를 흔히 T로 적지만, 이 책에서는 물리의 온도 T와 헷갈리지 않게 글자만 τ로 바꿔 쓴다. 로짓 z = (2, 1, 0)에 대해 온도를 바꿔 가며 계산하면 다음과 같다.

τ	확률	엔트로피 (nat)	유효 후보 수 e^H
0.5	0.867, 0.117, 0.016	0.44	1.55
1	0.665, 0.245, 0.090	0.83	2.30
2	0.506, 0.307, 0.186	1.02	2.77
→ 0	1, 0, 0 (argmax)	0	1
→ ∞	1/3, 1/3, 1/3	ln 3 ≈ 1.10	3

표의 첫 두 후보를 견주어 보자. 로짓 차이는 늘 1인데, τ = 1에서 두 확률의 비는 0.665/0.245 ≈ 2.72, 곧 e배이고 로그로 재면 1 nat이다. τ = 0.5에서는 0.867/0.117 ≈ 7.4 = e²배(2 nat), τ = 2에서는 0.506/0.307 ≈ 1.65 = e^0.5배(0.5 nat)다. 같은 로짓 차이 1이 몇 nat으로 바뀌는지를 온도가 정하는 셈이다.

표의 마지막 칸 e^H는 엔트로피를 「확률이 고르게 나뉜 후보 몇 개와 같은 퍼짐인가」로 바꾼 값이다. 후보 세 개에 확률이 고르게 나뉘면 엔트로피가 ln 3이므로 e^H = 3이 되고, 두 개에 고르게 나뉘면 2가 된다. 그래서 τ = 1의 2.30은 후보 두세 개에 고르게 나뉜 것과 같은 퍼짐이라는 뜻이고, 이 값을 유효 후보 수라 부른다.

확률 비의 로그

표에서 본 관계를 식으로 쓰면, 두 후보의 확률 비를 로그로 쓴 값은 다음과 같다.

\ln \frac{\textcolor{#e377c2}{p_i}}{\textcolor{#e377c2}{p_j}} = \frac{\textcolor{#ff7f0e}{z_i} - \textcolor{#ff7f0e}{z_j}}{\textcolor{#8c564b}{\tau}}

\begin{array}{ll} \textcolor{#e377c2}{p_i} & \text{후보 i가 뽑힐 확률} \\ \textcolor{#e377c2}{p_j} & \text{후보 j가 뽑힐 확률} \\ \textcolor{#ff7f0e}{z_i} & \text{후보 i의 로짓} \\ \textcolor{#ff7f0e}{z_j} & \text{후보 j의 로짓} \\ \textcolor{#8c564b}{\tau} & \text{softmax의 온도} \\ i, j & \text{비교하는 두 후보의 번호} \end{array}

이 식을 읽어 보면, 로짓 차이가 τ만큼 벌어질 때 확률 비는 1 nat, 곧 e배가 된다. 다시 말해 τ는 "로짓 몇 단위가 1 nat인가"를 정하는 환율이다. 물리에서 kT가 "에너지 몇 J(줄)이 1 nat인가"를 정하는 것과 같은 자리에 있는 셈이다.

표에서 보듯 온도를 올리면 같은 로짓 차이가 더 적은 nat으로 환산되므로 분포가 퍼지고, 반대로 온도를 낮추면 확률이 가장 큰 로짓에 모인다. LLM에서 샘플링할 때 조절하는 temperature가 바로 이 τ다.

ML에서: 지식 증류

높은 온도가 쓸모 있는 예로 지식 증류 (knowledge distillation, Hinton 외 2015)가 있다. 지식 증류에서는 큰 모델의 출력을 높은 온도로 펼쳐서 작은 모델에게 가르친다. 왜 굳이 온도를 높일까? τ = 1에서는 오답들의 확률이 거의 0이라서 그 사이의 순서가 잘 보이지 않는다. 이를테면 개 사진이 고양이와는 조금 닮았고 트럭과는 닮지 않았다는 정보가 그렇다. 온도를 높이면 이렇게 묻혀 있던 "오답 사이의 순서"가 드러난다.

개 사진에 대한 로짓 (개, 고양이, 트럭) = (9, 4, 0)(예시 값)을 τ = 1과 τ = 4로 softmax한 결과. 온도를 높이면 고양이가 트럭보다 확률이 높다는 순서가 보인다.

문제 10. 체스 레이팅과 승률

체스나 온라인 게임의 엘로(Elo) 레이팅에서는 두 사람의 승산(이길 확률 ÷ 질 확률)을 10^(레이팅 차 / 400)으로 둔다. 레이팅 1600인 사람이 1400인 사람과 둘 때 이길 확률은 얼마인가? 식의 400을 800으로 바꾸면 이길 확률은 어떻게 되는가?

승산이 10^(200/400) = 10^0.5 ≈ 3.16이니까 이길 확률은 3.16/4.16 ≈ 0.76이에요. 400을 두 배로 하면 승산을 반으로 나누면 되니까 1.58, 이길 확률은 1.58/2.58 ≈ 0.61이요.

400은 식의 어디에서 나누고 있죠?

지수 안이요. 레이팅 차를 나누네요. 그럼 10^(200/800) = 10^0.25 ≈ 1.78이고, 이길 확률은 1.78/2.78 ≈ 0.64예요. 승산은 반이 아니라 제곱근이 되고요.

맞아요. 400은 승산이 아니라 승산의 로그, 곧 레이팅 차를 나눠요. 이 수를 키우면 같은 레이팅 차가 더 작은 로그 승산으로 바뀌어서 승부가 운에 가까워지죠. 지진 규모도 그래요. 규모가 1 차이 나면 에너지는 약 32배인데, 규모 차이를 반으로 줄이면 에너지 비는 반이 아니라 제곱근인 약 5.6배가 돼요.

로그 쪽에 걸린 숫자를 바꾸면, 비는 나눗셈이 아니라 거듭제곱으로 바뀌는 거네요.

문제 11. softmax 온도 버그

다음 코드는 로짓 z = (3, 1, 0)에서 온도 τ = 0.5로 샘플링하려는 코드다. 무엇이 틀렸고, 바르게 고치면 확률은 얼마인가? (풀어 본 뒤 위젯 3의 「문제 11 불러오기」로 확인해 보자.)

p = torch.softmax(z, dim=-1)
p = p / tau
p = p / p.sum(dim=-1, keepdim=True)

제가 과제에서 짠 코드랑 똑같은데요. 온도로 나눠 주고 합이 1이 되게 다시 나눴어요(정규화).

문제의 z = (3, 1, 0)으로 돌려 봐요. τ = 1일 때와 결과가 다른가요?

softmax가 (0.844, 0.114, 0.042)이니까 이걸 모두 2로 나누고 다시 합이 1이 되게 나누면… 그대로네요. 모든 확률을 같은 수로 나누면 합을 1로 맞출 때 사라지니까요.

아까 체스 문제에서 내가 승산을 반으로 나눈 거랑 같은 실수네. 400이 레이팅 차를 나누듯이, 온도도 로그 쪽을 나눠야 해.

맞아요. 온도는 확률이 아니라 로짓, 즉 에너지 쪽을 나눠야 해요. torch.softmax(z / tau, dim=-1). 로짓을 2로 나눈 (6, 2, 0)에 softmax를 하는 셈이니까 몇이 나오죠?

(0.980, 0.018, 0.002)요. 첫 후보에 훨씬 몰리네요.

왜 하필 이름이 온도예요? 그냥 로짓의 배율이잖아요.

첫 번째와 두 번째 후보의 확률 비를 로그로 재 봐요.

ln(0.980/0.018) ≈ 4.0 nat이에요. 로짓 차이는 2이고요. 로짓 차이를 τ로 나눈 게 nat이 되네요.

그게 kT의 역할과 같아요. 에너지를 nat으로 바꾸는 환율이에요. 그럼 τ = 0은요?

0으로 나누면 에러죠. 극한으로는 가장 큰 로짓에 확률이 모두 가니까, 코드에서는 argmax로 따로 처리해야 하고요. 이게 β = ∞, 절대영도네요.