Understanding Diffusion Models - 브런치

논문 톺아보기 9 : 생성모델의 수식적 이해 이번에 리뷰할 논문은 'Understanding Diffusion Models: A Unified Perspective'입니다. 구체적으로 이미지 생성과 관련된 기본적인 사항에 대해서 수식적으로 잘 정리된 논문이어서 가져왔습니다. 이번 포스팅에서는 자체적으로 코드를 작성하는 것보다는 수식 전개를 살펴보고, 그 개념을 이해하는 것에 초점을 맞춰 진행해보고자 합니다. 또한 해당 논문의 수식과 관련된 코드를 찾아, 분석한 사항에 맞는지 확인하고 잘못된 점이 있다면 수정한 뒤 결과를 살펴보도록 하겠습니다. 본 논문은 이미지 생성 모델(Diffusion Models)의 백그라운드가 되는 여러 알고리즘 및 수식을 통합적으로 잘 정리한 메타분석 논문으로 아래와 같이 생성 모델 논문에서 기본 전제 조건으로 깔고 가는 개념에 대해 수식적으로 해설해 여러 개념들을 잘 합쳐놓은 논문이라 할 수 있습니다. 결국 생성모델은 위 그림과 같이 '정답분포 p(x) 및 그로부터 샘플링한 결과를 얻을 수 없는 상태에서, AI모델(neural net)을 이용해 실제 p분포에 근사시킬 수 있는 q분포를 이용하는 것'이며, 구체적으로 어떻게 p분포와 근사한 q분포를 생성할 것인가에 초점을 맞추고 있습니다. 제가 이 부분을 읽을 때마다 드는 생각 중 하나는 바로 '무한 원숭이 정리'입니다. 무한 원숭이 정리란 키보드를 칠 수 있는 무한한 원숭이가 있다고 가정했을 때 이들이 임의로 한 글자씩 입력하면 당연히 그 출력 결과의 대부분은 제대로 된 단어나 구문조차 형성하지 못하겠지만, 그 중 일부는 셰익스피어의 햄릿과 같은 결과물을 재현내낼 수도 있다는 것입니다. 간단히 표현하자면 우리가 완전히 랜덤한 확률로 키보드의 특정 버튼을 클릭하는 프로그램을 만들고, 이를 무한한 시간동안 실행시킨다면 그 중에는 우리가 잘 아는 유명한 드라마나 영화의 대본이 나올 수도 있다는 것이죠. 하지만 이러한 가정은 수학적으로는 모르나, 실질적으로는 의미가 없습니다. '무한'이라는 개념을 현실로 가져와 사용할 수가 없기 때문이죠. 생성 모델은 이러한 '무한'의 개념을 '유한'하게 만드는 방법을 다루고 있습니다. 완벽하게 랜덤한 것이 아닌 앞뒤의 글자를 확인하거나, 이전 페이지에 등장한 인물의 이름이나 특징을 기억하거나 함으로써 말이죠. 본 논문에서는 이를 가능하게 하는 수식적 근거를 제시합니다. 하나씩 읽어보며 진행해보도록 하겠습니다. 생성 모델을 이해하는 첫 단계는 바로 생성 모델의 핵심 개념을 이해하는 것입니다. Diffusion 기반의 생성 모델을 포함한 모든 종류의 생성 모델은 입력을 통해 출력이 결정되며, 이를 위해 입력한 값을 낮은 차원의 Latent vector로 embedding했다가 다시 원상복구시키는 과정을 학습함으로써 적절한 Latent Vector를 입력하는 것만으로 우리가 원하는 결과물을 생성할 수 있게 됩니다. 이러한 과정은 기존의 VAE(Variational AutoEncoder)와 동일한 방식입니다. 모델의 구조 자체는 encoder와 decoder로 구성된 AE(AutoEncoder)와 동일하지만, Encoder를 통해 압축된 벡터들의 분포까지도 고려한다는 점에서 VAE는 차별성을 가지며 생성 모델로서의 핵심 개념을 포함하고 있다고 할 수 있습니다. 물론 세상에 존재하는 모든 종류와 형상의 데이터를 확보할 수 없고, 확보한 일부의 데이터만으로 전체 데이터의 특징을 반영하는 Encoding 규칙을 구현하는 것은 무척 어려운 일입니다. 본 포스팅에서는 그러면 어떻게 관측한 데이터(학습가능한 데이터)로부터 전체 데이터의 분포를 근사할 수 있는지에 대해 알아볼 것입니다. 기본적으로 생성모델은 실제 분포 p(x)와 근사한 q(x)를 표현하는 잠재 구조(latent structure)를 구현하고, 해당 잠재구조를 통해 특정 입력이 들어왔을 때 의도한 형태로 재생성해내는 것을 목표로 하고 있습니다. 재생성하는 과정은 우리가 Decoder라고 명명하는 인공지능 모델에 할당한다면, 우리에게 남은 것은 근사한 분포 q(x)를 만들어내는 방법에 대한 고민이며, 그 첫 시작은 ELBO(Evidence Lower Bound)입니다. 위 그림에서와 같이 우리는 관측된 데이터와 잠재 분포를 Joint distribution을 통해 상상할 수 있고(근사시킬 수 있고), 관찰된 데이터의 가능도 p(x)를 최대화함으로서 생성모델의 구현이 가능하도록 할 수 있습니다. 이때 관찰된 데이터로부터 joint distribution p(x, z)를 통해 가능도 p(x)를 만들어내는 방법은 두 가지가 있으며, 위 그림의 두 가지 수식에 해당합니다. 만약 p(x)를 이렇게 한 번에 구해서 최대화할 수 있다면 괜찮겠지만 현실적 한계로 인해 위 공식들 중 어느 하나를 이용해 p(x)를 구현하는 것은 한계가 있습니다. 구체적으로 Equation (1)을 사용하려면 복잡하기 그지없는 모델 내부의 내재된 변수들(Latent Variables) 간 가능한 모든 조합에 대해 적분해야 하나 현실적으로 불가능하며, Equation (2)를 사용하려면 Ground Truth Latent Encoder인 P(z|x)가 필요하지만 이를 사전에 알 수 있는 방법이 없기 때문입니다. 이로 인해 위 두 가지 공식을 이용해 관찰된 데이터(Evidence)로부터 실제 데이터를 근사해낼 수 있는 공식인 ELBO(Evidence Lower Bound)를 구현하며, 도출하는 과정은 아래와 같습니다. 위 그림은 Equation (1)에서부터 ELBO를 도출하고, 왜 ELBO를 최대화하는 것이 실제 p(x)와 유사해질 수 있는지를 표현하는 과정입니다. 위 그림에서 Jensen 부등식이란 다음과 같습니다. 즉, 위 도출 과정에서 log 함수가 볼록함수 f가 되어 적용된 결과로 볼 수 있습니다. 다시 정리해서 보자면 ELBO의 수식은 다음과 같습니다. 하지만 이렇게 도출한 수식은 ELBO에 대한 수식적 이해 정도면 몰라도 직관적으로 왜 ELBO를 최대화해야하고, 정말 ELBO를 최대화하면 실제와 같아지는지에 대한 납득이 어렵습니다. 다시말해 관찰된 데이터와 ELBO 사이의 관계가 명확하게 와닿지 않습니다. 때문에 이번엔 Equation (2)까지 추가로 적용해 조금 다른 방법으로 공식을 유도하면 아래와 같습니다. 위 수식 전개 과정을 보면 알 수 있듯이 Equation (2)를 적용하고, 이후 KL Divergence의 정의를 이용함으로써 보다 직관적으로 ELBO 값이 왜 항상 본래의 값 이하로 나타나는지 알 수 있습니다. 참고로 KL Divergence는 두 확률분포 간의 차이를 의미하며 Cross Entropy와 Entropy의 결합으로 이뤄집니다. 구체적으로 KL Divergence = Cross Entropy - Entropy와 같은 식이 성립하며, 이를 통해 p분포와 q분포 간의 거리를 측정하게 됩니다[3]. Entropy와 Cross Entropy의 수식은 위 그림과 같습니다. 표기 방식이 달라 모호해보이지만, ground truth인 p분포끼리의 확률분포, 즉 H(p, p)가 곧 Entropy인 셈입니다. 따라서 Cross Entropy와 Entropy는 다음과 같은 식이 성립하며, 두 값의 차이는 p와 q 분포의 차이에 의해서 발생하는 것으로 해석할 수 있게 됩니다. 때문에 두 확률분포가 동일하다면 0이 되고 차이가 커질수록 값이 커지게 됩니다. 즉 최소값이 0인 것이죠. 이러한 점을 고려하면 'log P분포는 ELBO+KL Divergence 값이며, 이는 항상 ELBO 이상이다'라고 말할 수 있습니다. 또한 데이터의 likelihood가 상수라는 것을 고려한다면, ELBO를 최대화되는 경우는 KD Divergence가 최소가 되는 경우와 일맥상통하는 의미가 됩니다. 때문에 우리는 최적화 목표를 ELBO를 최대화시키는 것 혹은 KL Divergence를 최소화시키는 것이라고 정리할 수 있으며, 목적함수가 되는 것입니다. VAE의 공식에서는 ELBO를 바로 최대화시킵니다. VAE는 자기 자신을 압축했다가 다시 재구성하는 AutoEncoder 계열의 알고리즘이며, ELBO와는 아래와 같이 연결고리가 있습니다. 이때 위 수식 전개의 마지막에서 reconstruction term과 prior matching term은 각각 decoder와 encoder를 최적화하는 것을 의미합니다. 구체적으로 reconstruction term은 encoder를 통해 압축된 latent vector z가 decoder를 통해 우리가 학습한 데이터 z로 재구성하는 것을 의미합니다. prior matching term은 우리가 encoder로 압축한 (learned) variational distribution(qΦ(z|x))가 latent variable distribution(p(z))와 얼마나 유사한지를 측정하게 되며, 이를 통해 모델이 잠재 변수에 대한 불확실성과 다양성을 유지하도록 유도해 Dirac delta function으로 붕괴되지 않도록 합니다. 참고로 위 수식에서 θ와 ϕ는 각각 Decoder와 Encoder의 파라미터를 의미하며, pθ(x∣z)는 잠재변수 z로부터 원본 데이터 x를 재구성하는 조건부 확률분포를, qϕ(z∣x)는 주어진 입력 x에 대해 가능한 잠재 변수 z의 분포를 나타내며, p(x)는 잠재변수에 대한 사전분포로 우리가 잠재공간에 대해 사전에 가진 믿음을 의미한다 할 수 있습니다. 또한 우리가 정의한 목적함수는 몬테 카를로 추정에 의해 근사될 수 있습니다. 몬테 카를로 추정은 잠재변수 z의 샘플들을 qϕ(z∣x)로부터 추출해 기대값을 추정하는 방법입니다. 이렇게 근사된 식은 다음과 같습니다. 하지만 이때 몬테카를로 추정으로 샘플링한 z(l)은 확률적으로 생성되기 때문에 미분 불가능하며, 이는 머신러닝의 목적함수로서의 기능을 하지 못하는 문제를 발생시킵니다. 이러한 문제를 해결하기 위해 VAE에서는 미분이 가능한 특정 분포(ex. 다변량 가우시안)대로 모델링하도록 설계를 진행하고 재매개변수화 기법(Reparameterazation Trick)을 적용합니다. 구체적으로 VAE에 ELBO를 적용함에 있어 qϕ(z∣x)는 대각 공분산을 가진