Flow Matching : Diffusion과의 차이점을 중심으로

Flow matching (https://zzziito.tistory.com/153)은 디퓨젼과 함께 최근 확률적 러닝 모델 기반 생성 분야에서 주목받고 있는 방법이다. Flow matching은 확률 흐름 상미분 방정식 (Ordinary Differential Equation, ODE)을 기반으로 하며, 이 방법은 복잡한 데이터 분포를 단순한 분포(Ex. 가우시안)로 연속적으로 변형하는 벡터 필드를 학습한다. 데이터 분포와 단순한 분포 사이의 확률 흐름을 모델링하여 이 흐름을 따라 이동하는 궤적을 학습하고, 연속 정규화 흐름(Continuous Normalizing Flow, CNF)의 개념을 확장하여 적용했다. 

디퓨젼 모델이 노이즈가 없는 이미지에 점진적으로 가우시안 노이즈를 추가하여 데이터를 손상시키고, 이 노이징 과정을 학습하여 디노이징 과정을 예측한다면, Flow matching은 노이즈를 추가하는 대신 데이터 분포를 직접 변형시키고, 원본 데이터 분포에서 가우시안 분포로 (혹은 그 반대로) 변환되는 과정에서 속도 벡터 필드를 학습한다. 이 벡터 필드는 두 분포 사이의 최적 경로를 따라 이동하는 방법을 정의하며, 생성 시에는 단순 분포에서 샘플링한 후 학습된 벡터 필드를 따라 ODE 솔버로 적분하여 원본 데이터 분포의 샘플을 생성한다. 

FLOW MATCHING FOR GENERATIVE MODELING

여기서, "노이즈를 추가하는 데이터 분포를 직접 변형시킨다"는 것의 의미는 다음과 같다.

디퓨젼 모델의 경우 원본 이미지에 점진적으로 무작위 가우시안 노이즈를 추가하지만 Flow matching은 데이터 포인트의 위치를 직접 이동시키는 변환(transformation)을 정의하고, 노이즈를 추가하는 것이 아니라 확률 공간에서 이동 경로(path)를 설계한다. 따라서 디퓨젼의 경우 이미지에 무작위 노이즈를 더해서 점점 알아볼 수 없게 만들지만 Flow matching은 이미지의 각 픽셀이 어떤 경로를 따라 이동해야 하는지 정의하고, 그 경로를 따라 이미지를 변형한다. 

디퓨젼은 이미지에 안개를 점점 더 짙게 끼게 하는 방식이라면 Flow matching은 이미지의 각 픽셀을 특정 방향으로 조금씩 밀어내는 방식이라고 생각할 수 있다. 

 

시뮬레이션 과정에서의 차이

디퓨젼은 1) 순방향 과정(Forward process), 2) 역방향 과정(Inverse process)이라는 두 단계로 작동한다. 먼저 순방향 과정에서는 실제 데이터(ex. 깨끗한 이미지)에 점진적으로 노이즈를 추가한다. 이 과정은 미리 정해진 노이즈 스케줄에 따라 단계별로 진행되며, 각 단계는 이전 단계에만 의존하는 마르코프 특성을 가진다. 훈련 중에는 이 과정을 모델링하여 노이즈 예측기를 학습한다. 이후 역방향 과정에서는 완전한 노이즈 상태에서 시작해 학습된 노이즈 예측기를 사용하여 단계별로 노이즈를 제거한다. 이 과정은 보통 수십, 수백 번의 반복 단계가 필요하기 때문에 계산 비용이 높은 편이다. 

Flow matching은 먼저 1) 경로 정의 단계에서는 실제 데이터 분포와 단순한 노이즈 분포 (ex.가우시안) 사이의 연속적인 경로를 명시적으로 설계한다. 다음으로 2) 벡터 필드 학습 단계에서는 이 경로를 따라 움직이는 속도 벡터 필드를 직접 학습한다. 학습 과정에서는 경로 상의 임의의 시간 지점에서 속도 벡터를 예측하도록 모델을 훈련시킨다. 마지막으로 3) 생성 과정에서는 간단한 노이즈 분포에서 샘플을 추출한 후, 학습된 속도 벡터 필드를 따라 ODE 솔버를 사용해 적분함으로써 샘플을 생성한다. 

디퓨젼은 노이즈 추가와 제거의 이산적인 단계를 거치는 반면 Flow matching은 연속적인 변환 경로를 따라 이동하는 벡터 필드를 모델링한다. 

 

Diffusion vs. Flow matching : 언제 어떤 모델을 선택할까 ?

 

디퓨젼 모델은 복잡한 조건부 생성이 필요할 때나 높은 샘플 퀄리티가 최우선일 때 적합한 것 같고, Flow matching은 빠른 샘플링 속도가 중요할 때 적합한 것 같다. 또한 Flow matching은 변환 과정을 좀 더 직접적으로 제어할 수 있기 때문에 경로를 좀 더 세밀하게 제어해야 하거나 특정 도메인 지식을 변환 경로에 통합해야 할 때 적합하다.