Gumbel-Softmax 리뷰

본 논문은 2개를 확실히 다 읽어야 이해가 가능한데, [1] [2] 이다.

 

Overview

 

forward, backward

Gumbel-Softmax는 간단하게 정리하면 아래와 같다.

1) sampling을 하고 싶은데, neural network에서 backpropagation시에 불가능하다. 이를 해결하기 위해 Gumbel-Max Trick을 사용하여 backpropagation이 흐르도록 해주자

2) argmax를 사용하였더니 backpropagation이 흐르지 않는다. 이를 어떻게 해결할까? Softmax를 취하여 해결함과 동시에, continuous하게 relaxation을 하기 위해 temperature $\tau$를 쓰자

 

Method

 

Gumbel-Max Trick & Gumbel-Softmax Trick

1. Gumbel-Max Trick

1) 이전 VQ-VAE에서는 online K-means를 사용하여 codebook을 learning 시켜서 reconstruction하였고, 이 때 reshape된 input이 (32,32,64)였다. 여기에서 이어서 함

- Encoder가 logit을 뱉으면, categorical 한 (32,32,64)이 됨

- 이 logit을 다음으로 넘어가려면 one-hot으로 sampling 해야됨

- Categorical한 distribution이 나오면, 이 중 가장 큰 값이 one-hot vector가 되겠지만, one-hot vector가 항상 deterministic하다는 단점이 있음

 

2) Sampling 사용 이유

- 왜 sampling을 하는가?

- 벡터를 quntization만 하게 되면 noise가 필요 없고 샘플링도 할 필요 없지만, enc-dec가 배우는 도중에 sampling을 하게 되면 확률적으로 노이즈에 의해 다른 vector들이 decoder 입장에서 학습이 잘 될 수 있음

- 이로 인해 loss에 의해 더 많이 뽑히게 되면서 빠르게 수렴할 수 있다는 장점이 있음.

 

3) Sampling 사용이 아닌 deterministic한 one-hot vector만 사용하게 될 경우 

- Sampling 기법이 없으면 매번 큰 애들만 하게 되니까, 그 representation이 만약 정답이 아니라면, optimal solution을 못 구할 수 있다는 우려 및 단점이 있음

- 즉, sampling을 하면 이를 해결할 수 있음. 따라서, noise를 첨가한 뒤, argmax를 취하면 확률적으로 다른 애가 one-hot vector가 될 수 있음

 

4) Method

- 첫번째로, Sampling 된 것을 one-hot vecto로 사용하고 싶은데, sampling 을 하면 gradient가 흐르지 않음

- 이것을 흐르게 하는게 gumbel-max tirick 임

- Motivation은 Logit에 uniform한 noise를 더하고, 제일 큰 놈을 sampling 하면, 이것이 sampling 한것이랑 같지 않은가?

- 이후 argmax로 뽑으면 되는것이 아닌가? 라는 아이디어임.
- 여기에 Gumbel noise를 넣어야 categorical에서 sampling한 것이랑 Gumbel noise를 합한 뒤 max 인 것을 뽑은 것이 같다고 증명되어 있음

Gumbel Noise

- 이 내용을 사용하고 싶은 것이 이 두 논문임 [1-2]

2. Gumbel-Softmax Trick

- 위의 Gumbel-max trick으로 sampling의 문제인 미분불가능 문제를 해결하면, argmax operation 문제가 남아 있음

- Argmax는 neural networks에서 미분 불가능. 즉 역전파가 흐를 수 없음

- Argmax 문제를 제거하기 위해, Softmax 트릭을 이용하여 제거함

- 즉, Softmax가 argmax를 제거하기 위해 나온 것이므로 (Softmax: 미분 가능, Argmax: 미분 불가능)

 

정리하면, 

- Argmax는 backpropagation이 안되고, Softmax는 이러한 argmax를 제거함과 동시에 미분이 가능하도록 함

- 논문에 나와있는 $\tau$는 temperature임

- 즉, Continuous 하게 relaxation을 하기 위해 temperature가 들어옴

Temperature $\tau$의 역할

- 위의 그림 처럼, $\tau$값이 너무 높으면 무작위 sampling을 하는 역할을 하게 됨

- 반면에, 너무 낮게 되면 One-hot vector처럼 categorical한 값을 갖게 됨

- 논문에서는 처음에 낮은 값에서 시작하여 maximum 2.0값으로 annealing하여 종료함

- 논문에서는 처음에 높은 값에서 시작하여 (maximum 2.0) 낮은 값(0 값 아님)으로 annealing 하여 종료 (수정)

- 이를 통해 아래의 식처럼 Sampling을 자유 자제로 함과 동시에 codebook에서 online K-means로 낮은 vector의 개수(max: 64) 밖에 가지지 못했다면, Gumbel-Softmax Trick을 활용하여 더 많은 값을 사용할 수 있다고 밝힘

 

- 해당 방법을 통해 quantization 기법을 활용함

 

 

[1] Jang, Eric, Shixiang Gu, and Ben Poole. "Categorical reparameterization with gumbel-softmax." arXiv preprint arXiv:1611.01144 (2016).

[2] Maddison, Chris J., Andriy Mnih, and Yee Whye Teh. "The concrete distribution: A continuous relaxation of discrete random variables." arXiv preprint arXiv:1611.00712 (2016).