생성 AI - Ch6. 노멀라이징 플로 모델(Normalizing flow model)

생각보다 성실한 나날들의 연속이다.
벌써 생성AI 6장을 배우고 이에 관해 정리하고 있다.

자! 그럼 오늘도 힘내면서! 6장인 Normalizing flow 모델에 대해 배워보자.

공부를 본격적으로 시작하기에 앞서 normalizing flow 모델이 우리가 전에 배웠던 autoregressive 모델, VAE와의 공통점과 차이점을 간단히 정리해보겠다.

Autogressive 모델과의 공통점

• 다루기 쉽고 명시적인 데이터 생성 분포 p(x)를 모델링 가능

VAE 모델과의 공통점

• 데이터를 가우스분포와 같은 간단한 분포에 매핑

차이점

• 매핑함수는 역함수가 가능해야함.

오늘의 목차는

6.2 Normalizing Flow

6.3 RealNVP

6.4 다른 Normalizing flow 모델

과 같다.

6.2 Normalizing flow

노멀라이징 플로 모델은 VAE와 유사한 형태로 구성된다. 즉 인코더와 디코더로 구성된다. 다만 디코더가 인코더의 역함수라는 점이 다르다.
그렇다면 인코더의 역함수로 사용하려면 어떻게 해야할까?

이 질문에 대한 답을 얻기 위해 변수 변환(change of variables)에 대해 배워보자.

6.2.1 변수 변환(change of variables)

변수 변환은 x를 z로, 다시 z를 x로 변환하는 과정을 말한다.

우리의 목적은 원본 데이터 x를 latent space인 Z의 한 포인트로 매핑하는 것이다. (이때 Z는 X에서 1. 분포 이동 2. 스케일 조정을 거쳐 매핑되는 곳이다.)

이를 위해 함수 f를 정의해보자!
f는

• z=f(x)
• 변수 z=(z1, z2)와 X의 각 포인트를 정확히 Z의 한 포인트에 매핑하는 함수이다.
  *이때 z와 함수 f(x)는 그림 속 분홍색 박스에 정의되어있다.
• 함수 f를 가역함수(invertible function)라 부른다!

z를 x로 매핑하는 함수 g는

• x=g(z)

자 그럼 이제 X에서 Z로의 변수 변환이 확률변수 Px(x)에 미치는 영향을 알기 위해 확률분포 Pz(z)에 대해 알아보자.

위 그림에서 Pz(z)를 적분해 1/6이란 값을 얻었다. 그런데 확률분포를 샘플링이 가능한 분포로 변환하기 위해서는 적분 결과가 1이 되어야하기 때문에 Pz(z)가 더 이상 유효하지 않다는 사실을 알 수 있다.

그렇다면 Pz(z)가 유효한 확률분포가 되려면 어떻게 해야할까?
답은 정규화 계수를 곱하는 것이다.
즉 정규화 계수를 곱해 원본 데이터의 면적과 같은 데이터로 만들어줘야한다!

다음 절에서는 정규화 계수를 쉽게 구할 수 있는 야코비안 행렬에 대해 배워보자.

6.2.2 야코비 행렬식(Jacobian matrix)

Jacobian matrix에 대한 설명은 말보다 수식이 더 간단할 것 같아 내가 정리한 수식으로 대체하겠다.

6.2.3 변수 변환 방정식(change of variable equation)

변수 변환 방정식이란?

X와 Z 사이의 변수 변환 과정을 설명하는 하나의 방정식

이다.
이 역시 수식으로 설명하는 게 편할 것 같아 관련 설명은 내가 정리한 수식으로 대체하겠다.

그런데 이렇게 변수 변환 방정식을 적용할 때 문제점 2가지가 발생한다.
1. 고차원 행렬의 행렬식을 계산할 때 매우 많은 비용이 발생
2. f(x)의 역함수를 계산하는 방법이 불명확

이 2가지 문제점을 해결하기 위해 RealNVP가 등장했다.

6.3 RealNVP

이 RealNVP의 신경망에는 2가지 목표가 있다.
1. 역변환 가능(역함수 성립 가능)
2. 야코비 행렬 쉽게 계산 가능

