기술 요약: FlowConsist

문제 정의

Consistency Models 및 MeanFlow를 포함한 Fast Flow 모델들은 ODE(상미분 방정식) 경로 적분을 직접 예측하도록 학습함으로써, 1-스텝 또는 적은 수의 스텝만으로 생성을 가능하게 하여 생성 샘플링을 가속화하는 것을 목표로 한다. 그러나 저자들은 현재의 학습 패러다임이 성능을 저하시키는 두 가지 근본적인 문제점을 가지고 있다고 주장한다:

1. 체계적 궤적 드리프트 (Systematic Trajectory Drift): 현재의 방법들은 노이즈($\epsilon$)와 데이터($x$) 샘플을 무작위로 쌍을 지어 조건부 속도(conditional velocities)를 구성한다. 이러한 조건부 속도의 기댓값은 주변 속도(marginal velocity)와 일치하지만, 개별적인 조건부 경로는 서로 다른 여러 주변 궤적과 교차한다. 이러한 조건부 속도로 학습하면 최적화 대상이 진정한 단일 주변 ODE 경로로부터 벗어나게 되어, 결과적으로 일관되지 않은 궤적을 초래한다.
2. 오차 누적 (Error Accumulation): 모델에 의해 도입된 근사 오차는 긴 시간 간격에 걸쳐 누적된다. 스스로 예측한 속도에 의존하여 학습하는 Fast Flow 모델의 경우, 이러한 오차가 중첩되어 생성된 샘플과 실제 데이터 분포 사이의 상당한 편차를 발생시키며, 특히 단일 스텝(1-NFE) 생성 시에 두드ver진다.

방법론

저자들은 사전 학습된 교사(teacher) 모델 없이도 궤적 일관성을 강제하고 오차 누적을 교정하도록 설계된 학습 프레임워크인 FlowConsist를 제안한다.

1. 궤적 일관성 강제 (조건부 속도 대체)

핵심적인 이론적 통찰은 조건부 속도($v_t = \epsilon - x$)에 의존하는 것이 손실 함수에 분산 항($L_{var}$)을 도입하여, 모델이 고분산 방향의 그래디언트를 억제하게 만들고 곡선 궤적을 충실히 피팅하는 것을 방해한다는 점이다.

• 이론적 전환: 조건부 속도를 주변 속도의 대리물로 사용하는 대신, FlowConsist는 일관성 목적 함수 내의 조건부 속도를 모델 자체에 의해 예측된 주변 속도($u_\theta$)로 대체한다.
• 구현: 저자들은 표준 Flow Matching 항과 궤적 일관성 정규화 항을 결합한 새로운 학습 목적 함수(식 7)를 도출한다.
  • 이 목적 함수는 예측된 평균 속도와 모델 자신의 주변 속도 예측으로부터 유도된 타겟 사이의 차이를 최소화한다.
  • 결정적으로, 조건부 속도는 모델이 학습된 후 기대치 측면에서 주변 속도로 수렴하는 항에서만 사용되므로, 최적화 과정 중 발생하는 체계적인 드리프트를 효과적으로 제거한다.
• Classifier-Free Guidance (CFG): 이 방법은 가이던스 스케일 $\omega$를 보조 입력으로 취급하여 CFG를 통합하며, 이를 통해 추론 시 조정 가능한 통합된 표현을 학습할 수 있게 한다.

2. 궤적 교정 (오차 누적 해결)

시간 단계에 따른 누적 근사 오차로 인한 단조적 발산을 해결하기 위해, 저자들은 궤적 교정 전략을 도입한다.

