조건에 맞춘 소스 설계가 흐름 매칭을 가속한다

초록

본 논문은 텍스트‑이미지 생성에서 흐름 매칭(Flow Matching) 모델이 고정된 표준 가우시안 대신, 조건(텍스트 프롬프트)에 의존하는 학습 가능한 소스 분포를 사용하도록 설계하였다. 소스 분포의 평균은 자유롭게 이동시키고, 분산은 단위 가우시안에만 정규화하는 방식으로 붕괴를 방지한다. 또한 소스와 타깃 사이의 방향 정렬 손실을 도입해 학습 안정성을 높였다. 실험 결과, FID와 CLIP Score에서 기존 방식 대비 2~3배 빠른 수렴과 전반적인 품질 향상을 입증한다.

상세 분석

본 연구는 흐름 매칭이 “소스 → 타깃”을 연속적인 ODE로 모델링한다는 점에 착안해, 소스 분포 자체를 최적화 변수로 만든다. 기존 대부분의 흐름 매칭은 확산 모델에서 물려받은 표준 정규분포 N(0,I)를 그대로 사용했으며, 이는 조건부 생성에서 프롬프트 정보를 전혀 반영하지 못한다는 한계가 있었다. 저자들은 이를 극복하기 위해 조건 C(텍스트 임베딩) 에 의존하는 가우시안 pϕ(X₀|C)=𝒩(μϕ(C),σ²ϕ(C)I) 를 도입한다. 여기서 핵심은 두 가지 정규화 전략이다. 첫째, 분산 σ²ϕ(C) 가 훈련 과정에서 0으로 수렴하는 “분산 붕괴” 현상을 방지하기 위해, 평균을 고정하지 않은 채 σ²ϕ(C) 가 단위 분산에 가깝도록 KL 정규화 L_VarReg을 적용한다. 이는 KL(N(μ,σ²I)‖N(μ,I)) 형태로, 평균은 자유롭게 이동시켜 소스가 타깃 모드에 가까워질 수 있게 한다. 둘째, 소스와 타깃 샘플 사이의 방향을 맞추는 L_align = E