6.3.1 초승달 데이터셋

# 데이터 로드
#잡음이 있고 3000개의 포인트를 가진 초승달 데이터셋
data = datasets.make_moons(30000, noise=0.05)[0].astype("float32") 
norm = layers.Normalization()
norm.adapt(data)
#표준정규분포로 정규화. 
normalized_data = norm(data) 
plt.scatter(
    normalized_data.numpy()[:, 0], normalized_data.numpy()[:, 1], c="green"
)
plt.show()

6.3.2 커플링 층

커플링층은 입력의 각 원소(포인트)에 대해 scale factor와 translation factor을 만든다.
=>입력과 동일한 크기의 텐서 2개(scale factor, translation factor)를 만든다.

*이 모델은 이미지가 아니기 때문에 입력의 크기(shape)가 2차원이다.

코드에 주석을 달아놓았으니 천천히 보면서 이해해보자.

def Coupling(input_dim, coupling_dim, reg):
      #입력의 크기가 2차원. 
    input_layer = layers.Input(shape=2)

    #scale factor 산출하기. 
    #scale factor를 위해 크기가 256인 dense층 쌓기
    s_layer_1 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(input_layer)
    s_layer_2 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(s_layer_1)
    s_layer_3 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(s_layer_2)
    s_layer_4 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(s_layer_3)

    #마지막 층은 크기가 2이고 tanh활성화 함수 사용 
    s_layer_5 = layers.Dense(
        input_dim, activation="tanh", kernel_regularizer=regularizers.l2(reg)
    )(s_layer_4)

    #translation factor 산출하기. 
    #translation factor을 위해 크기가 256인 dense층 쌓기
    t_layer_1 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(input_layer)
    t_layer_2 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(t_layer_1)
    t_layer_3 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(t_layer_2)
    t_layer_4 = layers.Dense(
        coupling_dim, activation="relu", kernel_regularizer=regularizers.l2(reg)
    )(t_layer_3)

    #마지막 층은 크기가 2이고 linear 활성화 함수 사용. 
    t_layer_5 = layers.Dense(
        2, activation="linear", kernel_regularizer=regularizers.l2(reg)
    )(t_layer_4)


    #입력으로 input_layer을 넣으면서 scale factor와 translation factor 을 결과로 출력
    return models.Model(inputs=input_layer, outputs=[s_layer_5, t_layer_5])

커플링 층으로 데이터 전달하기.

사실 커플링 층의 구조는 특이하지 않지만 이 속에서 마스킹을 하고 변환되는 과정이 독특하기 때문에 이 과정에 대해 자세히 알아보자.

이 그림은 forward 과정을 보여준다.

• x의 d차원은 커플링 층에 주입 & 다음 단계로 이전

• x의 D-d차원은 완전히 마스킹

• 출력으로 scale factor와 translation factor가 나옴
  ->역 마스킹을 거친 뒤 (0,s2)와 (0, t2)가 출력됨.

• 출력된 계수를 x의 D-d차원에 원소별로 적용.

  여기서 마스킹을 왜 사용하는지는 그림의 오른쪽에 나와있고 이를 한 마디로 정리하면 계산의 편리성을 위해서이다!

  즉 jacobian 행렬의 결과가 하삼각 행렬이 되기 때문에 결과가 대각 원소의 곱으로 매우 간단해진다.
  이는 RealNVP의 두번째 목표를 충족시킨다.

  자 그럼 다음은 첫번째 목표인 역함수를 만들어보자.

  

그림 왼쪽에 정방향 계산을 재정렬한 역함수가 정리되어있다.

자 그럼 이제 이 forward 식과 inverse식을 결합해 커플링 층을 쌓아보자!
이때 핵심은 그대로 유지되던 x의 d차원을 업데이트해야한다는 점이다.

업데이트를 위해 마스킹을 매번 뒤집는다.

즉,
커플링 층을 쌓고 매번 마스킹을 뒤집으면([0,1]<->[1,0])
이전 층에서 변경되지 않았던 부분이 다음 층에서 업데이트된다.

이 과정을 반복하면 입력이었던 x가 전체 입력 텐서로 변환되는 과정울 만들 수 있다.

