The Amazing Stability of Flow Matching

논문 요약: Flow Matching 의 놀라운 안정성

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 최근 생성 모델 (Diffusion Models 등) 의 성공은 주로 대규모 데이터셋과 복잡한 아키텍처에 기인한다고 여겨집니다. Flow Matching (FM) 은 Diffusion 모델의 대안으로 효율성과 단순성을 제공하며 경쟁력 있는 성능을 보입니다.
• 문제: FM 모델을 훈련하는 데는 여전히 막대한 계산 자원과 데이터가 필요합니다. 따라서 데이터의 양을 줄이거나 (Pruning), 아키텍처를 변경하더라도 생성된 샘플의 품질과 다양성이 어떻게 유지되는지, 즉 **FM 모델의 안정성 (Stability)**을 규명하는 것은 중요합니다.
• 연구 목적: 데이터 분포의 교란 (Pruning, 데이터 교환) 과 모델 아키텍처의 변경 (용량 축소, 구조 변경) 하에서 Flow Matching 모델이 얼마나 견고한지 실험적으로 분석하고, 이를 통해 데이터 효율적인 훈련 전략을 모색합니다.

2. 방법론 (Methodology)

저자들은 CelebA-HQ 데이터셋을 기반으로 FM 모델을 훈련하고, 다음과 같은 변형 (Perturbations) 을 가하여 모델의 반응을 분석했습니다.

• 데이터 프루닝 (Data Pruning) 전략:
  전체 데이터셋 $S$에서 부분집합 $S'$을 선택하는 세 가지 지능형 (Informed) 방법을 제안하고, 이를 무작위 선택 (Random) 과 비교했습니다.

  1. Gradient-based (Grad): 작은 서로게이트 모델을 훈련하여 각 샘플의 그래디언트 노름 (Gradient Norm) 을 추정합니다. 그래디언트가 큰 샘플은 모델 학습에 큰 영향을 미치므로 이를 유지하거나 제거하는 전략을 테스트합니다.
  2. Loss-based (Loss): 공유된 노이즈 경로와 타임스텝을 기반으로 계산된 손실 (Loss) 값을 기준으로 샘플을 순위 매깁니다.
  3. Cluster-based (Clust): 사전 훈련된 CLIP 모델을 사용하여 이미지 임베딩을 추출하고 K-means 로 클러스터링합니다.
     • 클러스터 내 샘플 수: 전체 분포에 비례하거나 균등하게 할당.
     • 샘플 선택 기준: 클러스터 중심에 가장 가까운 샘플 (대표성) 또는 가장 먼 샘플 (희귀/어려운 샘플).
  • 역방향 테스트: 위 방법들의 역방향 (가장 낮은 점수를 가진 샘플 선택, $S^{-1}$) 을 적용하여 효과를 대조했습니다.

• 실험 설정:

  • 모델: Transformer 기반의 DiT (Diffusion Transformer) 와 U-Net 아키텍처를 사용하며, VQ-VAE 를 통해 이미지를 인코딩합니다.
  • 평가 지표:
    • FID (Fréchet Inception Distance): 생성된 이미지와 실제 데이터 분포 간의 품질 및 다양성 측정.
    • ArcFace Cosine Similarity: 동일한 시드 (Seed) 로 생성된 이미지 쌍 간의 시각적/의미적 유사성 측정 (안정성 지표).
  • 스트레스 테스트:
    1. 불연속 부분집합 (Disjoint Subsets): 데이터를 두 개의 불연속적인 무작위 부분집합으로 나누어 훈련.
    2. 클러스터 제거 (Cluster Removal): 특정 성별 (남성/여성) 로 분류된 데이터만 사용하여 훈련.
    3. 데이터 교환 (Data Swapping): CelebA-HQ 대신 FFHQ 데이터셋으로 훈련 (동일 도메인).
    4. 아키텍처 변경: 모델 용량 축소 (DiT-XL $\to$ DiT-S) 및 구조 변경 (DiT $\to$ U-Net).

3. 주요 기여 (Key Contributions)

