Reparameterization Trick에 대한 수학적 이해와 기댓값의 미분가능성

1. 기댓값 (Expectation)의 미분가능성

많은 머신러닝 문제에서 우리는 식 (1)과 같이 모델 매개변수 $\theta$에 대해 어떠한 기댓값을 최대화하고자 한다.

$$\begin{array}{}\text{(1)}& {\theta }^{\ast }=\underset{\theta }{\text{argmax}}{E}_{p(x)}[f_{\theta }(x)].\end{array}$$

식 (1)을 풀기 위해 $\theta$에 대해 $E_{p(x)}[f_{\theta }(x)]$를 미분하면 다음과 같다.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{E}_{p(x)}[f_{\theta }(x)]& ={\mathrm{\nabla }}_{\theta }[{\int }_{x}p(x)f_{\theta }(x)dx]\\ & ={\int }_{x}p(x)[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)]dx\\ \text{(2)}& & =E_{p(x)}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)].\end{array}$$

따라서 $f_{\theta }(x)$가 인공신경망 등과 같이 $\theta$에 대해 미분가능한 함수라면, 우리는 식 (1)의 기댓값에 대한 gradient를 식 (2)와 같이 gradient의 기댓값으로 계산할 수 있다.

식 (2)에서는 데이터의 참된 분포인 $p(x)$를 알고 있다고 가정했다. 그러나 대부분의 머신러닝 문제에서는 $p(x)$를 정확히 알 수 없으며, 이를 데이터로부터 추정하기 위해 $p(x)\approx p_{\theta }(x)$와 같이 매개변수화한다. 매개변수화된 데이터 분포 $p_{\theta }(x)$를 가정할 때, 식 (2)의 과정은  아래와 같이 변형된다.

$$\begin{array}{rl}{\mathrm{\nabla }}_{\theta }{E}_{p_{\theta }(x)}[f_{\theta }(x)]& ={\mathrm{\nabla }}_{\theta }[{\int }_x{p}_{\theta }(x)f_{\theta }(x)dx]\\ & ={\int }_x{\mathrm{\nabla }}_{\theta }[p_{\theta }(x)f_{\theta }(x)]dx\\ & ={\int }_x{f}_{\theta }(x){\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)dx+{\int }_x{p}_{\theta }(x)\mathrm{\nabla }{f}_{\theta }(x)dx\\ \text{(3)}& & ={\int }_x{f}_{\theta }(x){\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)dx+E_{p_{\theta }(x)}[{\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)].\end{array}$$

식 (3)의 두번째 항은 식 (2)와 같이 $p_{\theta }(x)$에 대한 ${\mathrm{\nabla }}_{\theta }{f}_{\theta }(x)$의 기댓값으로 정의된다.

그러나 식 (3)의 첫번째 항은 $p_{\theta }(x)$에 대한 어떠한 함수의 기댓값으로 표현되지 않는다. 따라서 $x$에 대한 샘플을 추출하여 몬테카를로 방법 (Monte Carlo method)을 기반으로 식 (3)의 첫번째 항을 근사할 수 없다. 만약 ${\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)$에 대한 analytic solution을 구할 수 있다면 몬테카를로 방법 없이도 계산이 가능하겠지만 대부분의 경우에 우리는 ${\mathrm{\nabla }}_{\theta }{p}_{\theta }(x)$에 대한 analytic solution을 찾을 수 없다.

 

2. Reparameterization Trick

많은 머신러닝 방법론에서 입력 데이터 $\mathbf{\text{x}}$의 숨겨진 정보를 나타내는 잠재 표현 (latent representation) $\mathbf{\text{z}}$를 가정하기 때문에 기댓값 $E_{p_{\theta }(\mathbf{\text{z}})}[f_{\theta }(\mathbf{\text{z}})]$에 대해 모델 매개변수 $\theta$를 최적화하는 문제가 자주 나타난다. 그러나 식 (3)에서 보이듯이 $p_{\theta }(\mathbf{\text{z}})$를 따르는 확률변수의 어떠한 기댓값 $E_{p_{\theta }(\mathbf{\text{z}})}[\cdot ]$에 대해 $\theta$의 gradient를 계산하는 것이 항상 가능한 것은 아니다. Reparameterization trick의 핵심은 확률 변수 $\mathbf{\text{z}}$를 deterministic 부분과 stochastic 부분으로 나누는 것이다.

