생성모델의 흐름 2(VAE to Diffusion)

seung·2025년 1월 12일
Generative ModelVAEdiffusion

Generative Model

목록 보기
6/7

VAE

Variational Autoencoders
생성모델의 흐름 1에서 배웠듯이 VAE는 latent variable model이고 Variational Bayes를 신경망을 이용해서 한 모델이다.

자세히 살펴보면


x: 관측한 값

  • decoder
    z가 주어졌을 때 x의 확률분포를 내보내주는 p(x|z)

  • encoder
    x가 주어졌을 때 그것을 잘 나타내는 z를 뽑아내는 q(z|x)

이 2개가 연결되어 있는 구조이다.

latent variable model이 z 모델링을 잘하면 복잡해보이는 데이터(x)의 분포를 잘 묘사할 수 있다.

하지만 z를 잘 모델링을 하는것이 어렵기 때문에
입력 데이터 x로부터 z를 알아서 학습하고 학습된 z로 부터 x를 복원하는 과정을 신경망으로 하는 VAE가 나오게 되었다.

VAE를 통해 우리는 성공적으로 생성모델을 학습할 수 있다는 것을 알게 되었다.

이렇게 학습된 Z들은 재미있게도 라벨이 주어지지 않아도 input을 클래스별로 잘 나누어서 인코딩하게 되고,
z space가 가우시안 분포로써 정의되어 샘플링하여 decoder에 넣으면 새로운 이미지를 생성하는 것이다!

HVAE

Hierarchical Variational Autoencoders

우리가 VAE를 보면 latent space를 하나의 구조로 만든 것이다.
즉 하나의 계층으로 되어 있는것이다.

생각해보면 사실 한 개의 계층을 구성하는 것보다 여러개의 계층으로 구성한다면 더 좋은 성능을 낼 수 있지 않을까? 라는 생각에서 나온 모델이
HVAE이다!

예를 들어 데이터가 사람의 얼굴이라고 하면
Z가 얼굴의 중요한 정보를 담는다고 하면 눈,코,입,피부 등을 하나의 계층에서 묶어서 표현하는 것보다 여러 계층으로 나누어서 세밀한 것부터 전체적인 것까지 표현하는 것이 훨씬 좋을 것이다!

위의 그림으로 보면,

기존의 Z가 하나인 VAE를 확장시켜서 Z를 더 글로벌하게 인코딩하고 그 Z를 또 인코딩하는 구조를 갖고 있다.

수식적인 형식은 크게 벗어나지 않지만 Z를 좀 더 복잡하게 디자인할 수 있다.

HVAE의 수식

t개의 계층으로 이루어진 HVAE의 수식을 살펴보면 아래와 같다.

여기서 t가 1인 HVAE는 VAE의 구조와 같다!

Markovian 가정을 통해 조건부 확률을 이용해서 계속해서 확장해 나갈 수 있는 구조를 띄고 있다.

  • Markovian이란?
    다음 단계의 variable은 그 전단계의 variable에만 종속되어 영향을 받는다는 것을 가정한다.
    수학적인 수식을 단순하게 하기 위한 가정이다.

위의 가정을 통해 ELBo 수식도 VAE에서 쉽게 확장할 수 있다.

전 게시글에서 배웠듯이 variational inference와 jensen's inequality를 이용해 위와 같은 식으로 정리가 될 수 있다.

vae처럼 ELBo부분을 최대화 하는 것이 목표이다!

여기까지 배웠다면 우리는 Diffusion을 배울 준비가 다 되었다!!

Diffusion(with vae)

Diffusion 모델을 이해하기 쉬운 방법은

3가지 특이제약 가정을 한 Markovian HVAE라고 생각하는 것이다!!!

3가지 제약 조건

  1. data 차원가 latent 차원이 일치한다.
    즉 x와 z의 사이즈가 동일하다.(vae는 고차원을 저차원으로 압축시킨다.)

  2. x에서 z로 가는 과정이 정해져 있다. (gaussian noise를 더해가는 형태)
    vae는 encoder를 학습했어야 한다.

  3. 타임스텝에 따라 뒤로 가면 마지막 t step에서 latent 분포는 표준정규분포를 따른다.
    vae도 분포가 뾰족해지지 않기 위해 정규분포라는 제약 조건을 거는 것과 일치한다.
    (step이 1이면 vae와 같은 구조인 것이다.)

