MAF IAF 논문 리뷰(Masked Autoregressive Flow for Density Estimation & Improved Variational Inference with Inverse Autoregressive Flow)

2017년 Papamakarios et at.의 MAF(Masked Autoregressive Flow)과 2016년 Kingma et al.의 IAF(Inverse Autoregressive Flow)를 리뷰할 예정이다. 두 논문은 사실상 완전히 대칭되는 아이디어를 가지고 쓰여진 논문으로 혼자 흥미로웠던 점은 IAF가 먼저 나오고 MAF가 나중에 나온 논문이라는 점이다.

 

IAF와 MAF는 완전히 반대의 아이디어를 사용하기 때문에 장단점도 정확히 반대이다.

 

IAF는 samlpling이 빠른 대신 likelihood를 계산하는데 오래 걸린다.

MAF는 반대로 likelihood 계산이 빠른 대신 sampling이 오래 걸린다는 단점이 있다.

 

MAF와 IAF 모두 Autoregressive model + Flow-based generative model이라서 두 모델의 특징을 가지고 있다.

 

* Autoregressive model은 model의 성질 때문에 그 자체로 invertible하고 computable한 Jacobian determinant를 가지고 있다. 즉 Autoregressive model을 일종의 Flow-based generative model로 볼 수 있다. 또한 이 두 논문은 Real NPV에서 더욱 유연한 구조를 가지고 있으며 더 좋은 퍼포먼스를 보인다고 알려져있다.

