[G]Part 2-4. Normalizing Flows

📌 목차

1. Normalizing Flow
2. RealNVP
3. GLOW  &  FFJORD
4. 요약

 

🧐  preview: 

앞서 설명한 VAE, GAN, AR모델은 분포 p(x)모델링문제해결을 위해
Sampling가능한 latent variable이나 (=VAE Decoder, GAN Generator)
이전 pixel값에 대한 함수로 분포를 모델링(=AR)한다.

Normalizing Flow와의 공통점:
∙ AR처럼 다루기 쉽고 명시적인 data생성분포 p(x)를 모델링
∙ VAE처럼 data를 Gaussian분포처럼 간단한 분포에 mapping

Normalizing Flow의 다른점:
∙ mapping함수의 형태에 제약을 둔다.
(= 이 함수는 반전가능해야하고, 이를 이용해 새 data point생성이 가능해야함.)

 

 

 

 

 

 

 

 

---

1. Normalizing Flow

Normalizing Flow는 VAE와 상당히 유사하다.

VAE

Encoder학습 → 복잡한 분포와 sampling가능한 간단한 분포를 mapping
→ Decoder학습 → 단순한분포에서 복잡한 분포로 mapping
∴ 단순한 분포에서 point z를 sampling, 학습된 변환을 적용 시, 새 data point를 얻음.
즉, Decoder는 p(x|z)를 모델링.
Encoder는 p(z|x)의 근사치인 q(z|x)로 둘은 완전히 다른 신경망임.

Normalizing Flows

Decoding함수 = Encoding함수의 역함수 (= 계산이 용이함.)
But❗️신경망은 기본적으로 반전가능한 함수가 아니다!
이를 위해 변수변환이라는 기법이 필요한 것.

---

Change of Variables

변수변환:

p_x(x)가 2차원 x = (x₁, x₂) 직사각형 X위에 정의되어있다 가정하자.

이 함수를 주어진 분포영역에서 적분하면 1이 된다. (x1의 범위: [1,4] , x2:[0.2])
따라서 아래와 같이 잘 정의된 확률분포를 정의할 수 있다:

이 분포를 이동하고 scaling하여 단위 정사각형 Z에 대해 정의한다 가정하자.
이를 위해 새로운 z=(z1, z2)변수와 X의 각 point를 정확히 Z의 한 Point에 mapping하는 함수 f를 정의할 수 있다:

이런 함수를 가역함수(invertible function)라 한다.
즉, 모든 z를 이에 해당하는 x로 다시 mapping가능한 g함수가 존재한다. (변수변환에 필수적.)
이 g함수가 없으면 두 공간사이를 mapping할 수 없다.

이때, p_z(z) = ((3z₁+1)-1)(2z₂)/9 = 2z₁z₂/3로 변환가능하다.
다만, p_z(z)를 단위면적에 대해 적분하면 아래와 같은 문제에 봉착한다:

1/6값이 나왔기에, 더이상 유효한 확률분포가 아니다. (적분결과가 1이 되어야함.)
원인: 변환된 확률분포 영역이 원본의 1/6로 작아졌기 때문.
해결: 새로운 확률분포의 상대적면적변화에 해당하는 정규화계수를 곱해야함
→ 변환에 의한 부피변환 = Jacobian Matrix의 절댓값.

---

Jacobian Matrix:

z = f(x)의 Jacobain Matrix는 1계 편미분도함수의 행렬이다.

x1에 대한 z1의 편미분은 1/3이고, x1에 대한 z2의 편미분은 0이며,
x2에 대한 z2의 편미분은 0이고, x2에 대한 z2의 편미분은 1/2이다.

따라서 함수 f(x)의 Jacobian Matrix는 아래와 같다.

행렬식(determinant)은 정방행렬(square matrix)에서만 정의된다.
해당행렬로 표현되는 변환을 단위 초입방체에 적용해 만들어진 평행육면체 부피와 같다.
즉, 2차원 행렬로 표현되는 변환을 단위 정사각형에 적용해 만들어진 평행사변형의 면적에 해당한다.


