Flow matching 논문 리뷰 및 설명

시작

• unknown target data distribution ,
하지만 우리는 해당 분포에서 나온 데이터포인트들 을 알고 있다.

• probability density path
는 time-dependent probability density function.
t가 0, 1일때 각각

• : vector field
• 의 의미 : 데이터 포인트 x가 시간 t일 때 이동해야하는 벡터
• A vector field can be used to construct a time-dependent diffeomorphic map, called a flow, , defined via the ordinary differential equation

• : flow
• 의 의미 : 데이터 포인트 x가 vector field를 따라 시간 t만큼 이동하였을 때 x의 위치
즉, 항상 (움직이지 않았기 때문에)

• flow는 vector field의 ODE(미분방정식)로 정의된다. contructs flow

• flow는 push-forward equation을 통해서 분포 자체를 변화시킬 수 있다.
여기서 push-forward equation은 아래 식으로 풀어서 적을 수 있다.



여기서 continuity equation을 사용해서

if v_t and p_t가 continuity equation을 만족시킨다면, v_t generates p_t
probability는 흐름에 따라 continuous하게 이동할 뿐, 갑자기 생기거나 사라지지 않는다.
 
결론적으로 : generates
→ 를 이용하면 를 closed form으로 정의할 수 있다 ❌
→ 를 알면 ODE를 풀어서 를 구할 수 있고, 이 flow를 이용해서 를 로 변화시킬 수 있다. ✅

1. 에서 로 이동가능한 의 선택지는 한개가 아니고, 무수히 많다.

2. 를 모르기 때문에 에 들어갈 x값을 샘플링할 수 없다.

3. 정답 vector field 를 모른다.

per-sample approach

• target probability path는 여러개의 간단한(한 샘플 단위의) probability path들의 조합으로 만들어낼 수 있다.

• 쉽게 생각해서 만약 타겟 분포 자체가 라고 할 때, 우리가 가진 특정 데이터 한개가 라고 하자.
• 그럼 를 mean으로 갖고 아주 작은 std값을 가진 Normal distribution을 데이터 수만큼 만들고 그걸 다 합치면 분포 이 만들어진다고 볼 수 있다.
• 즉

• 그럼 그럼 동일하게 전체 데이터 분포가 아니라 한개의 데이터 만 고려했을 때의 probability path도 정의할 수 있다.
Contitional probability path

• 그 다음 각 을 컨디션으로 conditional probability path들을 구한다음, Marginalizing the conditional probability paths over q(x1) give rise to the marginal probability path. (marginalizing over q(x1)의 뜻은 분포에 따른 확률을 고려해서 x1에 대해서 적분하겠다는 뜻)

• 그리고 conditional probability path에 대해 marginalize한 marginal vector field도 정의할 수 있다.
 

• Since pushes to , generates the , 그리고

• 이전에 정의한 CFM Loss

• 여기서 원래 인걸 x_0에 대한 식으로 reparameterization한다 by
그리고 위의 flow 식을 사용하면, Loss를 다르게 적을 수 있다.
이다.

• probability path가 Gaussian path이고, 그 변화를 만들어내는 flow가 위의 식에 따른 affine transformation(11)일 때 (둘다 위에서 가정)
그 flow를 정의하는 unique vector field를 closed form으로 정의할 수 있다.
증명



즉, given x1에 대한 만 모델링하면 된다.
어떻게 모델링하던지 간에 이 두 함수를 얻을 수만 있다면 목적을 달성한다. → 기존의 Diffusion process도 이걸 모델링하는 방법 중 하나.

Optimal Transport

다시 처음부터 쉽게 생각해보자

• 원하는건 데이터 생성 → 내가 가진 데이터셋에 대해 likelihood를 최대화하는 분포를 알아내는 것.

• 하지만 직접 target 분포를 모델링하는건 거의 불가능에 가깝다(intractable하다). 따라서 임의의 쉬운 분포를, 원하는 복잡한 분포로 매핑하는 방법을 찾으려고 한다.

• 방법은 여러 개가 있다. GAN, VAE는 모델이 한번에 p_0을 p_1로 만드는 방법을 학습하도록 하고, normalizing flow는 일정 스텝에 걸쳐서 discrete하게(= not ODE) 분포를 변환시키는 역변환 가능한 함수를 알아내려고 한다.

