Flow Matching For Generative Modeling

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}$ 로 정의되며, 이 논문에선 $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$ 에서 데이터 포인트들에게 작용하는 힘이 정의되어 있다. 다시말해 $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}$ 가 되겠다… 아무튼 반지의 제왕식 비유법이다. 그냥.. 각 개념이 뭔 느낌인지 참고만 하자… 약간 Score based SDE 에서 동일한 궤적을 그리는 ODE 를 다시 생각해보면, 여기서 말하는 Flow는 VF를 기준으로, 그리고 Score based ODE는 Score를 기준으로 궤적을 정의하므로, 어.. 비슷한 느낌이라고 볼 수 있다.

아무튼 위와 같이 $\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 는 다음과 같이 이 VF를 학습하게끔 정의된다.

\[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})\]

Score based model 에서 학습 대상이 Score 벡터장이 아닌 VF로 바뀌었을 뿐이다.

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})p_{t}(x_{1}|x) \, dx_{1} \\ &= \int u_{t}(x|x_{1})\frac{p_{t}(x|x_{1})q(x_{1})}{p_{t}(x)} \, dx_{1} \\ &= \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} \end{align}\]

여기서 왜 저기 위에 확률 밀도와 마찬가지로 $q(x_{1})$ 에 대해서 평균을 안내는지는 앞서 봤던 Continuity Equation 로 유도할 수 있다.

\[\begin{align} &\frac{\partial}{\partial t}p_{t}(x|x_{1}) + \nabla \cdot (p_{t}(x|x_{1})v_{t}(x|x_{1})) = 0 \\ &\frac{\partial}{\partial t}p_{t}(x|x_{1})q(x_{1}) + \nabla \cdot (p_{t}(x|x_{1})v_{t}(x|x_{1}))q(x_{1}) = 0 \\ &\int \frac{\partial}{\partial t}p_{t}(x|x_{1})q(x_{1}) \, dx_{1} + \int \nabla \cdot (p_{t}(x|x_{1})v_{t}(x|x_{1}))q(x_{1}) \, dx_{1} = 0 \\ &\frac{\partial}{\partial t} \underbrace{\int p_{t}(x|x_{1})q(x_{1}) \, dx_{1}}_{p_{t}(x)} + \nabla \cdot \int p_{t}(x|x_{1})v_{t}(x|x_{1})q(x_{1}) \, dx_{1} = 0 \end{align}\]

마지막 등식의 Continuity Equation 를 만족하기 위해서는 다음을 만족해야 한다.

\[p_{t}(x)v_{t}(x) = \int p_{t}(x|x_{1})v_{t}(x|x_{1})q(x_{1}) \, dx_{1}\]

즉, 여기다가 양변에 $p_{t}(x)$ 로 나눠주면 다음과 같다.

\[\begin{align} v_{t}(x) &= \int v_{t}(x|x_{1})\frac{p_{t}(x|x_{1})q(x_{1})}{p_{t}(x)} \, dx_{1} \\ &= \int v_{t}(x|x_{1})p_{t}(x_{1}|x) \, dx_{1} \end{align}\]

아무튼 이런 이유로 $q(x_{1})$ 로 평균을 내지 않는 것이다. 음.. 직관적으로 설명해보면… 우리가 조건부로 둔 $x_{1}$ 에 대해서, 모든 $x_{1}$ 이 $x$ 를 지나는게 아니기 때문이다. 어… 그러니까.. 씁.. 아무튼 $v_{t}(x)$ 는 $t$ 시점에서 $x$ 를 지나는 입자들의 속도인데? $p(x)$ 확률 밀도를 형성하는 입자들이 이 VF를 따라 이동하여 최종적으로 $p_{1}(x)$ 를 형성한다. 근데 $x_{1}$ 은 이 $p_{1}(x)$ 전체 확률 밀도를 의미하는게 아니라, 거기서 샘플링된 단일 입자로 볼 수 있다. 그러면 즉, $p(x)$ 의 $x$ 입자들이 이동하여 $p_{1}(x)$ 확률 밀도 전체를 형성하지만? 모든 $x$ 가 $p_{1}$ 의 일부분인 $x_{1}$ 로 이동하지는 않는다는 것이다. $x_{1}$ 로 이동하지 않는 $x$ 입자들도 있다는 것이다. 아 말로 설명하려니 좀 거시기 하네 아무튼 $t$ 시점에서 최종 도착지가 $x_{1}$ 일 때 $x$ 에서의 속도인 $v_{t}(x|x_{1})$ 를 평균 내기 위해서는 이 VF 의 조건부인 $x_{1}$, 요것이? 모든 $x_{1}$ 에 대해서 평균 내는게 아니라, $x$ 중 $x_{1}$ 에 도착하는 고것들만 가지고 평균을 내야 한다는 것이다. 때문에 $q(x_{1})$ 이 아닌 해당 $x$ 에 위치해 있을 때, $x_{1}$ 에 도착할 확률인 $p_{t}(x_{1}|x)$ 로 평균을 내는 것이다. 그냥 비문학이다 들여쓰기 하나도 없이 줄줄이 써서 죄송하다. 아무튼 내가 이해한바로는 이렇다. 계속 공부해보다 틀린거 같으면 수정하도록 하겠다.