[16' NeurIPS] IAF(Improved Variational Inference with Inverse Autoregressive Flow)(Kingma et al.)

* Introduction에 나온 VAE와 Normalizing Flow, Autoregressive model 설명 생략

 

Inverse Autoregressive Transformations

위와 같이 세팅된 경우를 생각해보자.

 

위와 같은 모델에서 샘플링을 한다는 것은 noise vector $\epsilon \sim N(0, I)$에서 $\bold{y}$로 보낸다는 의미이다.

$\bold{y}$를 샘플링 하는 방법

위 계산은 $\bold{y}$ 차원 D에 의존하게 된다. 

여기서 만약 모든 i에 대해 $\sigma_i > 0$이라면 sampling transformation은 1-1 transformation이 되고 아래와 같이 쓸 수 있다.

벡터로 보면 위와 같이 볼 수 있다.

여기서 눈여겨 보아야 할 점은 두가지이다.

1. This inverse transformation can be parallelized since computations of the individual elements $\epsilon_i$ do not depend on eachother. 

→ 역방향 변환이 병렬적으로 계산 가능하기 때문이 계산 속도를 향상 시킬 수 있다.

 

2. This inverse autoregressive operation has a simple Jacobian determinant.

→ 계산하기 쉬운 Jacobian determinant를 가지고 있다.

Inverse Autoregressive Flow(IAF)

위에서 제시한 transformation을 이용해 IAF를 만들 것이다.

$\epsilon \sim N(0,I)$에서 랜덤하게 고른 $\epsilon$을 IAF을 태운다. 

그러면 아래와 같이 분포를 나타낼 수 있다.

알고리즘은 아래와 같이 나타낼 수 있다.

 

결론

latent variable $z$에 대해서는 autoregressive property를 보전하기 위해 sequential하게 계산하기 때문이 likelihood를 계산하는 과정은 느리다. 하지만 inverse transformation의 경우 parallel하게 실행하기 때문에 sampling 과정이 상당히 빠르다는 사실을 알 수 있다.

 

* IAF 는 variational inference를 위해 처음 고안되었으며, density estimation에는 적합하지 않다

 

[17' NeurIPS] MAF(Masked Autoregressive Flow for Density Estimation)(Papamakarios et al.)

* Introduction에 나온 VAE와 Normalizing Flow, Autoregressive model 설명 생략

Masked Autoregressive Flows

사실상 IAF와 비슷한 설정이지만, 생성되는 평균과 표준편차가 $\bold{u}$가 아닌 $\bold{x}$에 의해 결정된다는 사실을 눈여겨 보자.

 

autoregressive model이기 때문에 함수 $f$는 invertible하며, inverse transformation은 다음과 같이 표현 가능하다.

Relationship with IAF

위에서 보이는 구조의 차이가 MAF와 IAF의 차이를 가져온다.

 

MAF

- calculate the density $p(x)$ of any datapoint x in one pass through the model

- sampling from MAF requires performing D(=dim of x) sequential passes

- suitable for density estimation.

 

☞ previous data variable로 평균과 분산을 구하기 때문에 density 계산은 빠르지만, sampling이 느리다.

 

IAF 

- generate sample and calculate their density with one pass

- calculate the density $p(x)$ of an externally provided datapoint x requires D passes to find the random numbers u associated with x.

- suitable as a recognition model for stochastic variational inference.

 

☞ previous random number로 평균과 분산을 구하기 때문에 sampling은 빠르지만 density 계산이 느리다.

참고

식 (7)번은 MAF로부터, 식 (8)번은 IAF로부터 나온 식이다.

우리는 MAF에서 $p_x(x)$와 $\phi_x(x)$의 KL Divergence 값이 가깝기를 바라는데, 이는  $p_u(x)$와 $\phi_u(x)$의 KL Divergence 값을 낮추는 것과 같다.

 

결론 :  Training a MAF with maximum likelihood corresponds to fitting an implicit IAF to the base density with stochastic variational inference.

 

Relationship with RealNVP

RealNVP

RealNVP는 MAF와 IAF에서 제약이 추가된 모형으로 MAF와 IAF가 좀 더 flexible한 모델이라고 볼 수 있다.

다만, RealNVP는 sampling과 density estimation 모두 one pass로 빠르게 계산이 가능하다는 장점을 가지고 있다. 

 

코드 비교(출처 : pytorch_flows/flows_04.ipynb at master · acids-ircam/pytorch_flows · GitHub)

# Masked Autoregressive Flow
class MAFlow(Flow):
    def __init__(self, dim, n_hidden=32, n_layers=2, activation=nn.ReLU):
        super(MAFlow, self).__init__()
        self.g_mu = self.transform_net(dim, dim, n_hidden, n_layers, activation)
        self.g_sig = self.transform_net(dim, dim, n_hidden, n_layers, activation)
        self.init_parameters()
        self.bijective = True

    def transform_net(self, nin, nout, nhidden, nlayer, activation):
        net = nn.ModuleList()
        for l in range(nlayer):
            net.append(MaskedLinearAR(l==0 and nin or nhidden, l==nlayer-1 and nout or nhidden))
            net.append(activation())
        return nn.Sequential(*net)
        
    def _call(self, z):
        zp = z * torch.exp(self.g_sig(z)) + self.g_mu(z)
        return zp

    def _inverse(self, z):
        z = (z - self.g_mu(z)) / self.g_sig(z)
        return z

    def log_abs_det_jacobian(self, z):
        return -torch.sum(torch.abs(self.g_sig(z)))

 

# Masked Autoregressive Flow
class IAFlow(Flow):
    def __init__(self, dim, n_hidden=32, n_layers=2, activation=nn.ReLU):
        super(IAFlow, self).__init__()
        self.ar_net = self.transform_net(dim * 2, dim * 2, n_hidden, n_layers, activation)
        self.init_parameters()
        self.bijective = True
        self.sigma = []

    def transform_net(self, nin, nout, nhidden, nlayer, activation):
        net = nn.ModuleList()
        for l in range(nlayer):
            net.append(MaskedLinearAR(l==0 and nin or nhidden, l==nlayer-1 and nout or nhidden))
            net.append(activation())
        return nn.Sequential(*net)
        
    def _call(self, z):
        out = self.ar_net(z)
        m, s = out.chunk(2, dim=1)
        z, h = z.chunk(2, dim=1)
        self.sigma = torch.sigmoid(s)
        z = sigma * z + (1 - sigma) * m
        return torch.cat((z, h), dim=1)

    def _inverse(self, z):
        z = (z - self.g_mu(z)) / self.g_sig(z)
        return z

    def log_abs_det_jacobian(self, z):
        return torch.sum(torch.abs(self.g_sig(z)))

 

두 코드의 비교

1. MAF는 g_mu와 g_sig 두 개의 독립적인 변환을 사용하지만, IAF는 단일 네트워크 ar_net을 사용한다.  

2. (_call 함수) MAF는 z 변환시 zp = z * torch.exp(self.g_sig(z)) + self.g_mu(z)을 계산하지만, 

IAF는 ar_net(z)의 출력을 m과 s로 나누고, z를 sigma * z + (1 - sigma) * m으로 변환한다.