Flow Matching For Generative Modeling
Flow Matching For Generative Modeling
Introduction
이 논문에서는 Flow Matching, FM 이라는 개념을 통해, 기존 Continuous Nomalizing Flows, CNFs (우리가 NICE, realNVP 에서 봤던거, 대신 discrete 이 아니라 continuous임) 를 어떻게 학습시키는지 소개한다.
..내가 애니메이션 공부를 아직 안해서 기하학적인 개념을 다 글로만 설명하려고 했더니.. 비문학이 된 느낌이다.
..일단 초안 상태이니, 계속해서 수정해가도록 하겠다…
Continous Normalizing Flows
먼저 CNFs 에 대해 알아보자면, 데이터 공간 $\mathbb{R}^{d}$ 에 속한 각 데이터 포인트 $x = (x^{1},…,x^{d}) \in \mathbb{R}^{d}$ 의 정의로 시작한다.
그리고 $t = [0,1]$ 에 대해서, $p_{t}(.) : t \times \mathbb{R}^{d} \to \mathbb{R}_{> 0}^{d}$ 로 정의되며, 시간 $t$ 에 따른 확률 밀도 함수이다. 이 논문에선 $p_{0}$ 을 완전한 노이즈 상태, $\mathcal{N}(x;0\mathbf{I})$ 로 정의하고, $p_{1}$ 을 모델의 분포로 나타낸다. 이때 모델의 분포는 실제 데이터 분포 $q(x)$ 를 근사한 것이다. 그래서 표준 정규 분포에서 모델 분포로 분포 자체가 변환되는 각 과정에 대해서, $t$ 시점에서의 분포 모양을 나타낸다고 볼 수 있다.
$v_{t}(.): t \times \mathbb{R}^{d} \to \mathbb{R}^{d}$ 는 시간 $t$ 에 의존한 Vector Field, VF로, 음.. Score based model 에서 본 확률 밀도의 최댓값을 가리키는 벡터 값들, Scores 처럼, 각 데이터들을 어떠한 방향으로 이동하게끔 하는 벡터 값들이다. 그러니까.. 각 지점에서 데이터 포인트들에게 작용하는 힘이 정의되어 있다. 다시말해 $t$ 에 따라서, 데이터 공간의 각 위치에서 데이터 포인트들이 어떻게 이동해야 하는지를 나타낸다는 것이다.
근데 여기서 아무리 혼자 생각해봐도, VF 가 $t$ 에 의존하는 이유를 명확히 모르겠는데, 뭐랄까.. 도저히 납득할 수 없달까. 그냥 직관적으로 생각해보면 예를 들어.. VF의 특정 위치에 임의의 두 데이터 포인트 $A,B$ 가 각각 시점 $t$, $t+1$ 에 도착한다고 해보자.. 그러니까 특정 위치에 $t$ 시점에서 $A$ 가 먼저 도착하고, 다음 $t+1$ 시점에선 $A$ 는 해당 VF 값에 따라 이동하고, 다음 $B$ 가 도착한다고 해보면… VF 가 $t$ 에 의존하지 않는 경우, 즉, 고정된 경우엔 이 두 $A,B$ 는 서로 다른 곳에서 해당 위치에 도착했더라도, 그 이후 경로는 같을 것이다. 따라서 $t$ 에 따라 이 VF 를 변화하게끔 해서, $A,B$ 가 다른 시점에서 같은 위치에 도착하더라도, 도착지가 다르게끔 해야하는데.. 뭐랄까.. 뭔가 납득이 되지 않는다. 그래서 이 부분에 대한 설명은 일단 여기선 킵하고, …명확한 이유를 찾으면 여따 추가하겠다..
아무튼 다시 돌아와서 이러한 VF 를 통해 시간 $t$ 에 의존한 Flow $\phi_{t}: t \times \mathbb{R}^{d} \to \mathbb{R}^{d}$ 를 정의하는데, 이는 각 데이터들이 $t$ 에 따라 이동하는 실제 궤적을 나타내며, $\phi_{0}(x)$ 는 초기 위치, $\phi_{1}(x)$ 는 초기 위치에서 최종적으로 도착한 위치가 되겠다. 그래서 이는 다음과 같이 ODE로 정의된다.
\[\begin{align} \frac{d}{dt}\phi_{t}(x) &= v_{t}(\phi_{t}(x)) \\ \phi_{0}(x) &= x \end{align}\]그래서 이러한 Flow는 여태 다른 생성형 모델에서 한 것 처럼 완전한 노이즈 $p_{0}$ 에서 임의의 복잡한 모델의 분포 $p_{1}$ 로 변환하는데 사용되는데, 다음과 같이 change of variables 를 통해 계산된다.
\[\begin{align} p_{t} &= [\phi_{t}]_{*}p_{0} \\ [\phi_{t}]_{*}p_{0}(x) &= p_{0}(\phi_{t}^{-1}(x))\det\left[\frac{\partial \phi_{t}^{-1}}{\partial x}(x)\right] \end{align}\]여기서 각 데이터 $x$ 의 이동 궤적을 통해 어떻게 모델의 분포 변해가는 과정에서 $p_{t}$ 와 연결되는지, 살짝 애매한데, 직관적으로 보면, 2차원 xy 평면에서, 중앙에 모인 노이즈 형태의 가우시안 분포를 나타내는 많은 데이터 포인트들이 있을 때, 이 데이터 포인트들이 각각 $\phi$ 를 따라 개별적으로 이동한다. 이렇게 이동해서 최종 도착지, $t=1$ 시점에서 이 데이터 포인트들의 구성이, 모델의 분포의 형태를 갖춘다고 볼 수 있다.
…아.. 이게 말로 설명하려니 좀 거시기한데, 아무래도 애니메이션 같은걸 한 번 공부해서 후에 더 발전해서 돌아오겠다…
대충 요런 느낌..? 각 기마병(=데이터 포인트) 들이 이동하는데, 그 방향은 적군을 향해! (= VF) 그리고 각 기마병들의 위치는 $\phi_{t}$ 로… 그래서 저 기마병들이 이루는 모양이 $p_{t}$ 가 되겠다… 아무튼 반지의 제왕식 비유법이다. 그냥.. 각 개념이 뭔 느낌인지 참고만 하자…
아무튼 위와 같이 $\phi_{t}$ 를 통해 노이즈 분포 $p_{0}$ 에서 $p_{t}$ 를 계산할 수 있으면, $\phi_{t}$ 를 $t$ 에 대해 미분한 값인 VF $v_{t}$ 가 $p_{t}$ 를 생성한다고 말할 수 있다. 이를 확인하는 방법은 다음을 만족하는 것과 같다.
\[\frac{d}{dt}p_{t}(x) + \text{div}(p_{t}(x)v_{t}(x)) = 0\]이때 $\text{div} = \sum_{i=1}^{d} \frac{\partial}{\partial x^{i}}$ 이고, 이 식 자체를 Continuity Equation 이라고 한다.
Flow Matching, FM
그래서 다시 상황을 정리해보면, 우리가 알고자 하는 데이터의 분포 $q$ 가 존재하고, 이를 근사한 모델의 분포 $p_{1} \approx q$ 이 있다. 그래서 우리는 $p_{0} = p = \mathcal{N}(x|0,\mathbf{I})$ 의 노이즈 부터 시작해서, 요 모델의 분포 $p_{1}$ 로 $t=[0,1]$ 에 따라 변환하는데, 이때 $t$ 시점에서의 $p_{t}$ 는 VF $u_{t}$ 에 의해 생성된다. 그러면 다시 정리하면, 우리가 이 VF $u_{t}$ 만 알고 있으면, $p_{t}$ 를 알 수 있다. 이러한 Flow Matching, FM loss 는 다음과 같이 정의된다.
\[L_{\text{FM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1),x \sim p_{t}(x)}(\|v_{t}(x;\theta) - u_{t}(x)\|^{2})\]즉, $t=[0,1]$ 에 따른 VF $u_{t}$ 를 학습하면 된다.
Conditional Flow Matching, CFM
…그래서 $u_{t}$ 가 뭔데, 더군다나 $x$ 를 뽑아오는 $p_{t}$ 도 도대체 어떻게 생겼는지 알지도 못한다.
논문 저자 형님들은 다음과 같이 이 부분을 다뤘다.
일단 실제 데이터 분포에서 데이터 샘플을 하나 뽑는다. $x_{1} \sim q_{1}(x)$. 그러면 $p_{t}(x|x_{1})$ 의 형태로, 조건부 확률 밀도를 나타내어, $p_{0}(x|x_{1}) = p_{0}(x) = p(x)$, 그리고 $p_{1}(x|x_{1}) = \mathcal{N}(x|x_{1}, \sigma^{2}\mathbf{I})$ 로, $t=1$ 에선 해당 데이터 샘플 $x_{1}$ 을 기준으로, 아주 작은 양수의 표준 편차 $\sigma$ 로 작게 가우시안으로 퍼진 형태로 $p_{1}(x|x_{1})$ 을 나타낸다.
그러면 기존 $p_{t}$ 는 다음과 같이 $x_{1}$ 에 대해 주변화하여 구할 수 있다.
\[\begin{align} p_{t}(x) &= \int p_{t}(x|x_{1})q(x_{1})\,dx_{1} \\ p_{1}(x) &= \int p_{1}(x|x_{1})q(x_{1})\,dx_{1} \approx q(x) \end{align}\]물론 intractable 하다.
VF 도 마찬가지로, 이러한 샘플 $x_{1}$ 을 조건부로 두어 $u_{t}(x|x_{1})$ 로 나타낼 수 있는데, 이 경우, 이러한 conditional VF는 기존 $p_{t}(x)$ 가 아닌 $p_{t}(x|x_{1})$ 을 생성한다.
물론 conditional VF 도 위와 같이 $x_{1}$ 에 대해 주변화하여 기존 VF를 구할 수 있다.
\[\begin{align} u_{t}(x) &= \int u_{t}(x|x_{1})\frac{p_{t}(x|x_{1})q(x_{1})}{\int p_{t}(x|x_{1})q(x_{1})dx_{1}}\,dx_{1} \\ &= \int u_{t}(x|x_{1})\frac{p_{t}(x|x_{1})q(x_{1})}{p_{t}(x)}\,dx_{1} \end{align}\]..마찬가지로 intractable 하다.
근데 여기서 논문의 저자 형님들이 서커스를 진행한다.
\[L_{\text{FM}}(\theta) = \mathbb{E}_{t \sim \mathcal{U}(0,1),x \sim p_{t}(x)}(\|v_{t}(x;\theta) - u_{t}(x)\|^{2})\]기존 FM 손실 함수 부분을 전개해보자.
\[\|v_{t}(x;\theta) - u_{t}(x)\|^{2} = \|v_{t}(x;\theta)\|^{2} - 2\langle v_{t}(x;\theta), u_{t}(x)\rangle + \|u_{t}(x)\|^{2}\]여기서 저 앞의 2개의 항에서 서커스를 다음과 같이 진행한다.
\[\begin{align} \mathbb{E}_{p_{t}(x)}(\|v_{t}(x;\theta)\|^{2}) &= \int \|v_{t}(x;\theta)\|^{2}p_{t}(x)\,dx \\ &= \int \|v_{t}(x;\theta)\|^{2}p_{t}(x|x_{1})q(x_{1})\,dx_{1}\,dx \\ &= \mathbb{E}_{q(x_{1}), p_{t}(x|x_{1})}(\|v_{t}(x;\theta)\|^{2}) \end{align}\] \[\begin{align} \mathbb{E}_{p_{t}(x)}(\langle v_{t}(x;\theta), u_{t}(x)\rangle) &= \int \left\langle v_{t}(x;\theta), \frac{\int u_{t}(x|x_{1})p_{t}(x|x_{1})q(x_{1})\,dx_{1}}{p_{t}(x)}\right\rangle p_{t}(x)\,dx \\ &= \int \left\langle v_{t}(x;\theta), \int u_{t}(x|x_{1})p_{t}(x|x_{1})q(x_{1})\,dx_{1}\right\rangle \,dx \\ &= \int \langle v_{t}(x;\theta), u_{t}(x|x_{1})\rangle p_{t}(x|x_{1})q(x_{1})\,dx_{1}\,dx \\ &= \mathbb{E}_{q(x_{1}), p_{t}(x|x_{1})}(\langle v_{t}(x;\theta), u_{t}(x|x_{1})\rangle) \end{align}\]마지막 항 $|u_{t}(x)|^{2}$ 이거는 모델 파라미터 $\theta$ 에 의존하지 않는다. 이를 다시 작성하면 다음과 같은데,
\[L_{\text{FM}}(\theta) = \mathbb{E}_{t,q(x_{1}), p_{t}(x|x_{1})}(\|v_{t}(x;\theta)\|^{2} - 2\langle v_{t}(x;\theta), u_{t}(x|x_{1})\rangle) + \mathbb{E}_{t, p_{t}(x)}(\|u_{t}(x)\|^{2})\]여기서 우리의 조건부 VF $u_{t}(x|x_{1})$ 을 대입해도, 마찬가지로 모델 파라미터 $\theta$ 에 의존하지 않는 상수로 취급되어, 결국 다음과 같은 손실 함수도 최적화했을 때, 기존 FM 손실과 같은 파라미터 $\theta^{*}$ 을 가진다.
\[L_{\text{CFM}}(\theta) = \mathbb{E}_{t,q(x_{1}), p_{t}(x|x_{1})}(\|v_{t}(x;\theta) - u_{t}(x|x_{1})\|^{2})\]이를 Conditional Flow Matching, CFM loss 라고 한다.
즉, 다시 말해서 기존 FM loss 를 최적화하는 것과 CFM loss 를 최적화하는 것은 같은 결과를 가진다는 것이다. 그래서 우리는 이 conditional VF 를 모델이 학습하게끔 하여, $p_{t}(x)$ 를 학습할 수 있다!!!!!!
Conditional Probability Path and VF
그러면 이제 conditional probability path 와 conditional vector field 를 어떻게 정의하는지만 알면 된다.
논문에선 일단 이 조건부 확률 밀도를 조건부 가우시안의 형태로 나타낸다.
\[p_{t}(x|x_{1}) = \mathcal{N}(x| \mu_{t}(x_{1}), \sigma_{t}(x_{1})^{2}\mathbf{I})\]이때 $\mu_{t}(.)$ 는 $t$ 에 따른 가우시안 분포의 평균, $\sigma_{t}$ 는 $t$ 에 따른 스칼라 값의 표준 편차이다. $t=0$ 에서의 두 값을 $\mu_{0}(x_{1}) = 0, \sigma_{0}(x_{1}) = 1$ 로 설정하여, $t=0$ 에서 표준 정규 분포 형태의 완전한 노이즈로 나타내고, $\mu_{1}(x_{1}) = x_{1}, \sigma_{1}(x_{1}) = \sigma_{\min}$ 으로 나타내어, $t=1$ 에서의 해당 조건부 확률 밀도는 $x_{1}$ 를 중심으로 가우시안의 형태로 분포된 모양을 가진다.
이제 VF 를 정해야 한다. 물론 수많은 VF 들을 정의할 수 있지만,
위 사진과 같이, 전체적인 분포 모양이 유지되면서 데이터가 회전하기만 하는 것처럼, 쓰잘데기 없는 계산을 최소화하기 위해서, 다음과 같은 간단한 형태로 conditional flow 를 정의한다.
\[\psi_{t}(x_{0}) = \sigma_{t}(x_{1})x_{0} + \mu_{t}(x_{1})\]여기서 $x_{0} \sim \mathcal{N}(x|0,\mathbf{I})$ 로 초기 위치이고, $\psi_{t}$ 는 이러한 초기 위치를 평균 $\mu_{t}(x_{1})$, 표준편차 $\sigma_{t}(x_{1})$ 인 가우시안 분포로 매핑하는 Affine 함수라고 볼 수 있다.
즉, 우리가 맨처음에 봤던 CNFs 에서 처럼, 다음과 같다.
\[\begin{align} p_{t}(x|x_{1}) &= [\psi_{t}]_{*}p_{0}(x) \\ [\psi_{t}]_{*}p_{0}(x) &= p_{0}(\psi_{t}^{-1}(x))\det\left[\frac{\partial \psi_{t}^{-1}}{\partial x}(x)\right] \end{align}\] \[\frac{d}{dt}\psi_{t}(x) = u_{t}(\psi_{t}(x)|x_{1})\]그럼 이제 $x \sim p_{t}(x|x_{1})$ 를 $\psi_{t}(x_{0}) = \sigma_{t}(x_{1})x_{0} + \mu_{t}(x_{1})$ 로 Reparameterize 하고, 이를 앞선 CFM loss에 다시 대입하면, 다음과 같이 나타낼 수 있다.
\[L_{\text{CFM}}(\theta) = \mathbb{E}_{t,q(x_{1}), p_{0}(x_{0})}\left(\left\|v_{t}(\psi_{t}(x_{0});\theta) - \frac{d}{dt}\psi_{t}(x_{0})\right\|^{2}\right)\]$\psi_{t}$ 를 통해 $t$ 에서의 각 데이터 포인트의 위치를 Affine 함수로 정의했으니, $t$ 에 대한 미분 값인 $u_{t}$ 를 쉽게 Closed-form 으로 구할 수 있다. 다음과 같이 계산한다.
일단 $\psi_{t}$ 를 $t$ 에 대해 미분하면 다음과 같다.
\[\frac{d}{dt}\psi_{t}(x) = \sigma_{t}'(x_{1})x + \mu_{t}'(x_{1}) = u_{t}(\psi_{t}(x)|x_{1})\]저기서 $x$ 는 $x=\psi_{t}^{-1}(\psi_{t}(x))$ 로 나타낼 수 있다. $\psi_{t}^{-1}$ 는 Affine 함수니까 아주 야무지게 쉽게 다음과 같이 계산된다.
\[\psi_{t}^{-1}(x) = \frac{x - \mu_{t}(x_{1})}{\sigma_{t}(x_{1})}\]이를 저 VF 에 때려 박으면, 다음과 같다.
\[u_{t}(\psi_{t}(x)|x_{1}) = \frac{\sigma_{t}'(x_{1})}{\sigma_{t}(x)}(\psi_{t}(x) - \mu_{t}(x_{1})) + \mu_{t}'(x_{1})\]즉, 최종 CFM loss 는 다음과 같다.
\[\therefore L_{\text{CFM}}(\theta) = \mathbb{E}_{t,q(x_{1}), p_{0}(x_{0})}\left(\left\|v_{t}(\psi_{t}(x_{0});\theta) - \left(\frac{\sigma_{t}'(x_{1})}{\sigma_{t}(x)}(\psi_{t}(x) - \mu_{t}(x_{1})) + \mu_{t}'(x_{1})\right)\right\|^{2}\right)\]Diffusion conditional VFs
이 논문 형님들은 우리가 본 Diffusion Processes 도 이러한 FM 의 특이 케이스일 뿐이라고 한다. 왜 그런지 알아보자.
..일단 전에 Score-based model 논문에서, Diffusion models 을 SDE 로 정의를 했었다. 막 각 데이터 포인트들이 $p_{1}$ 로 이동하되, 그 궤적이 노이즈 항이 섞여 약간 지그재그 형식으로 뺑 돌아가거나 하는 느낌이었는데, 이러한 복잡한 경로라 하더라도, 각 시점 $t$ 에서 데이터 포인트들로 이루어진 확률 분포의 모양 $p_{t}$ 를 유지하면서, 각 데이터 포인트들이 지그재그 발광을 하는게 아닌 부드러운 곡선의 궤적으로 움직이게 하는 결정론적인 단 하나의 ODE로 나타낼 수 있다고 했다.
그러니까, Diffusion SDE는 다음과 같다.
\[dx = f_{t}dt + g_{t}dw\]$f_{t}$ 는 Drift 계수, $g_{t}$ 는 Diffusion 계수, $w$ 는 표준 브라운 운동인데, 자세한건 위 블로그를 보자..
아무튼.. 저 SDE 식에 대해서, $p_{t}(x_{t})$ 는 다음과 같이 Fokker-Planck 방정식을 통해 나타낼 수 있다. (위키피디아 참고…)
\[\frac{dp_{t}}{dt} = -\text{div}(f_{t}p_{t}) + \frac{g_{t}^{2}}{2}\Delta p_{t}\]이때 $\Delta = \text{div}\nabla$ 는 $y$ 에 대한 Laplace 연산이다. 이제 저 식을 전개해보면…
\[\begin{align} \frac{dp_{t}}{dt} &= -\text{div}\left(f_{t}p_{t} - \frac{g_{t}^{2}}{2}\frac{\nabla p_{t}}{p_{t}}p_{t}\right) \\ &= -\text{div}\left((f_{t}-\frac{g_{t}^{2}}{2}\frac{\nabla p_{t}}{p_{t}})p_{t}\right) \\ &= -\text{div}(u_{t}p_{t}) \end{align}\]우리가 맨처음에 봤던 Continuity Equation 에 따라, 저 $u_{t}$ 는 $p_{t}$ 를 생성하는 VF 라고 할 수 있다.
즉, Diffusion VF 는 다음과 같다.
\[\begin{align} u_{t} &= f_{t}-\frac{g_{t}^{2}}{2}\frac{\nabla p_{t}}{p_{t}} \\ &= f_{t}-\frac{g_{t}^{2}}{2}\nabla \log p_{t} \quad (\nabla \log p_{t} = \frac{\nabla p_{t}}{p_{t}}) \end{align}\]아무튼 이렇게 Diffusion VF 는… 어쨌든 우리가 모르는 $p_{t}$ 를 다루고 있기 때문에, 이를 근사해야만 계산할 수 있다.
…딱 봐도 우리가 이 논문에서 본 VF $u_{t}$ 의 식과 확연히 다른데, 놀랍게도 Diffusion VF 에서 $p_{t}(x)$ 를 이 논문의 FM 방식 $p_{t}(x|x_{1})$ 로 바꾸면, 우리가 아는 그 형태가 나온다. 한번 알아보자.
Variance Exploding VF
먼저 Variance Exploding, VE 인 score-based models 의 경우, SDE 는 다음과 같다.
\[\begin{align} dx &= \sqrt{\frac{d}{dt}\sigma_{t}^{2}}dw \\ f_{t} &= 0 \\ g_{t} &= \sqrt{\frac{d}{dt}\sigma_{t}^{2}} = \sqrt{2\sigma_{t}\sigma_{t}'} \end{align}\]또한 노이즈에서 데이터분포로 변해가는 각 과정의 데이터의 위치, $x_{t}$ 를 다음과 같이 정의한다.
\[x_{t} \sim p_{t}(x_{t}|x_{1}) = \mathcal{N}(x_{t}|x_{1},\sigma_{t}^{2}\mathbf{I})\]이때 $\sigma_{0} \gg 1$, $\sigma_{1}=0$ 인 $t$ 에 따라 점차 감소하는 형태이다. 이 식에 따라, 우리가 $\mu_{t}(x_{1}) = x_{1}$ , $\sigma_{t}(x_{1}) = \sigma_{t}$ 로 나타낼 수 있는데, 그러면 conditional VF는 우리가 살펴본 정의에 따라 다음과 같다.
\[u_{t}(\psi_{t}(x_{0})|x_{1}) = - \frac{\sigma_{1-t}'}{\sigma_{1-t}}(x-x_{1})\]앞서 살펴본 것처럼, FM 을 통해 나타낸 VF가 만들어내는 데이터 포인트들의 궤적은, 저 앞에서 구한 Diffusion VF 에 각각 $f_{t}, g_{t}$ 와 조건부 가우시안 분포 값을 대입해서 보면, 정확히 같은 결과가 나온다.
Variance Preserving VF
Variance Preserving, VP, 즉 DDPM 의 경우도 알아보자.
SDE 와 $p_{t}$ 는 다음과 같이 정의된다.
\[\begin{align} dx &= - \frac{T'(t)}{2}xdt + \sqrt{T'(t)}dw \\ f_{t} &= - \frac{T'(t)}{2}x \\ g_{t} &= \sqrt{T'(t)} \end{align}\] \[\begin{align} x_{t} \sim p_{t}(x_{t}|x_{1}) &= \mathcal{N}(x_{t}|\alpha_{t} x_{1}, (1-\alpha_{t}^{2})\mathbf{I}) \\ \alpha_{t} &= \exp\left(-\frac{1}{2}T(t)\right) \\ T(t) &= \int_{0}^{t}\beta(s)\,ds \end{align}\]마찬가지로 각 파라미터에 대한 설명은 여기선 귀찮으니.. 다시 저 위 블로그 글을 참고하자.
그래서 아무튼 $\mu_{t}(x_{1}) = \alpha_{t} x_{1}$ , $\sigma_{t}(x_{1}) = \sqrt{1-\alpha_{t}^{2}}$ 에 대해 마찬가지로 conditional VF 를 구하면 다음과 같다.
\[\begin{align} u_{t}(x|x_{1}) &= \frac{\alpha_{t}'}{1-\alpha_{t}^{2}}(x_{1} - \alpha_{t}x) \\ &= - \frac{T'(t)}{2}\left(\frac{\exp\left(-\frac{1}{2}T(t)\right)x_{1} - \exp(-T(t))}{1 - \exp(-T(t))}\right) \end{align}\]역시나 Diffusion VF 에 $f_{t}, g_{t}$ 와 조건부 가우시안 분포 값을 대입해서 보면, 같은 결과가 나온다.
Compare score-based and FM-based for diffusion
아무튼 FM 의 방식으로도 기존 Diffusion Processes 와 정확히 같은 궤적을 그릴 수 있음을 보였는데, 그러면 어느 방식이 더 좋은가?
논문의 저자들은 Score matching 을 사용하는 것보다, Flow matching 을 사용하는게 더 안정적이라고 한다. 복잡하게 말고 간단하게 생각해도, 신경망이 Score 를 억지로 흉내내도록 하는 것 보단, 1차 식으로 딱 떨어지는 이 논문의 FM 을 통해 MSE 회귀 문제로 만들어 학습하는게 더 쉽고 안정적으로 보인다.
또한 DDPM 과 같은 diffusion VP 에선 $1 - \exp\left(-T(t)\right) \quad (T(t) = \int_{0}^{t}\beta(s)\,ds)$ 의 형태로 분산을 계산하였다.
여기서 완전한 노이즈를 나타내기 위해서는 분산의 값이 $1$ 이 되어야 하는데, 그러니 위해서는 $\exp\left(-T(t)\right) = 0$ 을 만족해야 한다. 즉, $T(t) \to \infty$ 를 만족해야 한다. 이는 $t \to \infty$ 를 의미하는 것과 같다. 따라서, $t=[0,1]$ 로 제한된 상황에서 DDPM 은 완전한 노이즈를 나타낼 수 없음을 의미하는데, 물론 반대로 FM 은 이러한 허점을 야무지게 보완했다.
Optimal Transport Conditional VFs
논문 저자 형님들은 $p_{t}$ 의 평균과 표준편차를 $t$ 에 따라 선형적으로 증가하는 직선의 간단한 형태로 다음과 같이 정의해봤다.
\[\begin{align} \mu_{t}(x_{1}) &= tx_{1} \\ \sigma_{t}(x_{1}) &= 1 - (1-\sigma_{\min})t \end{align}\]이렇게 정의하면 conditional VF 는 다음과 같다.
\[\begin{align} u_{t}(x|x_{1}) &= \frac{1-\sigma_{\min}}{1 - (1-\sigma_{\min})t}(x - tx_{1}) + x_{1} \\ &= \frac{x_{1} - (1-\sigma_{\min})x}{1 - (1-\sigma_{\min})t} \end{align}\]마찬가지로 $\psi_{t}$ 도 다음과 같다.
\[\psi_{t}(x) = (1-(1-\sigma_{\min})t)x + tx_{1}\]이제 요 두 식을 합체시키면 다음과 같은 CFM loss 가 나온다.
\[L_{\text{CFM}}(\theta) = \mathbb{E}_{t,q(x_{1}), p_{0}(x_{0})}\left(\left\|v_{t}(\psi_{t}(x_{0});\theta) - (x_{1} - (1-\sigma_{\min})x_{0})\right\|^{2}\right)\]평균과 표준 편차를 직선으로 선형적으로 증가하는 형태로 두니, 그냥 처음 위치 $x_{0}$ 에서 최종 위치 $x_{1}$ 을 직성 방햔만으로만 이동하게끔 모델을 학습할 수 있게 되었다.
이는 McCann 이라는 형님이 $p_{0}$ 에서 $p_{1}$ 로 가장 적은 비용으로 움직이도록 할 때, 그 중간의 $p_{t}$ 는 $p_{0}, p_{1}$ 요 두개를 그냥 직선으로 이어서, 그 직선을 따라 움직이는 것이 최적임을 밝혔는데, 이러한 Optimal Transport 는 식으로 다음과 같다.
\[p_{t} = [(1-t)\text{id} + t\psi]_{*}p_{0}\]여기서 $\text{id}(x) = x$ 는 그냥 항등 매핑 함수이다.
..아무튼 다시 정리하자면, 평균과 표준편차를 직선 형태로 정의했더니, 요 Optimal Transport, OT 와 정확히 일치했다는 것이다.
Compare with Diffusion processes
둘 다 동일한 $p_{0} \to p_{1}$ 으로 변화는 각 데이터 포인트들의 궤적을 나타내고 있는데, 같은 조건부 분포더라도, 그리는 궤적은 완전히 다르다.
Diffusion 은 약간 비유하자면 밥 그릇 던져주면 시골 똥개마냥 아주 이리저리 난리를 치면서 오는데, FM의 OT 는 약간 차분한 순종 귀족 강아지가 고개 딱 세우고 눈감은채 우아하게 걸어오는 느낌?
..장난 좀 쳐봤다. 그냥 저 사진으로만 다 이해돼서..
근데 물론 이 논문의 OT 는 학습할 때 임의의 $x_{1}$ 을 조건부로 둔 conditional VF 에 대해서 최적인거지, 모델 학습 후엔, $x_{1}$ 이 주어지지 않은 상태에서 학습된 VF 를 사용하여 나아가야 하는데, 모든 데이터에 대한 VF, 즉, marginal VF 에 대해서는 최적이 아님을 인지하자.
논문 형님들도 이를 언급했는데, 논문 형님들은 그럼에도 기존 Diffusion processes 보단 나을 거라고 기대한다. 실제로 그렇기도 하다.
찌라시시
일단 초안은 작성했는데, 약간.. 기하학적으로 설명하는데 글로만 쓰려다 보니 좀 거시기한데.. 그 부분은 한번 공부를 해봐야겠다..
아직 수정할 것도 많고 VF가 $t$ 에 의존하는 나를 납득시킬 이유같이 안다룬 내용도 있는데, 나름 핵심 정보들은 다 다뤘으니 차근차근 하도록 하겠다.
그리고 개인적으로 이 논문을 읽으면서 뭔가 굉장히 재밌었는데, 약간 Score-based 를 미리 공부해둔 상태에서 읽으니 논문 형님들이 말하는게 뭔지 쫙쫙 들어와서 약간 쾌감..? 을 느꼈다. 수학적으로도 나름…. 내딴에는 막힌 부분 없이 쫙 이해돼서 그런 것도 있다. 아니면 내 착각일 수도..
