[CS236n] 5. Normalizing Flow Models

---

💡 이 글은 『2024 YAI 봄전반기 생성모델팀』으로 진행되었으며, CS236n Fall 2023를 따라 정리했습니다.

1. Introduction

지금까지 다룬 생성모델 두 가지를 요약하면 다음과 같다.

• Autoregressive Models: $p_\theta(x) = \prod_{i=1}^n p_\theta(x_i \vert \mathbf{x}_ {\lt i})$
• Variational Autoencoders (VAEs): $p_\theta(x) = \int p_\theta (x \vert z)p(z)dz$

Autoregressive model은 tractable likelihood를 가지고 있지만, latent variable이 없어 특성(features) 학습을 명시적으로 하기 어렵다. 반면 VAE는 latent variable을 가지고 있지만, tractable likelihood를 가지고 있지 않다. Normalizing flow models에서는 이 두 모델의 장점을 결합해 tractable likelihood를 가지고 있으면서도 latent variable을 가진 모델을 제안한다.

일반적으로 생성모델에서의 모델의 확률분포 $p_\theta(x)$는 두 가지 특성을 만족하는 것이 좋다.

1. Easy-to-evaluate: 계산이 편리해야 한다. 즉, $p_\theta(x)$를 계산하는 것이 쉬워야 한다. 이는 학습 시에 유리하다. VAE는 이를 계산하기 어려운 문제가 있었다.
2. Easy-to-sample: 샘플링이 쉬워야 한다. 즉, $p_\theta(x)$에서 샘플링하는 것이 쉬워야 한다. 이는 생성 시에 유리하다.

아주 간단한 분포 (e.g. Gaussian, Uniform) 는 위 두 가지 조건을 만족시키지만, 보통 $p_\theta(x)$는 이보다 복잡한 분포이다. 따라서 normalizing flow models는 간단한 분포들을 복잡한 분포에 매핑(mapping)하는 것이 핵심 아이디어이고, 이 매핑 시 invertible transformation을 사용한다.

Invertible transformation을 사용하는 이유는 VAE에서의 문제를 생각해보면 알 수 있다. VAE에서는 $p_\theta(x) = \int p_\theta (x \vert z)p(z)dz$를 계산하기 어려웠다. 모든 $z$에 대해 $p_\theta(x \vert z)$를 계산해야 했기 때문이다. 그런데 만약 다음과 같은 관계가 있다면 어떨까?

\[x = f_\theta (z)\]

여기서 $f_\theta$는 deterministic, invertible function이다. 이러한 관계가 있다면 $x$ 하나에 대하여 $z$ 하나가 대응되기 때문에, 모든 $z$에 대해 $p_\theta(x \vert z)$를 계산할 필요가 없어진다. 그럼 이러한 대응 관계가 있을 때 $X$의 확률분포와 $Z$의 확률분포 사이의 관계는 어떻게 될까? 이를 다루는 것이 change of variables이다.

2. Change of Variables

2.1. 1D Case

$X = f(Z)$이고, 함수 $f(\cdot)$은 단조함수(monotone)이며 $Z=f^{-1}(X) = h(X)$라고 하면, $Z$의 확률분포와 $X$의 확률분포 사이의 관계는 다음과 같다.

\[p_X(x) = p_Z(h(x)) \left\vert \frac{dh(x)}{dx} \right\vert\]

이는 치환적분(change of variables)의 원리와 동일하다. 간단하게 증명해보자. 먼저 $X$의 확률분포와 $Z$의 확률분포 사이의 관계를 다음과 같이 나타낼 수 있다.

\[\begin{aligned} F_X(x) &= p \left[X \leq x \right] = p \left[f(Z) \leq x \right] \\ &= p \left[Z \leq f^{-1}(x) \right] = p \left[Z \leq h(x) \right] = F_Z(h(x)) \\ \end{aligned}\]

참고로 $F_X(x)$는 $X$의 누적분포함수(CDF)로 $p \left[X \leq x \right]$를 의미하며, 확률분포함수(PDF) $p_X(x)$는 $F_X(x)$를 미분한 것이다. 위 식을 미분하면 다음과 같다.

