MixFlow: 혼합 소스 분포를 통한 Rectified Flow 성능 개선

Diffusion 모델의 느린 샘플링 속도 원인인 생성 경로의 높은 곡률 문제를 해결하기 위해 소스 분포를 데이터와 정렬하는 새로운 학습 전략을 제시한다. 기존 Gaussian 분포에만 의존하던 방식에서 벗어나 조건부 분포를 혼합함으로써 적은 샘플링 단계에서도 고품질 이미지를 생성하고 학습 수렴 속도를 크게 높였다.

Diffusion 모델이나 Rectified Flow는 무작위 노이즈(Gaussian 분포)를 복잡한 이미지 데이터로 변환하는 과정을 학습한다. 이때 노이즈와 실제 데이터 사이의 관계가 독립적이면 변환 경로가 심하게 휘어지게 되는데, 이는 미분 방정식을 풀 때 많은 계산 단계를 요구하게 되어 샘플링 속도가 느려지는 원인이 된다.

MixFlow는 이 문제를 해결하기 위해 노이즈(소스) 자체를 데이터와 조금 더 닮게 만드는 '조건부 분포'를 활용한다. 하지만 조건부 분포에만 너무 의존하면 학습 데이터에 없는 새로운 이미지를 생성하는 능력이 떨어지는 'Prior Hole' 문제가 발생할 수 있다. 이를 방지하기 위해 MixFlow는 순수 노이즈와 조건부 노이즈를 섞어서 학습하는 방식을 택했다.

이러한 혼합 방식은 모델이 노이즈에서 데이터로 가는 길을 더 직선에 가깝게 인식하도록 유도한다. 결과적으로 구불구불한 길 대신 곧게 뻗은 길을 따라 데이터를 생성하게 되므로, 훨씬 적은 횟수의 계산만으로도 선명하고 다양한 이미지를 얻을 수 있게 된다.

전체적인 접근 방식은 소스 분포 p0와 타겟 분포 p1 사이의 결합(Coupling)을 최적화하여 생성 경로를 직선화하는 것이다. 이를 위해 임의의 신호 κ(클래스 레이블, 데이터 샘플 등)를 입력으로 받아 소스 분포의 평균과 분산을 예측하는 네트워크 qφ(x0|κ)를 도입하고, 이를 벡터 필드 vθ와 함께 공동 최적화한다.

핵심 메커니즘인 MixFlow 학습은 두 소스 분포의 선형 보간을 사용한다. [0, 1] 사이의 스칼라 가중치 w를 샘플링하여, 학습 가능한 가우시안 N(μφ(κ), Σφ(κ))와 표준 가우시안 N(0, I)을 혼합한 qφ(x|κ, w) = N(wμφ(κ), wΣφ(κ) + (1-w)I)를 소스로 사용한다. 이 과정에서 KL Divergence 손실 함수를 추가하여 소스 분포가 표준 가우시안에서 너무 멀어지지 않도록 규제한다.

학습 시에는 매 단계마다 w를 무작위로 선택하여 모델이 다양한 혼합 정도에 노출되게 한다. 이는 조건부 소스 분포와 타겟 분포 사이의 효율적인 결합 정보를 무조건부 소스 분포(표준 가우시안)로 전이시키는 역할을 하며, 추론 시에는 κ가 없어도 표준 가우시안에서 직선에 가까운 경로로 샘플링이 가능하게 한다.

CIFAR10 데이터셋 실험 결과, MixFlow는 RK45 솔버 사용 시 FID 2.27을 기록하여 표준 Rectified Flow(2.58) 및 Fast-ODE(2.45)보다 우수한 성능을 보였다. 특히 적은 샘플링 단계(NFE=5)에서 Heun의 2차 솔버를 사용했을 때 FID 19.29를 달성하여 기존 QAC(19.68)나 Fast-ODE(24.40)를 크게 앞질렀다.

FFHQ 및 AFHQv2 64x64 고해상도 데이터셋에서도 MixFlow는 모든 샘플링 단계에서 Fast-ODE보다 낮은 FID를 기록했다. 예를 들어 AFHQv2에서 10단계 샘플링 시 MixFlow는 FID 7.95를 기록한 반면, Fast-ODE는 설정에 따라 10.56~11.80 수준에 머물렀다.

학습 효율성 측면에서 MixFlow는 기존 강력한 베이스라인인 Fast-ODE가 도달한 성능을 단 60%의 학습 반복 횟수만으로 달성하며 수렴 속도의 우위를 입증했다. 또한 β(KL 규제 가중치) 값을 10^-5 수준으로 낮게 설정해도 학습 안정성이 유지됨을 확인했다.

MixFlow 아키텍처는 벡터 필드를 예측하는 메인 UNet과 소스 분포의 파라미터를 예측하는 소형 UNet(Source Prediction Network)으로 구성된다. 소스 예측 네트워크는 약 2M 파라미터 규모로 가벼우며, 입력 신호 κ의 특성에 따라 임베딩 레이어를 유연하게 변경할 수 있다.

수학적으로 MixFlow는 최적 운송(Optimal Transport) 이론에 기반한 경로 직선화 기법을 확장한다. 기존 Fast-ODE가 고정된 β 값에 민감하게 반응하여 성능 저하나 Prior Hole 문제를 겪는 것과 달리, MixFlow는 분포 혼합(Distribution Mixing)을 통해 매우 작은 β 값에서도 안정적인 학습이 가능하도록 설계되었다.

이론적으로 MixFlow의 손실 함수는 벡터 필드 학습을 위한 MSE 손실과 소스 분포 정규화를 위한 KL Divergence 손실의 합으로 정의된다. 이때 혼합 가중치 w는 모델이 조건부 정보의 구조를 무조건부 공간으로 매핑하도록 강제하는 브리지 역할을 수행하며, 이는 생성 경로의 교차(Intersection)를 최소화하는 효과를 낳는다.

κ-FC 수식은 일반적인 조건부 변수를 추상화하고 있지만, 본 논문에서는 주로 노이즈, 레이블, 데이터 샘플에 대해서만 실험을 진행했다. 텍스트 프롬프트와 같은 더 복잡한 조건부 입력에 대한 검증은 향후 과제로 남아 있으며, 여전히 가우시안 가정을 기반으로 하고 있어 더 자유로운 형태의 분포 완화 연구가 필요하다.