아무튼 결론은? 요 VF도 이런식으로 계산할 수는 있는데, 마찬가지로 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 \left( \int \left\langle v_{t}(x;\theta), u_{t}(x|x_{1})p_{t}(x|x_{1})q(x_{1})\right\rangle \,dx_{1}\right) \,dx \quad (\text{$v_{t}$는 $x_{1}$ 에 의존하지 않아 적분 기호를 밖으로 뺌})\\ &= \int \left( \int \left\langle v_{t}(x;\theta), u_{t}(x|x_{1})\right\rangle p_{t}(x|x_{1})q(x_{1}) \,dx_{1}\right) \,dx \quad (\text{$p_{t}, q$ 는 스칼라이므로, 내적 밖으로 뺌}) \\ &= \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)$ 를 우리의 조건부 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 를 최적화하는 것은 같은 결과를 가진다는 것이다.

이게 뭔소리냐. 간단한 예시로..음.. 이차항식을 생각해보자.

\[ax^{2} + bx + C \quad ( a > 0)\]

요거를 이제 최적화해야 하는데, 그러려면? 당연히 미분을 해서 해당 식이 $0$ 이 되게끔 하는 $x$ 찾고? 이게 최적의 파라미터 $x$ 일 것이다. 근데? 여기서 결국 $x$ 에 의존하지 않는 상수 $C$ 는? 이 최적 파라미터 $x$ 를 구하는데 아무런 도움도 되지 않는다. 아무튼 이런거랑 비슷한 느낌이다.

그래서 우리는 이 conditional VF 를 모델이 학습하게끔 하여, 결과적으로 $p_{t}(x)$ 를 학습할 수 있다!!!!!!

Conditional Probability Path and VF

그러면 이제 conditional probability path $p_{t}(x|x_{1})$ 와 conditional vector field $u_{t}(x|x_{1})$ 를 어떻게 정의하는지만 알면 된다.

논문에선 일단 이 조건부 확률 밀도를 조건부 가우시안의 형태로 나타낸다.

\[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_{1})}(\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_{1})}(\psi_{t}(x) - \mu_{t}(x_{1})) + \mu_{t}'(x_{1})\right)\right\|^{2}\right)\]

그럼 이제 진짜 마지막으로 남은건 $\mu_{t}(.)$ 와 $\sigma_{t}(.)$ 를 어떻게 정의하는지 뿐이다.

Diffusion conditional VFs

이 논문 형님들은 우리가 본 Diffusion Processes 도 이러한 FM 의 특이 케이스일 뿐이라고 한다. 왜 그런지 알아보자.

Score-based Generative Models 블로그 참고 글

..일단 전에 Score-based model 논문에서, Diffusion models 을 SDE 로 정의를 했었다. 막 각 데이터 포인트들이 $p_{1}$ 로 이동하되, 그 궤적이 노이즈 항이 섞여 약간 지그재그 형식으로 뺑 돌아가거나 하는 느낌이었는데, 이러한 복잡한 경로라 하더라도, 각 시점 $t$ 에서 데이터 포인트들로 이루어진 확률 분포의 모양 $p_{t}$ 를 유지하면서, 각 데이터 포인트들이 노이즈 항으로 인해 지그재그 발광을 하는게 1:1 로 대응되는 결정론적인 하나의 ODE로 나타낼 수 있다고 했었다.

그러니까, Diffusion SDE는 다음과 같은 형태였었다.

\[dx = f_{t}dt + g_{t}dw\]

각 항이 의미하는 바는 저 위의 블로그를 다시 보자..

아무튼.. 저 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$ 는 $x$ 에 대한 Laplace 연산이다. 이제 저 식을 전개해서 VF 가 뭔지 한번 알아보자..