Diffusion 구조

  • x0 : 입력 데이터

  • q : vae에서는 학습을 했지만 diffusion에서는 정해져 있다.
    입력 데이터에 가우시안 노이즈를 조금씩 추가하여 마지막 t step이 되면 가우시안 분포가 되게 한다.
    foward process라고 한다.

foward 구조는 HVAE와 같지만 다만 과정은 deterministic하다!

  • p : reverse process할 때 사용되며, 노이즈에서 데이터로 돌아오는 decoder과정이고, 이 과정은 학습을 하게 된다.

Diffusion 수식

HVAE에서 3가지 제약 사항을 걸게 되면 diffusion이 된다고 했는데
그것에 맞춰서 HVAE의 수식에서 이 제약 사항에 맞게 수정하여
diffusion 수식을 만들어보자!

1번째 제약사항

HVAE와 달리 diffusion은 x와 z의 차원이 같기 때문에 모두 다 x로써 표현할 수 있다.

2번째 제약사항

2번째 가정이 Markovian HVAE와 달리 diffusion의 인코딩은 가우시안 노이즈를 더해가는 과정으로 정해져있다는 것이다.

즉 위의 식처럼 timestep 한 단계 단계가 normal gaussian noise를 더하는 형태이다.

3번째 제약사항

마지막 step때 결과적으로 표준정규분포가 된다는 가정에 따라
x_T는 아래와 같이 생각할 수 있다.

이렇게 되기 위해서는 alpha가 step이 진행될수록 점점 작아져야 한다.

이러한 가정 때문에 diffusion의 인코딩은 variance-preserving이라고도 한다.

Markovian HVAE에 이 3가지 제약 사항을 가정하면, 우리는 diffusion의 수학적인 원리를 대부분 다 이해했다고 할 수 있다!

Diffusion objective function

diffusion의 ELBo를 구하는 방법은 2가지가 있다.

첫 번째 방법은 이해하기는 쉬우나 실제로 구현했을 때 문제가 생길 수 있다.
두 번째 방법은 조금 더 이해하기는 어려우나 실제로 구현했을 때 문제가 없다.

ELBO 1번째 접근


우리가 jensen's inequality를 이용해 위와 같은 식이 성립함을 알 수 있었다.

이 ELBO를 최대화 하기위해 Markovian을 이용해 식을 풀어보면

위와 같이 3가지 term으로써 분리할 수 있고

  1. likelihood
    x1을 latent z라고 하면 z가 주어졌을 때 현재 내 데이터(x0)를 가장 잘 나타내는 p_theta likelihood를 평균낸 것이 최대화 되게 하는 것!

  2. prior matching
    우리가 p(x_T)를 표준 정규분포로 정의했는데 우리가 설정한 q라는 분포가 너무 뾰족해지지 않도록 내가 설정한 p(x_T)의 유사한 거리가 줄어들게 할려고 하는 term

  3. consistency
    foward 에서 오는 분포랑 reverse process에서 오는 분포가 일치해야 된다는 텀인데
    그림으로 보면 즉 x_t의 분포가 foward에서 오든 reverse에서 오든 일관되야 된다는 것이다.

여기서 T가 1이 되면 vae와 동일한 구조가 된다!!
즉 vae에서 조금 더 복잡하고 여러 계층으로 latent를 구성한 것이 diffusion이다!

마지막 3번째 term을 계산할 때 문제가 있는데
그것은 3번째 term을 계산하기 위해서는 두개의 모델과
3개의 변수 x_t+1, x_t, x_t-1을 알고 있어야되며 제약을
여러개 넣어줘야 되다 보니 분산이 커져서 서로 안맞을 가능성이 높아지고,
계산 비용이 커지게 된다.

그래서 이 첫번째 방식은 실제로 구현하더라도 학습이 잘 안되고 오래걸리게 된다.

ELBO 2번째 접근

diffusion 모델을 보면 앞에 방식의 문제를 해결하기 위해 아래와 같은 식으로 objective function을 만들게 된다.

x_0을 추가할 수 있는 이유는 우리가 Markovian을 가정하기 때문에 조건부 부분에 추가해도 영향을 받는게 없기 때문이다.