1. Flow Matching 의 놀라운 안정성 입증: 데이터의 불연속적 분할, 전체 데이터셋 교체, 모델 아키텍처 축소 등 다양한 교란 하에서도 동일한 무작위 노이즈 ($x_0$) 에서 시작하면 시각적으로 매우 유사한 샘플이 생성됨을 확인했습니다.
2. 지능형 데이터 프루닝 방법론 도입: 생성 모델에 적용 가능한 세 가지 데이터 프루닝 방법 (Grad, Loss, Clust) 을 제안하고 정성적/정량적으로 분석했습니다.
3. 클러스터 기반 리샘플링의 효과: 클러스터 간 분포를 균형을 맞추는 리샘플링 방법 (Clustb) 은 오히려 생성된 이미지의 평가 지표 (FID) 를 개선할 수 있음을 보였습니다.

4. 실험 결과 (Results)

• 시각적 및 수치적 안정성:
  • 불연속 부분집합: 서로 다른 두 데이터 부분집합으로 훈련된 모델은 동일한 시드에서 거의 동일한 이미지를 생성했습니다 (ArcFace 유사도 $0.69 \pm 0.12$, 무관한 쌍 대비 $0.34$).
  • 데이터 교환 (CelebA-HQ vs FFHQ): 훈련 데이터셋을 완전히 다른 얼굴 데이터셋으로 바꾸더라도, 동일한 잠재 공간 (Latent Space) 에서 유사한 궤적을 학습하여 시각적 유사도가 유지되었습니다 ($0.58 \pm 0.15$).
  • 아키텍처 변경: 모델 크기를 줄이거나 (DiT-XL $\to$ DiT-S) U-Net 으로 변경하더라도 생성된 이미지의 핵심 속성은 유지되었습니다 (용량 변경 시 유사도 $0.81$, 구조 변경 시$0.55$).
• 프루닝 전략의 영향:
  • FID 변화: 무작위 프루닝은 FID 를 약간 악화시켰으나, Clustb(균형 잡힌 클러스터 샘플링) 는 FID 를 개선했습니다. 반면, Loss 기반 프루닝 (높은 손실 샘플 유지) 은 성능을 크게 저하시켰습니다. 이는 FM 에서 높은 손실 샘플이 밀도가 낮은 영역에 위치하여 흐름 (Flow) 을 방해할 수 있음을 시사합니다.
  • 안정성 유지: FID 가 악화되었던 방법 (Loss, Grad$^{-1}$) 으로 훈련된 모델조차도, Unpruned 모델과 비교했을 때 높은 유사도 ($0.79$ 이상) 를 보여주었습니다. 이는 생성된 이미지의 의미적 (Semantic) 구조가 데이터의 양적 변화에 덜 민감함을 의미합니다.

5. 의의 및 결론 (Significance)

• 이론적 통찰: Diffusion 모델의 안정성이 엔트로피 수송 (Entropic transport) 최적화와 관련이 있다면, Flow Matching 의 안정성은 속도장 (Velocity Field) $u(x, t)$의 국소적 조정과 전역적 구조의 불변성에서 비롯된 것으로 해석됩니다. 데이터가 제거되거나 변경되어도 흐름 (Flow) 은 필요한 부분만 국소적으로 조정되고 전역적 궤적은 유지됩니다.
• 실용적 가치:
  • 데이터 효율성: FM 모델은 데이터 양이 줄어들거나 분포가 변해도 견고하게 작동하므로, 계산 비용을 줄이기 위해 데이터 프루닝을 수행할 수 있음을 보여줍니다.
  • 모델 신뢰성: 다양한 교란 하에서도 일관된 생성 품질을 보장하므로, FM 모델의 신뢰성이 높음을 입증했습니다.
  • 향후 방향: 대규모 데이터로 훈련되는 차세대 생성 모델의 효율성을 극대화하기 위해 데이터 선택과 일반화 간의 상호작용을 이해하는 것이 중요함을 강조합니다.

이 논문은 Flow Matching 모델이 데이터와 아키텍처의 변화에 대해 놀라울 정도로 강건 (Robust) 하며, 이를 통해 더 효율적인 생성 모델 훈련 전략을 수립할 수 있음을 실증적으로 증명했습니다.