\[\begin{align} \frac{dp_{t}}{dt} &= -\text{div}(f_{t}p_{t}) + \frac{g_{t}^{2}}{2}\text{div}(\nabla p_{t}) \\ &= -\text{div}(f_{t}p_{t}) + \text{div}(\frac{g_{t}^{2}}{2} \nabla p_{t}) \\ &= -\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}) \\ \frac{dp_{t}}{dt} + \text{div}(u_{t}p_{t}) &= 0 \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 \sim p_{t}(x|x_{1}) = \mathcal{N}(x|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}(x|x_{1}) = \frac{\sigma_{t}'}{\sigma_{t}}(x-x_{1})\]

앞서 살펴본 것처럼, FM 을 통해 나타낸 VF가 만들어내는 데이터 포인트들의 궤적은, 저 앞에서 구한 Diffusion VF 에 각각 $f_{t}, g_{t}$ 와 조건부 가우시안 분포 값을 대입해서 보면, 정확히 같은 결과가 나온다. 한번 확인해보면…

\[\begin{align} u_{t} &= f_{t}-\frac{g_{t}^{2}}{2}\nabla \log p_{t} \quad (f_{t} = 0, g_{t} = \sqrt{2\sigma_{t}\sigma_{t}'}) \\ &= - \sigma_{t}\sigma_{t}' \cdot \nabla \left(- \frac{\|x - x_{1}\|^{2}}{2\sigma_{t}^{2}} + C\right) \\ &= - \sigma_{t}\sigma_{t}' \cdot \left(- \frac{x-x_{1}}{\sigma_{t}}\right) \\ &= \frac{\sigma_{t}'}{\sigma_{t}}(x-x_{1}) \end{align}\]

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))x}{1 - \exp(-T(t))}\right) \end{align}\]

역시나 Diffusion VF 에 $f_{t}, g_{t}$ 와 조건부 가우시안 분포 값을 대입해서 보면, 같은 결과가 나온다. 한번 확인해보면…

\[\begin{align} u_{t} &= f_{t}-\frac{g_{t}^{2}}{2}\nabla \log p_{t} \quad (f_{t} = - \frac{T'(t)}{2}x, g_{t} = \sqrt{T'(t)}) \\ &= - \frac{T'(t)}{2}x - \frac{T'(t)}{2} \cdot \nabla \left(-\frac{\|x - \alpha_{t}x_{1}\|^{2}}{2(1 - \alpha_{t}^{2})} + C\right) \\ &= - \frac{T'(t)}{2}x - \frac{T'(t)}{2} \cdot \left(- \frac{x - \alpha_{t}x_{1}}{1 - \alpha_{t}^{2}}\right) = - \frac{T'(t)}{2} \left(x - \frac{x - \alpha_{t}x_{1}}{1 - \alpha_{t}^{2}}\right) \\ &= - \frac{T'(t)}{2} \left(\frac{x(1 - \alpha_{t}^{2}) - x + \alpha_{t}x_{1}}{1 - \alpha_{t}^{2}}\right) = - \frac{T'(t)}{2} \left(\frac{\alpha_{t}x_{1} - \alpha_{t}^{2}x}{1 - \alpha_{t}^{2}}\right) \\ &=- \frac{T'(t)\alpha_{t}}{2} \left(\frac{x_{1} - \alpha_{t}x}{1 - \alpha_{t}^{2}}\right) \\ &= \frac{\alpha_{t}'}{1-\alpha_{t}^{2}}(x_{1} - \alpha_{t}x) \quad \left(-\frac{T'(t)}{2}\alpha_{t} = \alpha_{t}'\right) \end{align}\]

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 는 약간 차분한 순종 귀족 강아지가 고개 딱 세우고 눈감은채 우아하게 걸어오는 느낌?

..장난 좀 쳐봤다. 그냥 저 사진으로만 다 이해돼서..

아무튼 Score based SDE 즉, Diffusion 은 그 뭐냐 그거, 열역학? 열 확산? 아무튼 그거를 기반으로 하여 브라운 운동과 같은 확률 과정이 데이터 포인트들의 이동 궤적에 포함되어, 저렇게 삥 돌아가는 형태가 되는데, 이 논문의 VM은 뭐 앞서 본 포..폭커 플랑크..? 방정식 그거나 확률 과정이 포함된 SDE가 아닌, 두 확률 분포 사이의 VF 자체를 디자인함으로써, 좀 더 보다 덜 제약적이고, 효율적인 이동 궤적을 설계할 수 있게 된 것이다.

그러니까.. Score based SDE 가 DDPM 과 SMLD 를 하나로 일반화해서 나타냈다면, 이 논문의 FM 이 이를 더 일반화해서 나타낸 것이다. (정확하게는 Score based SDE의 확률 과정을 제거하여 더 효율적인 궤적으로 최적화? 한거라고 생각한다. 아무튼 난 이렇게 이해했다. )

어… 약관 관계를 다시 정리해보면..?

• Flow Matching (Continuous Normalizing Flow)
  • Gaussian Conditional Flow Matching
    • Diffusion Probability Flow ODE
      • VP-SDE
      • VE-SDE
      • sub-VP
    • Optimal Transport FM

약간 이런 느낌..?

아무튼 데이터 포인트들의 이동궤적이 가장 빠른 직선 형태로 이루어지기 때문에, 기존 SDE 를 통한 뻉 돌아가는 곡선 형태의 이동 궤적일 때보다 확실히 빠르다는 것이다. 즉, 샘플링 속도가 이빠이 빠르다. DDPM 논문에서 $T=1000$ 으로 두어 $1000$ 번의 단계를 거쳐 샘플링했던 것을 다시 상기해보면, 직접 구현해봐야 알겠지만 확실히 더 빠를 것이다.

추가로 우리가 $p_{0}$ 를 굳이 노이즈로 국한할 필요가 없다는 것이다. 기존 score based SDE 는 시간이 흐르면서 결국엔 기존 초기 데이터 $x_{0}$ 의 모든 정보가 사라지며 완전한 노이즈 분포로 수렴하도록 설계되어 있었는데, 여기서의 FM은 노이즈-분포 가 아닌 분포-분포 의 관계라는 것이다. 마찬가지로 $x_{1}$ 에서만 샘플링하여 조건부로 두는게 아니라, $x_{0}$ 에서도 샘플링하여 추가적인 조건부로 둠으로써 수행하면 된다는 것이다. 요거는 따로 구현해서 알아보도록 하겠다.

그리고 Score based SDE 논문에서 Probability ODE 그거, score based SDE 와 동일한 확률 밀도 궤적을 그리면서 확률과정 없이 결정론적으로 나타나지는 그거와 비교해볼 수 있는데, 간단하게도 이는 이미 살펴본 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}\]

