Generative Flow
1. Flow-Based Generatvie Model
Give a datapoint $x$ and latent varible $z\sim p_{\theta}(z)$. The function $g_{\theta}(..)$ is invertible such that
$$ \begin{aligned} x&=g_{\theta}(z) \\ z&=f_{\theta}(x) =g_{\theta}^{-1}(x) \end{aligned} $$
The probability density function of the model can be written as:
$$ \begin{aligned} \log p_{\theta}(x) &= \log p_{\theta}(z) + \log | det(dz/dx) | \\ &= \log p_{\theta}(z) + \sum_{i=1}^{K}\log | det(dh_i/dh_{i-1}) | \end{aligned} $$
To simplify the calculation, we can choose transformations with Jacobian $dh_{i}/dh_{i-1}$ is a triangular matrix,
$$ \log| det(dh_{i}/dh_{i-1}) | = sum(\log | diag(dh_{i}/dh_{i-1}) |) $$
where $\log$ is element-wise logarithm.
2. Generative Flow
One step flow is composed of a actnorm, a invertible 1x1 convolution and a coupling layer.
2.1 actnorm
$$ \forall i,j:y_{i,j}=s\odot x_{i,j}+b \\ \forall i,j:x_{i,j}=(y_{i,j}-b)/s \\ \log| \det | = h\cdot w\cdot sum(\log |s|) $$
For example, $[x1, x2, x3]\odot[w1, w2, w3]=[y1, y2, y3]$ where $\odot$ is element-wise operation. The Jacobian is
$$\left[ \begin{matrix} w1 & 0 & 0 \\ 0 & w2 & 0 \\ 0 & 0 & w3 \end{matrix} \right]$$
initialization..
2.2 invertible 1x1 convolution
$$ \forall i,j:y_{i,j}=Wx_{i,j} \\ \forall i,j:x_{i,j}=W^{-1}y_{i,j} \\ \log|\det| = h\cdot w\cdot \log|\det(W)| $$
where the Jacobian is $W:[c \times c]$. We can reduce the complexity from $O(c^3)$ to $O(c)$ by parameterizing $W$ in its PLU decomposition:
$$ W = PL(U+diag(s)) $$
then,
$$ \log|\det(W)| = sum(\log|s|) $$
2.3 coupling layer
$$\begin{aligned} x_a, x_b &= split(x) \\ \log s,t &= NN(x_b) \\ s &= \exp(\log s) \\ y_a &= s\odot x_a + t \\ y_b &= x_b \\ y &= concat(y_a, y_b) \end{aligned}$$
The inverse is:
$$\begin{aligned} y_a,y_b &= split(y) \\ (\log s,t) &= NN(y_b) \\ s &= \exp(\log s) \\ x_a &= (y_a-t)/s \\ x_b &= y_b \\ x &= concat(x_a,x_b) \end{aligned}$$
The log-determinant is :
$$\log|\det| = \log|\det(s)|$$
where the Jacobian is :
$$\left[ \begin{matrix} diag(s) & u \\ 0 & \mathbb{I} \end{matrix} \right]$$