만약 $\mathbf{\text{z}}$가 정규분포 $\mathcal{N}(\mu ,{\sigma }^2)$을 따른다면, $\mathbf{\text{z}}$는 $\mu$를 중심으로 $\sigma$ 정도의 변동성을 가진 값일 것이다. 따라서 $\mathcal{N}(\mu ,{\sigma }^2)$에서 $\mathbf{\text{z}}$를 바로 샘플링하는 것이 아니라, 먼저 표준정규분포인 $\mathcal{N}(0,1)$에서 $ϵ$를 샘플링하고 아래의 식 (4)와 같이 $\mathbf{\text{z}}$를 재구성할 수 있다.

$$\begin{array}{}\text{(4)}& \mathbf{\text{z}}=\mu +ϵ\times \sigma \end{array}$$

이러한 reparameterization trick을 일반화하면 아래의 식 (5)와 같다.

$$\begin{array}{}\text{(5)}& \mathbf{\text{z}}=g_{\theta }(\mathbf{\text{x}},ϵ)\end{array}$$

식 (5)에서 $ϵ$은 우리가 이미 알고있는 어떠한 확률분포를 따르는 확률변수이며, $\mathbf{\text{z}}$는 입력 데이터 $\mathbf{\text{x}}$와 난수 $ϵ$을 기반으로 매개변수화된 함수 $g_{\theta }$에 의해 $\mathbf{\text{z}}=g_{\theta }(\mathbf{\text{x}},ϵ)$으로 계산된다. 이러한 reparameterization trick을 적용하면 우리가 최대화하고자 하는 기댓값은 아래와 같이 변형된다.

$$\begin{array}{}\text{(6)}& E_{p_{\theta }(\mathbf{\text{z}})}[f_{\theta }(\mathbf{\text{z}})]=E_{p(ϵ)}[f_{\theta }(g_{\theta }(\mathbf{\text{x}},ϵ))]\end{array}$$

Reparameterization trick에서 가장 중요한 점은 우리가 이미 $ϵ$의 확률분포 $p(ϵ)$을 알고있기 때문에 매개변수화된 함수로 근사할 필요가 없다는 것이다. 따라서 $p(ϵ)$은 우리가 최적화해야하는 매개변수 $\theta$와 독립이며, 식 (6)에 대한 $\theta$의 gradient는 아래와 같이 계산된다.

$$\begin{array}{}\text{(7)}& {\mathrm{\nabla }}_{\theta }{E}_{p(ϵ)}[f_{\theta }(g_{\theta }(\mathbf{\text{x}},ϵ))]=E_{p(ϵ)}[{\mathrm{\nabla }}_{\theta }(f_{\theta }(g_{\theta }(\mathbf{\text{x}},ϵ)))]\end{array}$$

우리는 이미 $p(ϵ)$를 알고있기 때문에 $p(ϵ)$을 따르는 $ϵ$을 샘플링할 수 있으며, $N$을 샘플의 수라고 할 때 몬테카를로 방법에 의해 식 (7)은 아래와 같이 근사된다.

$$\begin{array}{}\text{(8)}& E_{p(ϵ)}[{\mathrm{\nabla }}_{\theta }(f_{\theta }(g_{\theta }(\mathbf{\text{x}},ϵ)))]\approx \frac{1}{N}\sum _{n=1}^N{\mathrm{\nabla }}_{\theta }{f}_{\theta }(g_{\theta }({\mathbf{\text{x}}}_n,{ϵ}_n)).\end{array}$$

일반적으로 머신러닝에서는 $f_{\theta }$와 $g_{\theta }$ 모두 미분가능한 인공신경망으로 정의되기 때문에 경사하강법 (gradient descent method)을 이용하면 (8)에 대한 최적의 매개변수 ${\theta }^{\ast }$ (정확히는 local optimum solution)를 찾을 수 있다.