\[\begin{aligned} p_X(x) &= \frac{dF_X(x)}{dx} = \frac{dF_Z(h(x))}{dx} \\ &= \frac{dF_Z(h(x))}{dh(x)} \frac{dh(x)}{dx} = p_Z(h(x)) \left\vert \frac{dh(x)}{dx} \right\vert \end{aligned}\]

$h = f^{-1}$이므로 다음과 같이 쓸 수도 있다.

\[p_X(x) = p_Z(z) \left\vert \frac{df^{-1}(x)}{dx} \right\vert = p_Z(z) \left\vert \frac{df(z)}{dz} \right\vert ^ {-1}\]

참고로 간단한 분포 (prior) $p_Z(z)$와 그렇지 않은 분포 $p_X(x)$의 모양(shape)은 다르다. 단지 분포가 invertible function에 의해 mapping될 뿐이다.

2.2. General Case

$Z$를 $\left[0,1 \right]^n$에서의 uniform random vector라고 가정하자. 그리고 invertible matrix $A$에 대해 $X=AZ$라 하면, matrix $A$는 hypercube $\left[0,1 \right]^n$를 $n$차원 공간의 parallelotope으로 매핑한다. 2차원에서의 시각적인 예시는 다음과 같다.

이때 parallelotope의 volume은 $\det(A)$이다. 따라서 $X=AZ$에서, $Z$의 공간이 1이라고 하면, $X$는 $\det(A)$만큼의 공간에 매핑되어 있다. 그러므로 다음과 같은 식이 자연스럽다.

\[p_X(x) = p_Z(z) / \left\vert \det(A) \right\vert = p_Z(z) \left\vert \det(A) \right\vert ^ {-1}\]

$A^{-1}$를 사용하여 다음과 같이 나타낼 수도 있다.

\[p_X(x) = p_Z(A^{-1}x) \left\vert \det(A^{-1}) \right\vert\]

따라서 이를 일반화하자. Invertible mapping $f: \mathbb{R}^ n \mapsto \mathbb{R}^ n$에 대해 $X= f(Z)$, $Z=f^{-1}(X)$라 하면, $X$의 확률분포와 $Z$의 확률분포 사이의 관계는 다음과 같다.

\[p_X(x) = p_Z(f^{-1}(x)) \left\vert \det \left( \frac{\partial f^{-1}(x)}{\partial x} \right) \right\vert\]

VAE와는 달리, 반드시 변수 $x, z$는 continuous variable이어야 하고, 동일한 차원을 가져야 한다. 그리고 다음과 같이 나타낼 수도 있다.

\[p_X(x) = p_Z(z) \left\vert \det \left( \frac{\partial f(z)}{\partial z} \right) \right\vert ^ {-1}\]

예를 들어 2차원에서 이 식이 어떻게 적용되는지 살펴보자. $Z=(Z_1, Z_2)$라 하고, $X=(X_1, X_2)$라 하자. 그리고 invertible transformation $u: \mathbb{R}^2 \mapsto \mathbb{R}^2$에 대해 $u = (u_1, u_2)$라 하고, $X_1 = u_1(Z_1, Z_2)$, $X_2 = u_2(Z_1, Z_2)$라 하자. 그럼 inverse transformation $v = (v_1, v_2)$에 대해 $Z_1 = v_1(X_1, X_2)$, $Z_2 = v_2(X_1, X_2)$라 할 수 있다. 결과적으로 p_X(x)는 다음과 같이 두 가지 방법(각각 inverse, forward라고 부를 수 있다)으로 나타낼 수 있다.