즉 q(x_t|x_t-1)를 알고 싶으면 베이즈 정리에 의한 식을 이용하면 되는 것이다!

왜 x_0를 끼워넣는걸까?

위에서 x_0를 끼워넣어서 변환시키는 것을 ELBO에 적용 시켜보면 아래와 같은 식이 되는것을 알 수 있다.

위에 식을 log를 좀 더 풀어서 계산하고 정리하면 아래와 같이 정리 할 수 있다.

1,2번째 term은 첫번째 접근과 같다.
3번째 term은 달라진것을 볼 수 있다.

  • Denoising matching
    x_0을 추가하니 consistency term이 x_t 가 주어졌을 때 x_t-1로 가는 과정이 q와 p 두 모델 다 일치해야된다는 것으로 바뀐것을 볼 수 있다.

3번째 term을 보면 방향이 같아진 것을 볼 수 있다.
방향이 같기 때문에 변수 2개만 이용해도 되기 때문에 메모리 효율이 올라가고 같은 방향으로만 학습하기 때문에 분산이 줄어드는 효과를 얻을 수 있다!

자 ELBO 식을 알았으니 이제 계산을 하면 될텐데 문제는
우리가 prior term과 reconstruction term은 VAE에서 배웠듯이 계산하면 되지만,

denoising matching term을 계산할 수 있어야 하는데
우리가 p(x_t-1|x_t)는 우리가 모델링했으니 계산이 가능하나, q()를 알아야 우리분포에 가까워지게 학습할텐데 그 분포를 몰라서 계산을 할 수 가 없다...

또한 x_t에서 x_t-1로 denoising하는 것을 우리가 알고 있다면 굳이 이 모델을 학습할 이유가 있나..?

q( x_t-1 | x_t, x_0 )

이 조건부 확률은 bayes 정리에 의해 아래와 같이 정리할 수 있다.

또한 우리는 q(x_t | x_t-1)의 형태를 이미 정해놓았다.

이때 여기서 우리가 x_t를 reparameterization trick을 사용하면,
표준정규분포에서 샘플링하여 분산을 곱하고 평균을 더한 형태로 바꿔줄 수 있다.

이 형태가 recursive한 형태이기 때문에 x_t를 x_0에 대한 식으로 표현할 수 있다.

정리

이제 우리는 q( x_t-1 | x_t, x_0 )에 필요한 모든 정보를 알게 되었다.
베이즈 정리로 나타나는 q( x_t-1 | x_t, x_0 ) 3가지 텀의 분포가 모두 정규분포를 따른다는 것을 위의 내용을 통해 알 수 있다.

우리는 통계학에서 정규분포끼리 곱하고 나누어도 다시 정규분포가 된다는 점을 알고 있다.

그것을 이용해서 q( x_t-1 | x_t, x_0 )의 분포를 계산해보면,

이제 우리가 알고자 했던 q( x_t-1 | x_t, x_0 )의 분포를 알 수 있게 되었고, 정확하게 계산할 수 있어서 denoising matching term을 최적화 할 수 있다!!

Denoising matching

denoising matching term을 계산하기 위해서는 KL divergence 함수를 계산해야 한다.

놀랍게도 우리가 계속 사용하고 있는 이 가우시안 분포끼리의 KL divergence를 구하는 것은 아래와 같이 깔끔하게 풀 수 있다.

가우시안 분포는 정말 얼마나 사랑스운지 알 수 있다 ㅠㅠ


이제 우리의 term에 적용하면 아래와 같이 쓸 수 있는데 우리가 q 분포의 평균과 분산을 정확하게 아는데 굳이 p 분포를 다른걸로 설정할 이유가 없다!

q( x_t-1 | x_t, x_0 )의 분포를 살펴보면 분산은 상수인 것을 알 수 있따. 그래서 p분포의 분산도 같은 것으로 설정해주면 된다.

하지만 평균은 우리가 예측해야되는 값(x_0)이 포함되어 있어서 불확실성이 추가된다. 그래서 우리는 두 분포의 평균 값만 비교해주면 된다.

분산이 같은 두 분포의 KL divergence는 더욱 간결하게 평균으로만 표현될 수 있다.