결국 Score based SDE 는 이 FM 의 손아귀 안에서 놀아나고 있다는 것이다.

근데 물론 이 논문의 OT 는 학습할 때 임의의 $x_{1}$ 을 조건부로 둔 conditional VF 에 대해서 최적인거지, 모델 학습 후엔, $x_{1}$ 이 주어지지 않은 상태에서 학습된 VF 를 사용하여 나아가야 하는데, 모든 데이터에 대한 VF, 즉, marginal VF 에 대해서는 최적이 아님을 인지하자.

논문 형님들도 이를 언급했는데, 논문 형님들은 그럼에도 기존 Diffusion processes 보단 나을 거라고 기대한다.

찌라시시

일단 초안은 작성했는데, 약간.. 기하학적으로 설명하는데 글로만 쓰려다 보니 좀 거시기한데.. 그 부분은 한번 공부를 해봐야겠다..

아직 수정할 것도 많고 VF가 $t$ 에 의존하는 나를 납득시킬 이유같이 안다룬 내용도 있는데, 나름 핵심 정보들은 다 다뤘으니 차근차근 하도록 하겠다.

그리고 개인적으로 이 논문을 읽으면서 뭔가 굉장히 재밌었는데, 약간 Score-based 를 미리 공부해둔 상태에서 읽으니 논문 형님들이 말하는게 뭔지 쫙쫙 들어와서 약간 쾌감..? 을 느꼈다. 수학적으로도 나름…. 내딴에는 막힌 부분 없이 쫙 이해돼서 그런 것도 있다. 아니면 내 착각일 수도..

(2026.08.05)

지금까지 정리한 블로그 글들 간단히 복습하면서 오인하고 있던 개념들 다시 바로 잡고 보니까, DDPM -> score based SDE -> FM 응 내가 짱이야 응 내가 짱이야 이 흐름이 자연스럽게 연결되면서 읽히니까 기부니가 좋다. 물론 이것도 현재의 나 기준이지 미래의 다시 공부하고 돌아온 내가 봤을 때 잘못된 방식으로 이해한 것 처럼 보일 수 있다. 이게 뭔소리지 아무튼 기부니가 좋다.