\[\begin{aligned} p_{X_1, X_2}(x_1, x_2) \\ &= p_{Z_1, Z_2}(v_1(x_1, x_2), v_2(x_1, x_2)) \left\vert \det \begin{pmatrix} \frac{\partial v_1(x_1, x_2)}{\partial x_1} & \frac{\partial v_1(x_1, x_2)}{\partial x_2} \\ \frac{\partial v_2(x_1, x_2)}{\partial x_1} & \frac{\partial v_2(x_1, x_2)}{\partial x_2} \end{pmatrix} \right\vert \\ &= p_{Z_1, Z_2}(z_1, z_2) \left\vert \det \begin{pmatrix} \frac{\partial u_1(z_1, z_2)}{\partial z_1} & \frac{\partial u_1(z_1, z_2)}{\partial z_2} \\ \frac{\partial u_2(z_1, z_2)}{\partial z_1} & \frac{\partial u_2(z_1, z_2)}{\partial z_2} \end{pmatrix} \right\vert ^ {-1} \end{aligned}\]

3. Normalizing Flow Models

3.1. Modeling

$Z$와 $X$간의 deterministic, invertible mapping $f_\theta: \mathbb{R}^ n \mapsto \mathbb{R}^ n$에 대해 $X=f_\theta(Z)$, $Z=f_\theta^{-1}(X)$라 하자. 그럼 다음과 같이 나타낼 수 있다.

\[p_X(x; \theta) = p_Z(f_\theta^{-1}(x)) \left\vert \det \left( \frac{\partial f_\theta^{-1}(x)}{\partial x} \right) \right\vert\]

그렇다면 이름이 왜 “Normalizing Flow Model”일까?

• Normalizing: invertible transformation을 적용하면 변수가 바뀐 뒤($Z \mapsto X$)에도 확률분포의 normalization이 유지된다.
• Flow: invertible transformation은 여러 번 적용될 수 있다. 즉 다음과 같이 표현할 수 있다.

\[z_m = f_ \theta ^m \circ \cdots \circ f_ \theta ^1 (z_0) = f_ \theta ^m (f_ \theta ^{m-1} (\cdots (f_ \theta ^1 (z_0) ))) := f_ \theta (z_0)\]

즉, gaussian과 같은 간단한 분포로부터 $z_0$를 시작하여 $M$번의 invertible transformation을 적용하여 $z_M$을 얻을 수 있다. 이때 $x=z_M$은 복잡한 분포를 가지게 된다.

이를 사용하여 $p_X(x; \theta)$를 다음과 같이 나타낼 수 있다.

\[p_X(x; \theta) = p_Z(f_ \theta ^{-1}(x)) \prod _{m=1}^M \left\vert \det \left( \frac{\partial (f_ \theta ^m)^ {-1}(z_m)}{\partial z_m} \right) \right\vert\]

시각화된 예시로 planar flows (Rezende & Mohamed, 2016)를 살펴보자.

3.2. Learning and Inference

학습은 MLE(Maximum Likelihood Estimation)을 사용하여 진행할 수 있다. 즉, 다음과 같이 나타낼 수 있다. 여기서 $\mathcal{D}$는 데이터셋이다.

\[\max_ \theta \log p_X(\mathcal{D}; \theta) = \sum_ {x \in \mathcal{D}} \left[\log p_Z(f_ \theta ^{-1}(x)) + \log \left\vert \det \left( \frac{\partial f_ \theta ^{-1}(x)}{\partial x} \right) \right\vert \right]\]

$p_\theta(x)$를 구할 수 없어 ELBO를 계산하는 등 어려운 과정을 거쳤던 VAE와는 달리 exact likelihood evaluation이 가능하다. 이는 inverse transformation $x \mapsto z$를 사용했기 때문이다. 즉, 오로지 하나의 $z$에 대해서만 값을 계산하면 된다. Sampling은 forward transformation $z \mapsto x$를 사용하여 진행할 수 있다. 그리고 다시 latent representation $z$를 얻기 위해 inverse transformation $x \mapsto z$를 사용할 수 있다. Latent representation을 얻기 위해 VAE와 같이 추가적인 inference network가 필요하지 않다. 정리하면 다음과 같다.

• Learning: $p_X(x; \theta)$를 계산하기 위해 inverse transformation $x \mapsto z$를 사용한다.
• Sampling: $p_Z(z)$로부터 forward transformation $z \mapsto x$를 사용한다.
• Latent representation: inverse transformation $x \mapsto z$를 사용한다.

3.3. Desiderata for flow models