• 주변 속도 정렬: 초기 노이즈만을 교정하는 기존 방식과 달리, FlowConsist는 궤적을 따라가는 모든 시간 단계에서 모델이 생성한 샘플의 주변 분포를 실제 데이터 분포와 정렬시킨다.
• 교사 없는 메커니즘 (Teacher-Free Mechanism): 실제 주변 속도를 추정하기 위해 사전 학습된 교사 모델에 의존하는 방식(DMD와 같이) 대신, FlowConsist는 메인 모델($F_\theta$)이 생성한 샘플의 주변 속도를 추정하도록 학습된 보조 모델($G_\psi$)을 활용한다.
• 교정 과정: 메인 모델은 자신의 예측된 주변 속도를 보조 모델이 추정하는 속도와 일치시키도록 학습된다. 이는 모델의 예측을 실제 데이터 분포 쪽으로 다시 "끌어당겨", 단일 스텝 매핑 시 발생할 수 있는 누적 오차를 교정하는 역할을 한다.

주요 기여

1. 궤적 드리프트 식별: 본 논문은 조건부 속도가 Flow Matching에서 본질적으로 비제로 공분산을 가지며, 이를 빠른 흐름 모델의 순간 속도 추정치로 사용할 때 필연적인 궤적 드리프트를 초래한다는 것을 이론적 증명(정리 1 & 2)을 통해 입증하였다.
2. 새로운 학습 목적 함수: 조건부 속도를 스스로 예측된 주변 속도로 대체하여, 최적화가 일관된 주변 궤적에 고정되도록 보장하는 원칙적인 대안(식 7)을 제안하였다.
3. 궤적 교정 전략: 모든 시간 단계에서 주변 분포를 정렬하기 위한 교사 없는 방법(식 11)을 도입하여, 긴 시간 간격 동안 발생하는 근사 오차의 누적을 완화하였다.
4. 최첨단 성능 (SOTA): 교사 모델로부터의 증류(distillation) 없이도, 단 **1회의 샘플링 스텝(1-NFE)**만으로 ImageNet 256×256에서 1.52의 Fréchet Inception Distance (FID)를 달성하며 새로운 SOTA를 기록하였다.

실험 결과

• 데이터셋 및 설정: 실험은 DiT 아키텍처와 사전 학습된 VAE 잠재 공간을 사용하여 ImageNet 256×256에서 수행되었다.
• 절제 연구 (Ablation Studies):
  • 표준 MeanFlow 손실을 제안된 주변 궤적 일관성 손실(식 7)로 교체했을 때, FID가 5.86에서 4.99로 개선되었다.
  • 주변 속도 정렬 손실(식 11)을 추가했을 때, FID는 4.31로 더욱 감소하였다.
  • 해당 방법은 다양한 CFG 스케일과 두 손실 구성 요소의 통합 비율에 대해 견고함을 보여주었다.
• 비교: FlowConsist-XL/2는 1-NFE에서 1.52의 FID를 달성하여, 이전의 최고 성능을 보였던 비증류(non-distilled) Fast Flow 모델(iMF-XL/2, 1.72)을 능가하였으며, 증류된 방법들의 성능에 근접하였다. 또한, 경쟁력 있는 품질을 유지하면서도 효율성 측면에서 다단계 확산 모델(250+ 스텝)보다 뛰어난 성능을 보였다.

의의

FlowConsist는 기존 프레임워크의 성능 제한의 근본 원인인 조건부 속도와 주변 속도 사이의 이론적 불일치, 그리고 실질적인 근사 오차의 누적 문제를 해결함으로써 Fast Flow 모델링에 대한 새로운 관점을 제시한다. 이러한 문제들을 교정함으로써, 저자들은 정교한 학습 목적 함수만으로도 고충실도(high-fidelity) 단일 스텝 생성이 가능하다는 것을 입증하였으며, 이를 통해 교사 없는(teacher-free) 패러다임과 교사 증류(teacher-distilled) 패러다임 사이의 간극을 메웠다. 이 연구는 궤적 일관성이 Fast Flow 모델이 시간 단계 간의 선형 매핑을 성공적으로 수행하기 위한 핵심 요구사항임을 강조한다.