검벨-소프트맥스 트릭 완벽 정리: 이산 분포를 미분 가능하게 샘플링하는 원리와 온도 어닐링

핵심 메시지

• 신경망 순전파 도중 범주형(이산) 분포에서 샘플을 뽑으면 argmax 때문에 미분이 끊겨 역전파로 학습할 수 없다.
• 검벨-맥스 트릭은 각 범주의 로그 확률에 검벨 노이즈를 더한 뒤 argmax를 취하면 원래 분포에서 그대로 샘플링한 것과 같아진다는 성질을 이용한다.
• argmax는 미분 불가능하므로 이를 부드러운 근사인 소프트맥스로 바꾼 것이 검벨-소프트맥스이며, 이로써 경사를 역전파할 수 있다.
• 온도(temperature) 파라미터로 출력이 원-핫(이산)에 가까운지, 부드러운(연속) 분포에 가까운지를 조절한다.
• 학습 초기에 큰 온도로 시작해 점차 낮추는 '온도 어닐링'으로 안정적으로 학습하면서 점점 이산 샘플에 수렴시킨다.

쉽게 이해하기

이 트릭을 이해하려면 먼저 재매개변수화(reparameterization) 트릭을 떠올리는 것이 좋다. 변분 오토인코더(VAE)에서는 입력을 저차원 잠재변수 z로 인코딩한 뒤 가우시안 분포에서 z를 샘플링해 다시 이미지를 복원한다. 이때 샘플링 노드가 끼면 도함수를 제대로 추정할 수 없어 역전파가 막히는데, 무작위성(엡실론)과 학습 파라미터(평균·표준편차)를 분리해 이 문제를 푼다.

검벨-맥스 트릭이 푸는 문제는 같은 '샘플링 때문에 미분이 끊기는' 문제이지만, 대상이 연속 분포가 아니라 범주형(이산) 분포라는 점이 다르다. 예를 들어 모델이 순전파 중 어떤 범주를 골라 캡션을 생성해야 한다면 범주형 분포에서 샘플을 뽑아야 하는데, 그 선택을 표현하는 argmax는 미분이 되지 않아 그대로는 학습할 수 없다.

해법은 각 범주의 로그 확률에 '검벨 노이즈'를 더하는 것이다. 검벨 분포는 원래 폭우나 나일강 범람 같은 극단값을 모델링하던 분포로, 위치(mu)와 척도(beta)를 갖는다. 여기서는 mu=0, beta=1인 표준 검벨 노이즈를 쓰며, 정규분포처럼 바로 뽑는 함수가 없어 균등분포 U에서 값을 뽑아 -log(-log(U)) 변환으로 만든다.

노이즈를 더한 뒤 argmax를 취하면 원래 확률대로 정확히 한 범주가 뽑히지만 여전히 미분이 안 된다. 그래서 argmax를 부드러운 근사인 소프트맥스로 바꾼다. 소프트맥스에는 온도 파라미터가 들어가는데, 온도가 0에 가까우면 출력이 원-핫에 가까워 진짜 범주형처럼 되고, 온도가 크면 출력이 부드러워져 연속 분포에 가까워진다.

온도가 작을수록 이산성은 좋아지지만 경사가 매우 불안정해 학습이 어렵다. 그래서 '온도 어닐링'을 쓴다. 학습 초기에는 온도를 크게 두어 매끄럽게 학습하고, 진행하면서 점차 온도를 낮춰 마지막에는 거의 원-핫 표현에 수렴시킨다. 이렇게 하면 역전파가 가능한 상태로 이산 분포에서 샘플링하는 학습을 끝까지 수행할 수 있다.

주요 인사이트

• 재매개변수화 트릭(연속)과 검벨-소프트맥스(이산)는 서로 다른 문제를 풀지만, '무작위성과 학습 파라미터를 분리한다'는 같은 아이디어에 뿌리를 둔다.
• 검벨 분포는 본래 극단값을 모델링하던 통계 도구인데, 여기서는 노이즈로 활용해 확률이 비슷한 범주 사이의 '동점'을 깨는 역할을 한다.
• 검벨 노이즈는 정규분포처럼 바로 뽑는 내장 함수가 없어, 균등분포에서 뽑은 값에 -log(-log(U)) 변환을 적용해 생성한다.
• 온도가 0에 가까우면 원-핫에 가까워 진짜 이산 분포 같지만 경사가 불안정하고, 온도가 크면 학습은 쉬우나 이산성이 약해진다 — 어닐링이 이 둘을 시간에 따라 절충한다.

자주 묻는 질문

관련 AI 소식