Normalizing flow models가 잘 작동하기 위해서는 다음 몇 가지를 만족하는 것이 좋다.

• Prior $p_Z(z)$가 간단한 분포이어야 한다. 간단한 분포여야 sampling과 tractable likelihood evaluation이 쉽다.
• Invertible transformation $f_\theta$가 tractable해야 한다. 즉, $f_\theta$를 계산하기 쉬워야 한다. 학습 시에는 빠르게 $x \mapsto z$를 계산할 수 있어야 하고, 생성 시에는 빠르게 $z \mapsto x$를 계산할 수 있어야 한다.
• $n \times n$ Jacobian matrices의 determinant 계산이 쉬워야 한다. 일반적인 $n \times n$ matrix의 determinant 계산은 $O(n^3)$이지만, 이를 줄이기 위해 Jacobian matrix가 determinant 계산이 편리한 형태가 되도록 설계할 수 있다.

4. Application: Diagonal Jacobians

여기서는 Jacobian matrix의 determinant 계산이 쉬운 형태를 가지는 몇 가지 normalizing flow models를 살펴보자. 일반적인 $n \times n$ Jacobian matrix의 determinant 계산은 $O(n^3)$이다. 즉 다음과 같은 normalizing flow model이 있을 때,

\[\mathbf{x} = \left(x_1, \cdots, x_n \right) = f(\mathbf{z}) = \left(f_1(\mathbf{z}), \cdots, f_n(\mathbf{z}) \right)\]

Jacobian matrix는 다음과 같이 나타낼 수 있다.

\[\mathbf{J} = \begin{pmatrix} \frac{\partial f_1}{\partial z_1} & \cdots & \frac{\partial f_1}{\partial z_n} \\ \vdots & \ddots & \vdots \\ \frac{\partial f_n}{\partial z_1} & \cdots & \frac{\partial f_n}{\partial z_n} \end{pmatrix}\]

그런데 만약 $x_i = f_i(\mathbf{z})$가 $\mathbf{z}_ {\leq i}$에만 의존한다면, Jacobian matrix는 다음과 같이 나타낼 수 있다.

\[\mathbf{J} = \begin{pmatrix} \frac{\partial f_1}{\partial z_1} & \cdots & 0 \\ \vdots & \ddots & \vdots \\ \frac{\partial f_n}{\partial z_1} & \cdots & \frac{\partial f_n}{\partial z_n} \end{pmatrix}\]

즉 Jacobian matrix가 lower triangular matrix가 되므로, determinant는 diagonal elements의 곱으로 나타낼 수 있다. 이는 $O(n)$의 시간복잡도를 가진다. 이러한 normalizing flow models를 살펴보자.

4.1. NICE (Nonlinear Independent Components Estimation)

NICE(Nonlinear Independent Components Estimation)에서는 additive coupling layer와 rescaling layer를 사용해 Jacobian matrix의 determinant를 1로 고정했다.

먼저 $\mathbf{z}$를 $1 \leq d \lt n$인 $d$에 대해 $\mathbf{z}_ {1:d}$와 $\mathbf{z}_ {d+1:n}$로 나눈다. 이때 forward mapping은 다음과 같이 나타내었다.

• $\mathbf{x}_ {1:d} = \mathbf{z}_ {1:d}$ (identity transformation)
• $\mathbf{x}_ {d+1:n} = \mathbf{z}_ {d+1:n} + m_ \theta (\mathbf{z}_ {1:d})$ ($m_ \theta$: neural network)

따라서 inverse mapping은 다음과 같이 나타낼 수 있다.

• $\mathbf{z}_ {1:d} = \mathbf{x}_ {1:d}$ (identity transformation)
• $\mathbf{z}_ {d+1:n} = \mathbf{x}_ {d+1:n} - m_ \theta (\mathbf{x}_ {1:d})$

이와 같이 forward mapping을 정의하면 Jacobian matrix는 다음과 같이 나타낼 수 있다.

\[\mathbf{J} = \begin{pmatrix} \mathbf{I}_ d & 0 \\ \frac{\partial \mathbf{x}_ {d+1:n}}{\partial \mathbf{z}_ {1:d}} & \mathbf{I}_ {n-d} \end{pmatrix}\]