• flow matching은 분포가 만약 t~[0, 1]에서 source에서 타겟으로 점점 변한다고 치고, probability density path p_t와 time-dependent vector field v_t를 생각해본다. 그리고 이 vector field construct flow, defined via a ODE
• , t=0일때 flow는 identity
• 그리고 이 flow는 source 분포(p_0)를 target 분포(p_t)로 push-forward equation을 통해 변환시킨다.
• 그럼 우리는 vector field가 probability density path를 generate한다고 볼 수 있다. (만약 참이라면 시간의 흐름에 따라 probability mass를 보존하고, continuity equation을 만족해야한다.)

• 그래서 뉴럴넷을 통해서 vector field를 모델링하고 regress한다. 즉, Loss는
하지만 2) p_t를 모르기 때문에 x_t를 샘플링할 수도 없다. 2) GT vector field를 알지도 못해서 Loss 식을 계산할 수 없다.

• 그럼 어떻게 하지? 전체 모든 데이터 범위에 대한 Ground Truth Vector field를 알 수는 없다. 하지만 Mixture of Gaussians의 시각으로 분포라는걸 바라보자. N개의 데이터셋으로 구성되는 전체 분포는, 데이터셋의 각 데이터 포인트 값을 mean으로, 작은 표준편차를 가진 Gaussian 분포 N개를 합치고 나눈 것과 같다고 볼 수 있다. 이건 복잡한 분포를 근사하는 유용한 방법 중 하나이다.



그래서 임의로 하나의 데이터 포인트만 정한다음, 이걸 해당 값을 평균으로 갖고 아주 작은 std를 가진 Gaussian 분포로 바라보고 source에서 target으로 이동하는 probability density path를 정의한다 - conditional ~~
• 그런 다음 이렇게 얻은 각각의 simpler probability path를 전체 데이터셋 q(x1)에 대해 marginalizing하면 원하는 probability path를 얻었다고 볼 수 있다!(A simple way to construct a target probability path is via a mixture of simpler probability paths)
• 따라서 하나의 데이터 포인트 x1을 condition으로 vector field를 예측하는 것으로 objective를 변경한다.
• 모든 데이터가 아니라 하나의 데이터 포인트에 대한 conditional vector field를 예측하는게 훨씬 간단하다 = x_t와 v_t 를 임의로 정하고 계산하기 훨씬 수월하다.

• 따라서 Conditional Flow Matching loss
이걸 최적화하기위한 gradient는 기존의 FM loss의 gradient와 같다.

• 일단 좀 더 간단하게 만들긴 했는데, 여기서도 pt와 vt에 대한 임의의 정의가 필요하다. 이렇게만 계산하면 중간의 p_t는 무슨 모양의 분포던지 될 수 있다고 가정한 상태임

• p0과 p1 둘다 Gaussian distribution이다. pt가 정의되는 방법은 너무 많은데, 가장 간단한 방법을 고려해보자 : mean과 std만 시간에 따라 변화하고, 계속 Gaussian 분포의 형태로 시간에 따라 이동하고, 제자리 변환 같은 무의미한 변화는 없는.

• 그럼 이 분포 이동을 만들어내는 flow는 affine transformation이다.
이걸 이용해서 다시 위의 CFM Loss식에 넣는다. 시간 t에 대해 미분한게 vector field니까,
이런 관점에서는 mu와 std 함수만 어떻게든 정의한다면, 그 방법에 상관없이 목적을 달성한다고 볼 수 있다.

• Diffusion도 같은 조건 하에 를 정의한 방법들 중 하나라고 본다. 따라서 현재의 개념에 속함.
대신 그 과정이 stochastic한. (Variance Exploding score matching, Variance Preserving Diffusion 둘다)

• t=0에서 mean = 0, std = 1 | t=1에서 mean = x1, std = sigma_min
만약 정의하는 방법에 제한이 없다면, 시간 t에 따라 linear하게 변화하는 flow를 생각해보자(가장 쉽고 간단한 이동)
그럼 given x1에 대한 도 interpolation만으로 쉽게 계산할 수 있고, vector field도 쉽게 정의할 수 있어서 Loss를 계산할 수 있다.
 
이렇게하면 모델링이 훨씬 간단해질 뿐더러, Gaussian 분포간의 이동에대한 최적의 변환 Optimal Transport를 달성하니까 좋다!