딥러닝 모델은 본질적으로 미분 가능한 함수들의 조합입니다.
그렇기 때문에 신경망을 학습시킬 때는 대부분의 연산이 미분 가능해야 합니다.
하지만, 종종 이산적인 선택(discrete decision)이 필요한 상황이 필요하게 됩니다.
예를 들어
a) 특정 클래스를 선택할 때
b) 텍스트 생성 모델에서 다음 토큰을 선택할 때
c) 강화학습에서 이산적인 행동(action)을 선택할 때
d) 그래프 신경망에서 노드를 선택할 때
가 있습니다.
이러한 이산 적인 선택은 argmax 나 sampling 등으로 구현이 되며, 미분이 불가능합니다. 이로 인해 역전파(Back propagation)가 끊기고, 모델 학습이 불가능해집니다.
이 문제를 해결하기 위해 제안된 방법이 바로 오늘 포스팅의 주제인
Gumbel Softmax
입니다.
오늘 포스팅에서는 Gumbel Softmax의 개념과, 수식, Pytorch를 활용한 시각화를 통해 다각적으로 이해해보도록 하겠습니다.
1. Gumbel Softmax란?
1) Gumbel Softmax가 필요한 상황
예를 들어 아래와 같이 확률 분포에서 클래스를 하나 샘플링 한다고 생각해보겠습니다.
이때 sample은 단 하나의 클래스 인덱스를 반환합니다. 예를 들어 2라면, 이를 one-hot 벡터로 바꾸어 [0,0,1]이 됩니다.
이때 문제는 이 과정이 연속적이지 않다는 것입니다.
입력 확률이 조금만 바뀌어도 출력은 확 바뀌기기에, 기울기(Gradient)를 계산할 수 없습니다.
기존에는 이를 해결하기 위해
REINFORCE와 같은 policy gradient 방식을 사용하지만,
* gradient의 variance가 크고
* 학습이 느리고 불안정
하다는 문제가 있습니다.
그래서 나온 대안이 바로 Gumbel-softmax 입니다.
이는 이산 샘플링을 연속적인 softmax 형태로 근사하여 미분 가능성을 확보해주는 방법입니다.
2) Gumbel-Max Trick + Softmax 근사
a) Gumbel-Max Trick
- 카테고리 확률 분포인 phi로 부터 샘플링 하고 싶다고 해보겠습니다.
- 이때 Gumbel-Max Trick은 다음을 이용합니다.
- 여기서 gi 는 Gumbel(0,1)분포에 뽑은 noise 입니다.
* Gumbel 분포 샘플링 방법의 수식은 아래와 같으며
즉, 샘플링을 argmax 형태로 바꿔주는 트릭입니다.
b) Softmax 근사: argmax => softmax
- argmax는 여전히 미분이 불가능하므로, softmax로 근사하게 됩니다.
- 이렇게 하게 되면 입력인 phi에 대해 출력인 y는 미분이 가능해지는 것입니다.
Gumbel-Max Trick을 활용한 샘플링을 충분히 활용시 원본의 확률과 같아지는 것은 아래 파이썬 코드를 통해서 확인해보실 수 있습니다.
여기서는 Gumbel Maxtrick의 결과와 Multinomial의 결과가 충분한 iteration을 거쳤을때 원본의 확률과 거의 동일하게 나옴을 알 수 있습니다.
import torch
import numpy as np
import matplotlib.pyplot as plt
def gumbel_max_sample ( logits ) :
"""
Gumbel-Max Trick을 사용하여 이산적으로 샘플링합니다.
logits: 각 클래스의 로그 확률 (또는 로짓)
"""
# 1. Uniform(0, 1)에서 노이즈 U 샘플링
U = torch.rand_like ( logits )
# 2. Gumbel 노이즈 G 계산: G = -log(-log(U))
# 안정성을 위해 아주 작은 값(1e-20)을 더합니다.
gumbel_noise = -torch.log ( -torch.log ( U + 1e-20 ) + 1e-20 )
# 3. 로짓 + Gumbel 노이즈
scores = logits + gumbel_noise
# 4. 가장 높은 점수를 가진 클래스의 인덱스 선택 (argmax)
# 이는 원-핫 벡터로 변환됩니다.
return torch.argmax ( scores , dim= -1 )
# 1. 예제 확률 분포 (3개 클래스)
# 실제 확률 분포: [0.2, 0.3, 0.5]
# 이를 로짓(log-probabilities)으로 변환합니다.
# log_softmax를 사용하면 안정적으로 로짓을 얻을 수 있습니다.
true_probs = torch.tensor ([ 0.2 , 0.3 , 0.5 ])
logits = torch.log ( true_probs )
num_samples = 100000 # 충분히 많은 샘플 수
# 결과를 저장할 리스트
multinomial_samples = []
gumbel_max_samples = []
print ( f "원본 확률 (True Probabilities): { true_probs.numpy ()} " )
print ( f "원본 로짓 (Logits): { logits.numpy ()} " )
# 2. torch.multinomial을 이용한 샘플링 (정통적인 방법)
for _ in range ( num_samples ):
# multinomial은 확률(prob_dist)을 기반으로 샘플링합니다.
sample = torch.multinomial ( true_probs , 1 ) .item ()
multinomial_samples.append ( sample )
# 3. Gumbel-Max Trick을 이용한 샘플링
for _ in range ( num_samples ):
sample = gumbel_max_sample ( logits ) .item ()
gumbel_max_samples.append ( sample )
# 샘플링 결과의 빈도 계산
multinomial_counts = np.bincount ( multinomial_samples , minlength= len ( true_probs ))
gumbel_max_counts = np.bincount ( gumbel_max_samples , minlength= len ( true_probs ))
# 빈도를 확률로 정규화
multinomial_dist = multinomial_counts / num_samples
gumbel_max_dist = gumbel_max_counts / num_samples
print ( f "\nMultinomial 샘플링 분포: { multinomial_dist } " )
print ( f "Gumbel-Max 샘플링 분포: { gumbel_max_dist } " )
# 4. 결과 시각화
fig , axes = plt.subplots ( 1 , 2 , figsize= ( 12 , 5 ), sharey= True )
fig.suptitle ( 'Comparison of Sampling Methods' , fontsize= 16 )
class_labels = [ f 'Class { i } ' for i in range ( len ( true_probs ))]
x = np.arange ( len ( true_probs )) # 클래스 인덱스를 위한 x축
# Multinomial 샘플링 결과
axes [ 0 ] .bar ( x , multinomial_dist , width= 0.4 , label= 'Sampled Frequencies' )
axes [ 0 ] .plot ( x , true_probs , 'o-' , color= 'red' , label= 'True Probabilities' )
axes [ 0 ] .set_title ( 'Multinomial Sampling' )
axes [ 0 ] .set_xlabel ( 'Class Index' )
axes [ 0 ] .set_ylabel ( 'Probability / Frequency' )
axes [ 0 ] .set_xticks ( x )
axes [ 0 ] .set_xticklabels ( class_labels )
axes [ 0 ] .set_ylim ( 0 , 0.6 )
axes [ 0 ] .legend ()
axes [ 0 ] .grid ( axis= 'y' , linestyle= '--' , alpha= 0.7 )
# Gumbel-Max 샘플링 결과
axes [ 1 ] .bar ( x , gumbel_max_dist , width= 0.4 , label= 'Sampled Frequencies' )
axes [ 1 ] .plot ( x , true_probs , 'o-' , color= 'red' , label= 'True Probabilities' )
axes [ 1 ] .set_title ( 'Gumbel-Max Sampling' )
axes [ 1 ] .set_xlabel ( 'Class Index' )
axes [ 1 ] .set_xticks ( x )
axes [ 1 ] .set_xticklabels ( class_labels )
axes [ 1 ] .set_ylim ( 0 , 0.6 )
axes [ 1 ] .legend ()
axes [ 1 ] .grid ( axis= 'y' , linestyle= '--' , alpha= 0.7 )
plt.tight_layout ( rect= [ 0 , 0.03 , 1 , 0.9 ])
plt.show ()
여기서 한 가지 의문이 들 수 있습니다.
신경망이 한 번의 epoch을 거칠 때마다 이렇게 충분히 많은 수의 샘플링이 이루어지지 않는데 그러면 샘플링에 따라 확률이 왜곡되는것이 아닌가?라고 말입니다.
- 실제 신경망 학습에서는 매 iteration 마다 batch 단위로 데이터가 처리됩니다. 이때 Gumbel-Softmax 샘플링은 배치 내의 각 데이터 포인트에 대해 한 번씩만 수행이 됩니다.
신경망 학습 과정에서의 샘플링은 다음과 같이 이루어 집니다.
a) 입력 데이터 배치 처리: 한 번에 N개의 데이터(N은 배치 크기)가 모델이 입력
b) 로짓 계산: 모델은 각 데이터 포인트에 대해 k개의 클래스에 대한 로짓을 출력. 즉 출력은 (N,k) 형태의 텐서
c) Gumbel Softmax 적용: 이 로짓 텐서에 대해 검벨 소프트 맥스 함수가 적용되며, 각 데이터 포인트(N개)마다 별 도의 검벨 노이즈가 추가되고 소프트 맥스 연산이 수행됨. 결과는 여전히 (N,k) 형태의 소프트한 확률 분포 텐서가 됨
d) 역전파: 이 소프트한 결과는 손실 계산에 사용되며, 기울기가 역전파 되어 모델의 파라미터 업데이트 됨
이때 대량의 샘플링 없이 적은 수의 샘플만으로도 모델이 잘 학습되는 이유는 다음과 같습니다.
a) 미분 가능한 근사: 검벨-소프트맥스의 핵심은 이산적인 선택을 미분 가능한 연속적인 값으로 근사한다는 점입니다. 이 연속적인 근사를 통해 매 스텝마다 비록 하나의 '샘플'에 해당하는 소프트 벡터를 얻지만, 이 벡터는 그 자체로 확률 분포에 대한 기대값 정보를 담고 있으며, 기울기가 끊기지 않고 흘러갈 수 있게 해줍니다.
2. Pytorch를 통한 Gumbel Softmax 구현
다음은 Pytorch를 통해서 Gumbel Softmax를 구현해보겠습니다.
import torch
import torch.nn.functional as F
def gumbel_softmax_sample ( logits , temperature ) :
U = torch.rand_like ( logits )
gumbel_noise = -torch.log ( -torch.log ( U + 1e-20 ))
y = logits + gumbel_noise
return F.softmax ( y / temperature , dim= -1 )
# 예제
logits = torch.tensor ([ 1.0 , 2.0 , 3.0 ], requires_grad= True )
sample = gumbel_softmax_sample ( logits , temperature= 0.5 )
print ( sample ) # soft one-hot 벡터
# 역전파 테스트
loss = sample. sum ()
loss.backward ()
print ( logits.grad )
위와 같이 1,2,3 이라는 discrete 한 텐서도 gumbel softmax를 통해 미분이 가능해짐을 알 수 있습니다.
Temperature에 따른 Gumbel Softmax의 확률에 미치는 영향은 아래와 같습니다.
import matplotlib.pyplot as plt
import numpy as np
def plot_gumbel_softmax ( logits , temperatures , num_samples = 10000 ) :
fig , axes = plt.subplots ( 1 , len ( temperatures ), figsize= ( 15 , 5 ), sharey= True )
for i , tau in enumerate ( temperatures ):
samples = [ F.gumbel_softmax ( logits , tau=tau , hard= False ) .detach () .numpy ()
for _ in range ( num_samples )]
samples = np.array ( samples )
mean_probs = samples.mean ( axis= 0 )
axes [ i ] .bar ( range ( len ( logits )), mean_probs )
axes [ i ] .set_title ( f "tau= { tau } " )
axes [ i ] .set_ylim ( 0 , 1.0 )
plt.tight_layout ()
plt.show ()
logits = torch.tensor ([ 0.5 , 1.5 , 0.1 , 2.5 ])
temperatures = [ 10 , 1 , 0.1 ]
plot_gumbel_softmax ( logits , temperatures )
Temperature가 높아질수록 부드러워지고 낮을수록 날카로워지는 모습을 보입니다.
댓글