따라서 $\det(\mathbf{J}) = 1$이다. 이를 “volume preserving transformation”이라고 한다.

이러한 additive coupling layer를 여러 번 쌓아서 deep normalizing flow를 만들 수 있고, 마지막 layer에서는 scaling을 위해 rescaling layer를 사용한다.

• Forward mapping $\mathbf{z} \mapsto \mathbf{x}$: $\mathbf{x}_ i = s_ i \mathbf{z}_ i$
• Inverse mapping $\mathbf{x} \mapsto \mathbf{z}$: $\mathbf{z}_ i = \frac{1}{s_ i} \mathbf{x}_ i$

이 layer의 Jacobian과 determinant는 다음과 같다.

\[\mathbf{J} = \begin{pmatrix} s_1 & 0 & \cdots & 0 \\ 0 & s_2 & \cdots & 0 \\ \vdots & \vdots & \ddots & \vdots \\ 0 & 0 & \cdots & s_n \end{pmatrix}, \quad \det(\mathbf{J}) = \prod_ {i=1}^n s_ i\]

이를 통해 NICE를 훈련한 결과는 다음과 같다.

4.2. RealNVP (Real-valued Non-Volume Preserving)

RealNVP는 NICE로부터 더 발전된 모델로, $X$와 $Z$ 간의 관계를 더 복잡한 수식으로 모델링하였다. Forward mapping $\mathbf{z} \mapsto \mathbf{x}$는 다음과 같이 나타낼 수 있다.

• $\mathbf{x}_ {1:d} = \mathbf{z}_ {1:d}$ (identity transformation)
• $\mathbf{x}_ {d+1:n} = \mathbf{z}_ {d+1:n} \odot \exp(\alpha_ \theta (\mathbf{z}_ {1:d})) + \mu_ \theta (\mathbf{z}_ {1:d})$ ($\mu_\theta, \alpha_\theta$: neural networks, $\odot$: element-wise product)

따라서 inverse mapping $\mathbf{x} \mapsto \mathbf{z}$는 다음과 같이 나타낼 수 있다.

• $\mathbf{z}_ {1:d} = \mathbf{x}_ {1:d}$ (identity transformation)
• $\mathbf{z}_ {d+1:n} = (\mathbf{x}_ {d+1:n} - \mu_ \theta (\mathbf{x}_ {1:d})) \odot \exp(-\alpha_ \theta (\mathbf{x}_ {1:d}))$

이때 Jacobian과 determinant는 다음과 같다.

\[\mathbf{J} = \begin{pmatrix} \mathbf{I}_ d & 0 \\ \frac{\partial \mathbf{x}_ {d+1:n}}{\partial \mathbf{z}_ {1:d}} & \text{diag}(\exp(\alpha_ \theta (\mathbf{z}_ {1:d}))) \end{pmatrix} \\ \det(\mathbf{J}) = \prod_ {i=d+1}^n \exp(\alpha_ \theta (\mathbf{z}_ {1:d}))_ i = \exp \left( \sum_ {i=d+1}^n \alpha_ \theta (\mathbf{z}_ {1:d})_ i \right)\]

Determinant가 1보다 크므로, 이를 “non-volume preserving transformation”이라고 한다. 이러한 RealNVP를 훈련한 결과는 다음과 같다.

Normalizing flow models는 latent variable을 자유롭게 조절할 수 있으므로, 4개의 latent variable $z^{(1)}, z^{(2)}, z^{(3)}, z^{(4)}$에 대해 interpolation을 수행할 수도 있다.

\[\mathbf{z} = \cos \phi (z^{(1)} \cos \phi^\prime + z^{(2)} \sin \phi^\prime) + \sin \phi (z^{(3)} \cos \phi^\prime + z^{(4)} \sin \phi^\prime)\]

이를 시각화하면 다음과 같다.

이처럼 normalizing flow models는 latent space interpolation을 수행하는 태스크에 유리하다.

5. Application: Autoregressive Flows