이 차이를 줄이기만 하면 두 분포의 차이가 줄어들고, 이렇게 모델링한 p_theta는 정확하게 우리가 원하는 denoising step을 하게 된다!!!

두 평균 값 비교

앞에 내용에서 M_q는 아래와 같은 식으로 표현될 수 있다고 배웠다.

이때 M_theta는 M_q와 똑같이 모델링을 하게 될텐데 이때 모르는 값 x_0를 예측하게 될 것이다.

그러면 결국 우리는 평균에서 같은 것을 제외하면 x_0만 다른 것을 알 수 있다.
즉 평균 대신 모르는 x값과 x의 예측값의 차이로써 계산할 수 있을 것이다.

하지만 논문을 살펴보면 x를 비교하는 것이 아니라 입실론을 비교하는 것을 볼 수 있다.
그것은 x가 아래와 같은 형식을 따르기 때문에

x_0를 다시 표현하면

이것을 M 식에 대입하면

따라서 p도 위와 같이 대입하여 KL식을 정리하면

이처럼 우리는 denoising matching term을 3가지로 표현할 수 있고,
논문의 저자는 x로 했을 때보다 epsilon으로 모델링 했을 때 더 좋은 성능을 보였다고 한다!

ELBO 3번째 접근

마지막으로 ELBo를 계산하기 위해서 Tweedie's Formula를 이용할 수 있다.

나는 처음 들었지만
diffusion 모델이 사실은 score based model과 같은 것이다.
(DDPM 논문에서도 증명이 되어 있다.)

우리가 앞에서 배운 formula를 Tweedie's Formula를 사용해서 다시 표현하면 이것이 score에 기반한 objective function으로 표현될 수 있다.

(즉, 두 모델은 사실상 같은 일을 하는 것이라는 거다!!)

Tweedie's Formula란?

gaussian distribution인 z라는 변수가 있을 때,
M_z|z의 함수를 예측하고 싶으면, z를 하나 샘플하고 거기에 분산 x score(gradient)를 더해주는 것이다.

p(z), 즉 likelihood가 높아지는 방향으로 분산 step만큼 조금씩 조금씩 업데이트 하면, 언젠가는 M_z로 수렴할 것이다 라는 것이 Tweedie's Formula다!!

이 공식을 아까 q(x_t|x_0)을 구할 때에 적용해보자

x_0를 위와 같이 표현할 수 있고, M_q에서 x_0자리에 이것을 대입하면,

따라서 아까 우리가 KL divergence식을 평균 또는 입실론으로 표현했는데 이번엔 score로써 표현할 수 있다.

이렇게 정리하면 이제 내 모델이 x가 높아지는 방향의 gradient가 있다고 가정했을 때 이 score가 그것을 예측하는 것을 학습하는 것과 다름없어졌다.
즉, diffusion에서 사용하는 epsilon을 이용한 objective function과,
score based model에서 사용하는 function이 같다는 것을 알 수 있다!!

정리

  • diffusion 구조
    epsilon 예측이나 x 예측을 통해 학습할 수 있다.

  • score generative model 구조
    분산 스텝만큼 score를 더해주다 보면 density가 높은 방향으로 모이게 될 것이다.
    real한 데이터를 뽑을 순 있겠지만 두번째 처럼 하게 되면 항상 같은 샘플링만 될 것이다. (다양하지 못함)

    그래서 세번째 처럼 가우시안에서 z를 뽑아서 매 스텝마다 조금씩 random하게 움직이게 함으로써 좀더 다양한 샘플링이 될 것이다.
    (이 샘플링을 Langevin Dynamics 샘플링이라고 한다!)

  • score based model with SDEs
    diffusion이 t step만큼 진행하는데 step을 점점 얇게 자르다 보면 연속적으로 데이터에서 노이즈로 가게 모델링할 수 있다.
    그리고 그것을 미분방정식으로 표현할 수 있게 된다.
    (이때 그냥 미분방정식이 아니라 stochastic differential equation을 할 수 있게 된다.)

    그렇게 foward와 reverse를 정의하면 score가 등장하는 것을 볼 수 있다.

즉 정리하면 이 3가지 구조가 같은 모델 family안에서 설명가능하다는 것이다!!!

profile
seung
이전 포스트

생성 모델의 흐름(variation bayes to VAE)

다음 포스트

Score based models

0개의 댓글