6.3.3 RealNVP 모델 훈련

코드를 살펴보자.

class RealNVP(models.Model):
    def __init__(
        self, input_dim, coupling_layers, coupling_dim, regularization
    ):
        super(RealNVP, self).__init__()
        self.coupling_layers = coupling_layers

         #타깃 분포는 표준 2D 가우스 분포이다.
        self.distribution = tfp.distributions.MultivariateNormalDiag(
            loc=[0.0, 0.0], scale_diag=[1.0, 1.0]
        )

        #번갈아 바뀌는 마스크 패턴을 만든다.
        #[0,1]=>(x1, 0), [1,0]=>(0, x2)
        self.masks = np.array(
            [[0, 1], [1, 0]] * (coupling_layers // 2), dtype="float32"
        )
        self.loss_tracker = metrics.Mean(name="loss")

        #coupling 층 리스트로 RealNVP 신경망 정의. 
        self.layers_list = [
            Coupling(input_dim, coupling_dim, regularization)
            for i in range(coupling_layers)
        ]

    @property
    def metrics(self):
        return [self.loss_tracker]


    #coupling 층 순회하는 함수
    #training=True -> 정방향으로 층 통과(x->z)
    #training=False ->역방향으로 층 통과(z->x)
    def call(self, x, training=True):
        log_det_inv = 0
        direction = 1
        if training:
            direction = -1
        for i in range(self.coupling_layers)[::direction]:
            x_masked = x * self.masks[i]
            reversed_mask = 1 - self.masks[i]
            s, t = self.layers_list[i](x_masked)
            s *= reversed_mask
            t *= reversed_mask
            gate = (direction - 1) / 2

            #direction에 따라 정방향/역방향 식 구별
            x = (
                reversed_mask
                * (x * tf.exp(direction * s) + direction * t * tf.exp(gate * s))
                + x_masked
            )
            #손실함수는 스케일링 계수의 합
            log_det_inv += gate * tf.reduce_sum(s, axis=1)
        return x, log_det_inv

    def log_loss(self, x):
        y, logdet = self(x)

        #손실함수는 타깃 가우스 분포와 야코비 행렬식의 로그 값으로 결정되는 변환된 데이터의 음의 로그 확률의 합. 
        log_likelihood = self.distribution.log_prob(y) + logdet
        return -tf.reduce_mean(log_likelihood)

    def train_step(self, data):
        with tf.GradientTape() as tape:
            loss = self.log_loss(data)
        g = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(g, self.trainable_variables))
        self.loss_tracker.update_state(loss)
        return {"loss": self.loss_tracker.result()}

    def test_step(self, data):
        loss = self.log_loss(data)
        self.loss_tracker.update_state(loss)
        return {"loss": self.loss_tracker.result()}


model = RealNVP(
    input_dim=INPUT_DIM,
    coupling_layers=COUPLING_LAYERS,
    coupling_dim=COUPLING_DIM,
    regularization=REGULARIZATION,
)

출력 결과

<훈련하기 전>

<훈련한 후>

가우스 분포에서 샘플링한 포인트가 원본 데이터인 x와 매우 유사함을 알 수 있다!

6.4 다른 노멀라이징 플로 모델

다른 모델들은 이것에서 조금더 발전시킨 것이고 다른 모델들까지 설명하기에는 너무 복잡해서 차이점만 간단히 정리하고 넘어가겠다.

6.4.1 GLOW 모델

기존 RealNVP에 이미지를 적용하면 단계마다 채널의 순서가 바뀌면서 신경망이 모든 입력을 변환할 가능성이 있다는 단점이 있다.
이를 개선하기 위해
GLOW에서는

1X1 합성곱을 적용해서 모델이 원하는 채널 순서대로 조합을 생성한다.

6.4.2 FFJORD

FFJORD는 RealNVP와 GLOW의 한계였던 discrete time을 continuous time으로 발절시켰다.
즉,

연속적인 시간 과정으로 변환을 모델링하는 과정을 담았다. 이때
p(z(t))함수를 사용해 데이터 분포를 표준 가우스 분포로 변환한다.