2장에서 배운 autoregressive models와 normalizing flow models를 결합한 autoregressive flows를 살펴보자. Autoregressive models는 다음과 같이 나타낼 수 있다.

\[p(\mathbf{x}) = \prod_ {i=1}^n p(x_ i \vert \mathbf{x}_ {\lt i})\]

이때 다음과 같이 가정하자.

\[p(x_ i \vert \mathbf{x}_ {\lt i}) = \mathcal{N}(\mu_ i(\mathbf{x}_ {\lt i}), \exp(\alpha_i(\mathbf{x}_ {\lt i}))^ 2)\]

여기서 $\mu_ i$, $\alpha_ i$는 neural networks이고, $\mu_ 1$, $\alpha_ 1$은 상수이다. 이때 각각의 $x_i$가 latent variable $z_i \sim \mathcal{N} (0, 1)$에 대응된다고 가정하고 다음과 같이 모델링하자.

\[x_ i = \exp (\alpha_ i(\mathbf{x}_ {\lt i})) z_ i + \mu_ i(\mathbf{x}_ {\lt i})\]

이는 Real-NVP에서 보았던 식과 거의 동일하다. 우리는 standard gaussian $(z_1, z_2, \cdots, z_n)$으로부터 invertible transformation을 통해 $(x_1, x_2, \cdots, x_n)$을 얻을 수 있다. 샘플링 과정을 간략하게 요약하면 다음과 같다.

1. $z_i \sim \mathcal{N} (0, 1)$을 샘플링한다.
2. $x_1 = \exp (\alpha_ 1) z_1 + \mu_ 1$로 $x_1$을 얻고, $\mu_ 2(x_ 1)$, $\alpha_ 2(x_ 1)$을 계산한다.
3. $x_2 = \exp (\alpha_ 2) z_2 + \mu_ 2$로 $x_2$을 얻고, $\mu_ 3(x_ 1, x_ 2)$, $\alpha_ 3(x_ 1, x_ 2)$을 계산한다.
4. $x_n$까지 반복한다.

5.1. Masked Autoregressive Flow (MAF)

위 과정을 이용한 모델이 Masked Autoregressive Flow (MAF)이다. MAF의 forward mapping $\mathbf{z} \mapsto \mathbf{x}$는 다음과 같이 나타낼 수 있다.

• $z_ i \sim \mathcal{N} (0, 1)$을 샘플링한다.
• $x_ 1 = \exp (\alpha_ 1) z_ 1 + \mu_ 1$로부터 $\mu_ 2(x_ 1)$, $\alpha_ 2(x_ 1)$을 계산한다.
• $x_ 2 = \exp (\alpha_ 2) z_ 2 + \mu_ 2$로부터 $\mu_ 3(x_ 1, x_ 2)$, $\alpha_ 3(x_ 1, x_ 2)$을 계산한다.
• 반복한다.

그러나 이러한 방식은 parallelization이 어렵다. 즉, $\mathbf{x}_ i$를 계산하기 위해 $\mathbf{x}_ {\lt i}$를 하나씩 계산해야 한다. Autoregressive model들의 단점과 동일하게, 샘플링에 $O(n)$의 시간이 걸린다.

한편 inverse mapping $\mathbf{x} \mapsto \mathbf{z}$는 다음과 같이 나타낼 수 있다.

• MADE와 같은 모델을 사용하면, 이미 모든 $x_i$가 주어져 있으므로 모든 $\mu_ i$, $\alpha_ i$를 한 번에 계산할 수 있다. 이는 parallelization이 가능하다는 것을 의미한다.
• $z_1 = (x_1 - \mu_ 1) / \exp (\alpha_ 1)$를 계산한다.
• $z_2 = (x_2 - \mu_ 2) / \exp (\alpha_ 2)$를 계산한다.
• 반복한다.

따라서 MAF의 장단점은 명확하다.

• 장점: likelihood evaluation이 쉽다. 즉, inverse mapping $\mathbf{x} \mapsto \mathbf{z}$ 과정의 parallelization이 가능하다.
• 단점: 샘플링이 어렵다. 즉, forward mapping $\mathbf{z} \mapsto \mathbf{x}$ 과정을 sequential하게 수행해야 한다.

