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

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

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

$$ \theta^* = \underset{\theta}{\text{argmax}} \; E_{p(x)}[f_{\theta}(x)]. \tag{1} $$

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

$$\begin{align*}\nabla_{\theta} E_{p(x)}[f_{\theta}(x)] &= \nabla_{\theta} \left[ \int_x p(x) f_{\theta}(x) dx \right]\\ &= \int_x p(x) \left[\nabla_{\theta} f_{\theta}(x) \right] dx\\ &= E_{p(x)} \left[ \nabla_{\theta} f_{\theta}(x) \right]. \tag{2}\end{align*}$$

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

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

$$\begin{align*}\nabla_{\theta} E_{p_{\theta}(x)}[f_{\theta}(x)] &= \nabla_{\theta} \left[\int_x p_{\theta}(x) f_{\theta}(x) dx \right]\\ &= \int_x \nabla_{\theta} \left[p_{\theta}(x)f_{\theta}(x)\right] dx\\ &= \int_x f_{\theta}(x)\nabla_{\theta}p_{\theta}(x) dx + \int_x p_{\theta}(x)\nabla f_{\theta}(x) dx\\ &= \int_x f_{\theta}(x)\nabla_{\theta}p_{\theta}(x) dx + E_{p_{\theta}(x)} [\nabla_{\theta}f_{\theta}(x)]. \tag{3}\end{align*}$$

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

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

 

2. Reparameterization Trick

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

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

$$\begin{equation} \textbf{z} = \mu + \epsilon \times \sigma \tag{4} \end{equation}$$

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

$$\begin{equation} \textbf{z} = g_{\theta}(\textbf{x}, \epsilon) \tag{5} \end{equation}$$

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

$$\begin{align*} E_{p_{\theta}(\textbf{z})}\left[f_{\theta}(\textbf{z}) \right] = E_{p(\epsilon)} \left[f_{\theta}(g_{\theta}(\textbf{x}, \epsilon))\right] \tag{6} \end{align*}$$

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

$$\nabla_\theta E_{p(\epsilon)} [f_{\theta}(g_\theta(\textbf{x}, \epsilon))] = E_{p(\epsilon)} [\nabla_\theta (f_\theta (g_\theta (\textbf{x}, \epsilon)))] \tag{7}$$

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

$$E_{p(\epsilon)} [\nabla_\theta (f_\theta (g_\theta (\textbf{x}, \epsilon)))] \approx \frac{1}{N} \sum_{n=1}^N \nabla_\theta f_\theta(g_\theta(\textbf{x}_n, \epsilon_n)). \tag{8}$$

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