앞선 예시를 보자. det(a b c d) = ad - bc이므로 
위 예시에서 Jacobian 행렬식은 1/3 × 1/2 = 1/6이다.
이값이 바로 scaling계수이다.

---

변수 변환 방정식

하나의 방정식으로 X와 Z사이 변수변환과정을 설명할 수 있다.
이를 "변수변환방정식(change of variables equation)"이라 한다.

간단한 분포라면, X→Z로 mapping할 적절한 가역함수 f(x)와 
sampling된 z를 기존 domain point x로 다시 mapping 시 사용할 역함수 g(z)를 찾으면 된다.
이를 위해, 필요한 2가지 문제가 존재한다.

① 고차원 행렬식은 Too Hight Cost. O(n³)시간이 든다.
② f(x)의 역함수 계산이 명확하지 않다.

위의 문제들 해결을 위해 RealNVP라는 변환기법으로 수행가능하다.

 

 

 

 

 

 

 

 

 

 

 

 

 

 

---

2.  RealNVP

prev.

RealNVP는 복잡한 data분포를 간단한 Gaussian분포로 변환하는 신경망을 만들 수 있다.
또한, 역변환이 가능하고 Jacobian Matrix를 쉽게 계산할 수 있다.

 

Coupling Layer

Coupling층은 input원소에 대해 scale계수와 translation계수를 만든다.

아래 예시처럼 Linear층으로 Scale출력을 만들고,
또다른 Linearfh translation계수를 만든다.

class CouplingLayer(nn.Module):
    def __init__(self, input_dim, output_dim, hid_dim, mask):
        super().__init__()
        self.s_fc1 = nn.Linear(input_dim, hid_dim)
        self.s_fc2 = nn.Linear(hid_dim, hid_dim)
        self.s_fc3 = nn.Linear(hid_dim, output_dim)
        self.t_fc1 = nn.Linear(input_dim, hid_dim)
        self.t_fc2 = nn.Linear(hid_dim, hid_dim)
        self.t_fc3 = nn.Linear(hid_dim, output_dim)
        self.mask = mask

    def forward(self, x):
        x_m = x * self.mask
        s_out = torch.tanh(self.s_fc3(F.relu(self.s_fc2(F.relu(self.s_fc1(x_m))))))
        t_out = self.t_fc3(F.relu(self.t_fc2(F.relu(self.t_fc1(x_m)))))
        y = x_m + (1-self.mask)*(x*torch.exp(s_out)+t_out)
        log_det_jacobian = s_out.sum(dim=1)
        return y, log_det_jacobian

    def backward(self, y):
        y_m = y * self.mask
        s_out = torch.tanh(self.s_fc3(F.relu(self.s_fc2(F.relu(self.s_fc1(y_m))))))
        t_out = self.t_fc3(F.relu(self.t_fc2(F.relu(self.t_fc1(y_m)))))
        x = y_m + (1-self.mask)*(y-t_out)*torch.exp(-s_out)
        return x​

더 복잡한 표현을 학습하기 위해 차원을 늘렸다 다시 원본차원으로 축소한다.

Coupling Layer는 input data가 들어갈 때, Masking 후 변환되는 방식이 독특하다.

Step 1. 처음 d차원만 Coupling Layer에 주입.
Step 2. 남은 D-d차원은 완전히 Masking(= 0으로 설정.)

이때, 많은 정보를 Masking하는 이유가 뭘까?
이를 알기위해, 이 함수의 Jacobian Matrix를 살펴보면 알 수 있다.

장점 1

하삼각행렬형태로 하삼각행렬의 행렬식은 단순히 대각원소의 곱과 같다.
즉, 좌측하단의 복잡한 도함수와 상관이 없어진다!

장점 2

