1. 문제 정의 (Problem Statement)

최근 Flow Matching (FM) 은 이미지 및 비디오 생성 분야에서 큰 성공을 거두었으나, 다음과 같은 근본적인 한계가 존재합니다.

• 유클리드 거리 기반의 한계: 기존 FM 은 예측된 속도와 목표 속도 간의 제곱 오차 (MSE, L2 Loss) 를 최소화합니다. 이는 유클리드 공간에서의 거리를 측정하지만, 실제 데이터는 고차원 공간 내의 낮은 차원 매니폴드 (manifold) 에 존재합니다.
• 잘못된 일반화 (Incorrect Generalization): 유한한 용량을 가진 모델은 훈련 데이터를 단순히 복제하는 것이 아니라 새로운 데이터를 생성해야 합니다. 그러나 MSE 손실 함수는 매니폴드 구조를 고려하지 않는 등방성 (isotropic) 유클리드 거리를 측정하므로, 모델이 데이터 분포와 일치하지 않는 아웃 - 오브 - 디스트리뷰션 (out-of-distribution) 샘플을 생성하게 됩니다.
• 가이드 (Guidance) 의존성: 이러한 문제를 해결하기 위해 Classifier-Free Guidance (CFG) 를 사용하면 샘플 품질이 향상되지만, 이는 원래 데이터 분포를 변경하여 왜곡된 분포를 생성하게 됩니다. 따라서 가이드 없이도 원본 데이터 분포에 충실한 고품질 샘플을 생성하는 것이 중요한 과제로 남았습니다.

2. 방법론 (Methodology)

저자들은 Continuous Adversarial Flow Models (CAFM) 을 제안하여 연속 시간 흐름 모델에 적대적 훈련 (Adversarial Training) 을 도입했습니다.

핵심 아이디어

• 학습된 판별자 (Learned Discriminator) 의 도입: 고정된 MSE 기준 대신, 생성자 (Generator) 와 함께 학습되는 판별자 (Discriminator) 를 사용하여 훈련 목표를 설정합니다. 이는 GAN 의 적대적 훈련 원리를 연속 시간 흐름 모델에 적용한 것입니다.
• JVP (Jacobian-Vector Product) 기반 판별:
  • 기존 이산 시간 적대적 흐름 모델 (AFM) 과 달리, CAFM 은 속도 벡터 ($v_t$) 자체를 판별하는 것이 아니라, 판별자 네트워크 $D(x_t, t)$ 의 방향 미분 (Directional Derivative) 을 판별합니다.
  • 구체적으로, $D$ 의 JVP 를 계산하여 $D_{\text{jvp}}(x_t, t, \dot{x}_t, \dot{t}) = \frac{\partial D}{\partial x_t}\dot{x}_t + \frac{\partial D}{\partial t}\dot{t}$ 형태로 정의합니다. 여기서 $\dot{x}_t$ 는 실제 속도 ($\bar{v}_t$) 또는 생성된 속도 ($G(x_t, t)$) 입니다.
  • 이는 속도가 위치의 미분이라는 물리적 성질을 반영하며, 연속 시간 흐름의 특성을 효과적으로 포착합니다.

손실 함수 (Loss Functions)

• 적대적 목표: 생성자 $G$ 와 판별자 $D$ 는 최소 - 최대 (minimax) 게임을 수행합니다.
  • $D$ 는 실제 속도 ($\bar{v}_t$) 에 대해 높은 점수, 생성된 속도 ($G$) 에 대해 낮은 점수를 주도록 학습합니다.
  • $G$ 는 $D$ 가 높은 점수를 주도록 속도를 조정합니다.
  • 대조 함수 (Contrastive Function) 로는 $f(a, b) = (a-1)^2 + (b+1)^2$ 형태의 최소 제곱 GAN (LSGAN) 스타일 함수를 사용합니다.
• 중심화 페널티 (Centering Penalty): $D$ 의 절대값이 무한히 치우치는 것을 방지하기 위해 $D(x_t, t)^2$ 항을 추가합니다.
• 최적 수송 정규화 (Optimal Transport Regularization): 고차원 공간에서 생성자가 영공간 (null space) 을 악용하는 것을 방지하고 최소 노름 해를 유도하기 위해 $\|G(x_t, t)\|^2_2$ 항을 추가합니다. (후처리 (Post-training) 시에는 이 항을 제거하여 유클리드 편향을 완전히 제거합니다.)

