Coupling Layer

Ref : NICE, RealNVP

General coupling layer

$x \in \mathcal{X}$, $I_{1}, I_{2}$은 $[1, D]$의 분할로 $d=\vert I_{1} \vert$이고, $m$은 $\mathbb{R}^{d}$에서 정의된 함수이다.

$$
\begin{align}
y_{I_{1}}&=x_{I_{1}}\\
y_{I_{2}}&=g(x_{I_{2}};m(x_{I_{1}}))
\end{align}
$$

• $g:\mathbb{R}^{D-d}\times m(\mathbb{R}^{d})\rightarrow \mathbb{R}^{D-d}$ : coupling 법칙이며, 첫번째 인수에 대해 두번째 인수를 고려한 상태에서 역함수가 존재하는 변환

$I_{1}=[1,d], I_{2}=[d, D]$를 고려한다면, Jacobian은 아래와 같이 표현할 수 있다.
$$
\frac{\partial y}{\partial x}= \begin{bmatrix} I_{d} & 0 \\ \frac{\partial y_{I_{2}}}{\partial x_{I_{1}}} & \frac{\partial y_{I_{2}}}{\partial x_{I_{2}}} \end{bmatrix}
$$

• $I_{d}$는 크기 $d$의 identity matrix
• $\det \left(\frac{\partial y}{\partial x}\right)=\det \left(\frac{\partial y_{I_{2}}}{\partial x_{I_2}}\right)$

또한, 아래와 같이 역함수를 사용할 수 있다.

$$
\begin{align}
x_{I_{1}}&=y_{I_{1}} \\
x_{I_{2}}&=g^{-1}(y_{I_{2}};m(y_{I_{1}}))
\end{align}
$$
이러한 변환을 coupling function $m$을 사용하는 coupling layer라고 한다.

Affine Coupling layer

Coupling layer에서 함수 $f$를 신중하게 설계함으로써 계산가능하면서도 매우 유연한 bijective model을 학습할 수 있다. 이를 위해서는 Jacobian 행렬식 계산이 매우 중요하다.

삼각행렬의 행렬식을 삼각행렬의 대각선 원소들의 곱으로 효율적으로 계산할 수 있다.

아래와 같은 간단한 bijection을 affine coupling layer라고 한다.
$$
\begin{align}
y_{1:d}&=x_{1:d}\\
y_{d+1:D}&=x_{d+1:D}\odot \exp(s(x_{1:d}))+t(x_{1:d})
\end{align}
$$

• s : scale
• t : translation
• $R^{d}\mapsto R^{D-d}$ : s, t function
• $\odot$ : 원소별 곱셈

위 변환의 Jacobian은 아래와 같다.
$$
\frac{\partial y}{\partial x^{T}}= \begin{bmatrix} \mathbb{I}_{d} & 0 \\ \frac{\partial y_{d+1:D}}{\partial x^{T}_{1:d}} & \text{diag}(\exp[s(x_{1:d})]) \end{bmatrix}
$$
따라서 행렬식은 $\exp \left[\sum_{j}s(x_{1:d})_{j}\right]$ 로 계산된다.

coupling layer는 역함수가 간단하게 계산될 수 있는 특성을 가지고 있으며, 샘플링 과정에서도 효율적으로 사용된다.