5.2. Inverse Autoregressive Flow (IAF)

IAF는 이를 거꾸로 뒤집은 모델이다. 즉, $\mu_i$, $\alpha_i$를 $\mathbf{x}$의 함수로 보는 것이 아니라 $\mathbf{z}$의 함수로 보는 것이다. 이를 “inverse autoregressive flow”라고 한다. IAF의 forward mapping $\mathbf{z} \mapsto \mathbf{x}$는 다음과 같이 나타낼 수 있다.

• $z_ i \sim \mathcal{N} (0, 1)$을 샘플링한다.
• 이미 모든 $z_ i$가 주어져 있으므로 모든 $\mu_ i$, $\alpha_ i$를 한 번에 계산할 수 있다.
• $x_ 1 = \exp (\alpha_ 1) z_ 1 + \mu_ 1$을 계산한다.
• $x_ 2 = \exp (\alpha_ 2) z_ 2 + \mu_ 2$을 계산한다.
• 반복한다.

Inverse mapping $\mathbf{x} \mapsto \mathbf{z}$는 다음과 같이 나타낼 수 있다.

• $z_1 = (x_1 - \mu_ 1) / \exp (\alpha_ 1)$로부터 $\mu_ 2(x_1)$, $\alpha_ 2(x_1)$을 계산한다.
• $z_2 = (x_2 - \mu_ 2) / \exp (\alpha_ 2)$로부터 $\mu_ 3(x_1, x_2)$, $\alpha_ 3(x_1, x_2)$을 계산한다.
• 반복한다.

그러므로 IAF의 장단점은 다음과 같다.

• 장점: 샘플링이 쉽다. 즉, forward mapping $\mathbf{z} \mapsto \mathbf{x}$ 과정의 parallelization이 가능하다.
• 단점: likelihood evaluation이 어렵다. 즉, inverse mapping $\mathbf{x} \mapsto \mathbf{z}$ 과정을 sequential하게 수행해야 한다. (단, generated point $\mathbf{z}$에 대해 likelihood evaluation은 쉽다.)

5.3. Parallel Wavenet

MAF의 장점은 training, 즉 likelihood evaluation이 빠르다는 것이다. 반면 IAF의 장점은 generation, 즉 sampling이 빠르다는 것이다. 이러한 장점을 결합한 모델이 Parallel Wavenet이다. Parallel Wavenet은 teacher-and-student model을 사용해 두 단계의 학습을 수행한다. 이때 teacher model은 MAF, student model은 IAF를 사용한다.

1. Teacher model: MAF를 사용하여 빠르게 학습한다.
2. Student model: teacher model로부터 학습한 후, IAF를 사용하여 빠르게 샘플링한다.

가장 큰 문제는 teacher model로부터 student model을 학습시키는 과정이다. 여기서의 아이디어는 IAF도 generative point $\mathbf{z}$에 대해서는 likelihood evaluation이 쉽다는 것이다.

Student model에게 teacher model의 학습 결과를 전달(distillation)하기 위해서는 probability density distillation을 사용한다. 즉 다음 KL-divergence를 최소화하는 것을 목표로 학습시킨다.

\[D_{KL}(s,t) = \mathbb{E}_ {x \sim s} \left[ \log s(x) - \log t(x) \right]\]

여기서 각각 $s(\cdot)$, $t(\cdot)$는 student model, teacher model의 확률밀도함수이다. 위 KL-divergence를 계산하기 위해서는 아래 3가지 과정이 필요하다.

1. Student model로부터 $\mathbf{z}$를 샘플링하여 샘플 $\mathbf{x}$를 생성한다.
2. Student model로부터 density $s(\mathbf{x})$를 계산한다.
3. Teacher model로부터 density $t(\mathbf{x})$를 계산한다.