구현 및 효율성

• JVP 계산: PyTorch 의 torch.func.jvp 와 vmap 을 사용하여 효율적으로 계산하며, DDP/FSDP 와 호환됩니다.
• 아키텍처: 생성자 $G$ 는 기존 Flow Matching 모델 (예: SiT, JiT) 을 그대로 사용하며, 판별자 $D$ 는 $G$ 와 동일한 아키텍처를 가지되 LayerNorm 을 RMSNorm 으로 교체하여 훈련 안정성을 높였습니다.
• 후처리 (Post-training) 전략: CAFM 은 처음부터 훈련하는 것보다 기존 FM 모델을 후처리 (Post-training) 하는 데 가장 효과적입니다. FM 으로 사전 훈련된 모델에 CAFM 목적 함수를 적용하여 10~20 에포크만 추가 훈련하면 성능이 크게 향상됩니다.

3. 주요 기여 (Key Contributions)

1. 연속 시간 적대적 흐름 모델 (CAFM) 제안: 이산 시간 AFM 을 연속 시간으로 확장하여, 연속 시간 흐름 모델에 적대적 훈련을 성공적으로 통합했습니다.
2. 매니폴드 인식 일반화: 고정된 유클리드 손실 대신 학습된 판별자를 통해 데이터 매니폴드 구조를 더 잘 포착하도록 유도하여, 가이드 없이도 데이터 분포에 더 잘 부합하는 샘플을 생성합니다.
3. 효율적인 후처리 프레임워크: 기존 Flow Matching 모델 (SiT, JiT 등) 을 재학습 없이, 소수의 에포크만 추가 훈련하여 성능을 획기적으로 개선하는 실용적인 방법을 제시했습니다.
4. JVP 기반 적대적 훈련 설계: 속도를 직접 판별하는 대신 JVP 를 통해 속도의 방향성을 판별함으로써, 연속 시간 흐름의 물리적 특성을 보존하고 그래디언트 소실 문제를 완화했습니다.

4. 실험 결과 (Results)

ImageNet 256px 생성 (클래스 조건부)

• SiT (Latent Space): 가이드 없는 (Guidance-free) FID 가 8.26 에서 3.63으로, 가이드 있는 경우 (CFG 1.5) FID 가 2.06 에서 1.53으로 개선되었습니다.
• JiT (Pixel Space): 가이드 없는 FID 가 7.17 에서 3.57로, 가이드 있는 경우 FID 가 1.86 에서 1.80으로 개선되었습니다.
• 비교: 기존 최첨단 (SOTA) 모델들보다 가이드 없는 설정에서 가장 우수한 성능을 보였으며, 가이드 설정에서도 경쟁력 있는 결과를 달성했습니다.

Text-to-Image 생성 (Z-Image 모델)

• GenEval 벤치마크: 가이드 없는 설정에서 점수가 0.81 에서 0.85로 향상되었습니다.
• DPG-Bench 벤치마크: 전체 점수가 83.7 에서 85.2로 향상되었습니다.
• 시각적 품질: 텍스트 프롬프트에 대한 객체 배치, 색상, 속성 일치도가 개선되었으며, 더 자연스러운 질감과 선명도를 보여줍니다.

5. 의의 및 결론 (Significance)

• 분포 일치 (Distribution Matching) 의 혁신: 기존 Flow Matching 이 가진 "유클리드 거리 기반의 잘못된 일반화" 문제를 해결하여, 가이드 없이도 원본 데이터 분포에 충실한 고품질 샘플을 생성할 수 있음을 입증했습니다.
• 실용성: 새로운 모델을 처음부터 훈련하는 것보다 기존 Flow Matching 모델을 CAFM 으로 미세 조정 (Post-training) 하는 것이 매우 효율적이며, 소량의 추가 훈련으로 큰 성능 향상을 얻을 수 있습니다.
• 미래 연구 방향: 이 연구는 연속 시간 흐름 모델과 적대적 훈련의 결합 가능성을 열었으며, 향후 더 정교한 매니폴드 학습과 고품질 생성 모델 개발의 새로운 기준을 제시합니다.

요약하자면, CAFM은 Flow Matching 의 이론적 한계를 적대적 훈련과 JVP 기반의 연속 시간 판별을 통해 극복하여, 가이드 의존성을 줄이고 데이터 분포에 더 충실한 생성 모델을 가능하게 한 획기적인 방법론입니다.