아래 그림 및 식을 보면, 쉽게 역전할 수 있는 함수라는 목표달성이 가능하다.
정방향계산을 재정렬하면, 아래와 같은 역함수식을 만들 수 있다.

Step 3. 그러면 이제 처음 입력의 d개 원소를 어떻게 update해야할까?

 

Coupling Layer 쌓기

선형대수학의 행렬식 조건

따라서 위의 선형대수학 행렬식조건에 따라 Coupling Layer를 쌓고, 
매번 Masking을 뒤집으면, 간단한 Jacobian Matrix와 가역성이라는 필수속성을 유지하며
전체 input tensor를 변환하는 신경망을 만들 수 있다.

RealNVP 모델훈련

class RealNVP(nn.Module):
    def __init__(self, input_dim, output_dim, hid_dim, mask, n_layers = 6):
        super().__init__()
        assert n_layers >= 2, 'num of coupling layers should be greater or equal to 2'

        self.modules = []
        self.modules.append(CouplingLayer(input_dim, output_dim, hid_dim, mask))
        for _ in range(n_layers-2):
            mask = 1 - mask
            self.modules.append(CouplingLayer(input_dim, output_dim, hid_dim, mask))
        self.modules.append(CouplingLayer(input_dim, output_dim, hid_dim, 1 - mask))
        self.module_list = nn.ModuleList(self.modules)
        
    def forward(self, x):
        ldj_sum = 0 # sum of log determinant of jacobian
        for module in self.module_list:
            x, ldj= module(x)
            ldj_sum += ldj
        return x, ldj_sum

    def backward(self, z):
        for module in reversed(self.module_list):
            z = module.backward(z)
        return z


mask = torch.from_numpy(np.array([0, 1]).astype(np.float32))
model = RealNVP(INPUT_DIM, OUTPUT_DIM, HIDDEN_DIM, mask, N_COUPLE_LAYERS)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
prior_z = distributions.MultivariateNormal(torch.zeros(2), torch.eye(2))

Loss functions: NLL Loss

def train(epoch):
    model.train()
    train_loss = 0
    for batch_idx, data in enumerate(train_loader):
        optimizer.zero_grad()
        z, log_det_j_sum = model(data)
        loss = -(prior_z.log_prob(z)+log_det_j_sum).mean()
        loss.backward()
        cur_loss = loss.item()
        train_loss += cur_loss
        optimizer.step()
        if batch_idx % LOG_INTERVAL == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100.*batch_idx / len(train_loader),
                cur_loss/len(data)))

    print('====> Epoch: {} Average loss: {:.4f}'.format(
        epoch, train_loss / len(train_loader.dataset)
    ))

[train 결과]:

복잡한 input이 단순한 Gaussian분포형태로 mapping됨.

 

 

 

 

 

 

 

 

 

 

 

---

3. GLOW  &. FFJORD

GLOW

고품질 sample생성 및 유의미한 latent space 생성간으.
핵심: 역마스킹 설정을 반전가능한 1×1 conv로 대체
오죽하면 논문제목도 Glow: Generative Flow with Invertible 1x1 Convolutions

이를 이용해 모델이 원하는 채널순으로 조합생성 가능.

---

FFJORD

Discrete time Normalizing Flows: RealNVP, GLOW
Continuous time Normalizing Flows: FFJORD

즉, Normalizing Flow의 단계수가 무한대이고, 단계크기가 0에 가깝기에
신경망에 의해 parameter가 정의되는 상미분방정식(ODE)을 사용해
data분포와 표준정규분포사이 변환을 모델링:

 

 

 

 

 

 

 

 

4. 요약

Normalizing Flow란, 다루기 쉬운 data밀도함수를 만들면서 고품질 sample을 생성하는 강력한 생성모델링방법이다.
이를 위해 신경망의 형태를 제한해 가역성과 계산하기 쉬운 Jacobian determinant를 계산한다.