1, 2 과정은 student model (IAF)의 forward mapping $\mathbf{z} \mapsto \mathbf{x}$이고, 3 과정은 teacher model (MAF)의 inverse mapping $\mathbf{x} \mapsto \mathbf{z}$이다. 이러한 과정을 통해 Parallel Wavenet은 MAF의 장점과 IAF의 장점을 결합하여 빠르게 학습하고 빠르게 샘플링할 수 있다. 이로부터 기존 autoregressive model이었던 Wavenet보다 샘플링의 효율이 약 1000배 향상되었다고 한다.

6. Others

6.1. MintNet (Song et al., 2019)

MintNet은 masked convolution을 사용하여 CNN 기반의 normalizing flow model을 제안했다. 일반적인 CNN은 invertible하지 않고, Jacobian matrix의 determinant를 계산하기 어렵다. 따라서 MintNet은 PixelCNN과 같은 masked convolution을 사용하여 invertible한 CNN을 만들었다. 따라서 Jacobian도 triangular matrix가 된다. 만약 Jacobian matrix의 diagonal elements가 모두 양수라면, CNN은 invertible transformation이 된다. 이를 시각화하면 다음과 같다.

이를 활용하여 얻은 생성 결과는 다음과 같다.

6.2. Gaussianization Flow (Meng et al., 2020)

Gaussianization flow model에서는 $p_{data}$ 분포를 가우시안 분포로 바꾸는 gaussianization을 flow model로 설명한다. 먼저 $X = f_\theta(Z)$는 gaussian prior $Z \sim \mathcal{N}(0, I) = p_Z$를 가진 flow model이라고 하고, $\tilde{X} \sim p_{data}$라 하자. 이때 KL-divergence를 다음과 같이 나타낼 수 있다.

\[\begin{aligned} D_{KL}(p_{data} \Vert p_\theta(x)) &= D_{KL}(p_{\tilde{X}} \Vert p_X) \\ &= D_{KL}(p_{f_\theta^{-1}(\tilde{X})} \Vert p_{f_\theta^{-1}(X)}) \\ &= D_{KL}(p_{f_\theta^{-1}(\tilde{X})} \Vert p_Z) \\ \end{aligned}\]

즉 $f_\theta^{-1}$을 통해 data samples $\tilde{X}$는 gaussian distribution이 되어야 한다. 그렇다면 적절한 $f_\theta^{-1}$을 어떻게 찾을 것인가? 여기서는 inverse CDF trick을 사용한다.

CDF $F_{data}$가 알려져 있다면, $p_{data}$에서 $\tilde{X}$를 샘플링하는 것은 $\tilde{X} = F_{data}^{-1}(U)$로부터 $U \sim \mathcal{U}(0, 1)$을 샘플링하는 것과 동일하다. 이를 다시 해석하면 $U = F_{data}(\tilde{X})$는 uniform distribution이라는 것이다. 다시 여기에 inverse gaussian CDF $\Phi^{-1}$을 적용하면 gaussian distribution을 얻을 수 있다. 즉 $\Phi^{-1}(U) = \Phi^{-1}(F_{data}(\tilde{X}))$는 가우시안 분포가 된다. 따라서, invertible transformation $\Phi^{-1} \circ F_{data}$는 데이터를 가우시안화(gaussianize)한다.

실제 gaussianization flow model의 과정은 다음과 같다.

1. 모든 차원에 대하여 gaussianization을 진행하기는 어려우므로, dimension-wise gaussianization을 진행한다.

   

2. Rotation matrix를 적용한다.

   

3. 1과 2를 반복하여 데이터를 normal distribution으로 바꾼다.

   

이러한 gaussianization flow model을 통해 실험한 결과는 다음과 같다. 참고로 Glow는 GF(Gaussianization Flow) 이전의 모델이다.

7. Summary

[CS236n] 5. Normalizing Flow Models에서는 normalizing flow models에 대해 살펴보았다. Normalizing flow models는 deterministic, invertible mapping을 사용하여 latent variable과 observed variable 간의 관계를 모델링한다. 이를 통해 exact likelihood evaluation이 가능하다. 그리고 그 예시로 NICE, RealNVP, MAF, IAF, Parallel Wavenet, MintNet, Gaussianization Flow를 살펴보았다.

CS236n 카테고리 내 다른 글 보러가기