기술 요약: 잠재 흐름 매칭에서의 안정성 탐색 및 활용

1. 문제 제기

생성 모델, 특히 확산 모델 (Diffusion Models) 과 그 흐름 매칭 (Flow Matching, FM) 대응 모델의 학습은 여전히 계산 비용이 많이 든다. 이러한 모델들은 일반적으로 방대한 데이터셋, 광범위한 컴퓨팅 자원, 긴 학습 기간, 그리고 조건부 설정에서는 광범위한 주석이 필요하다. 흐름 매칭은 확률 미분 방정식 (SDE) 에 비해 종종 샘플링 단계를 줄이는 상미분 방정식 (ODE) 을 통한 샘플 생성에 대한 결정론적 공식을 제공하지만, 전체 규모 데이터셋에 대한 이러한 모델의 학습 비용은 여전히 병목 현상이다.

이 연구가 다루는 핵심 질문은 생성 품질을 희생하지 않고 데이터셋 크기와 모델 용량을 크게 줄일 수 있는지 여부이다. 구체적으로, 저자들은 데이터셋 부분 표본 추출 (가지치기), 아키텍처 변경, 그리고 변경된 학습 구성을 포함한 주요 교란 하에서 잠재 흐름 매칭 (LFM) 모델의 안정성을 조사한다. 저자들은 LFM 수송 궤적이 본질적으로 안정적이라면, 이 특성을 활용하여 더 효율적인 학습 및 추론 알고리즘을 도출할 수 있다고 가정한다.

2. 방법론

2.1. 이론적 기반: 흐름 매칭에서의 안정성

이 연구는 직교 흐름 (Rectified Flow) 이 속도장에 대한 폐쇄형 표현을 허용한다는 관찰에 기반한다. 유한 데이터셋의 경우, 최적 속도장은 현재 상태와 스케일링된 학습 샘플 간의 제곱 거리를 소프트맥스 함수로 계산한 가중치에 의해 결정되는 학습 샘플들의 가중 합이다. 저자들은 이 소프트맥스가 수송 단계 초기에 "뾰족해진다 (peaked)"고 지적하며, 이는 특정 궤적에 대해 단일 학습 샘플이 종종 속도장을 지배한다는 것을 의미한다.

이 이론적 통찰은 수송 궤적이 특정 "지배적" 샘플의 존재에 의해 크게 결정됨을 시사한다. 결과적으로, 비지배적 샘플을 제거 (가지치기) 하는 것은 대부분의 궤적을 변경하지 않아야 한다. 저자들은 이를 공식화하여 **할당 지표 (assignment metric)**를 정의한다: 노이즈 $x_0$에서 시작하는 궤적에 대해, 해당 궤적이 어떤 학습 샘플 $x_i$에서 끝나는지 결정한다. 저자들은 해당 궤적의 지배적 샘플이 유지되는 한, 데이터셋의 큰 부분이 제거되더라도 이 할당이 불변일 것이라고 가정한다.

2.2. 데이터 가지치기 전략

이 안정성을 활용하기 위해 저자들은 판별 모델에서 차용되었으나 공유 노이즈 경로와 시간 단계를 활용하는 FM 목적 함수에 맞게 재형식화된 세 가지 데이터 가지치기 기준을 제안한다:

1. 기반 점수 매기기 (G): 손실의 기울기 노름 크기로 샘플을 순위 매긴다.
2. 손실 기반 점수 매기기 (L): 공유 노이즈 경로를 따라 계산된 손실 크기로 샘플을 순위 매긴다.
3. 클러스터링 기반 점수 매기기 (C): 사전 훈련된 임베딩 공간 (예: CLIP) 에서 샘플을 클러스터링한다. 저자들은 두 가지 구체적인 변형을 제안한다:
   • 비례 ($C_p$): 클러스터 크기에 비례하여 샘플을 선택한다.
   • 균형 ($C_b$): 분포 불균형을 완화하기 위해 각 클러스터에서 동일한 수의 샘플을 선택한다.
   • 커널 기반 ($C_b^\kappa$): 잠재 공간에서 클러스터의 가우시안 커널 평균과 일치하도록 샘플을 탐욕적으로 선택한다.

저자들은 또한 안정성 주장의 견고성을 테스트하기 위해 이러한 기준의 "역" (가장 낮은 점수를 선택) 을 평가한다.

2.3. coarse-to-fine (C2F) 추론

모델 용량 전반에 걸친 궤적의 안정성을 활용하여, 저자들은 2 단계 추론 전략을 제안한다:

1. Coarse 모델: 가지치기된 데이터 부분 집합으로 훈련된 경량 아키텍처 (예: DiT-S/2) 가 궤적의 첫 번째 단계 ($t \in [0, t_0)$) 를 처리한다.
2. Fine 모델: 전체 데이터셋으로 사전 훈련된 고용량 아키텍처 (예: DiT-XL/2) 가 두 번째 단계 ($t \in [t_0, 1]$) 를 처리한다.

"이음점 (seam)" $t_0$에서 원활한 전환을 보장하기 위해 Coarse 모델은 **연속성 손실 (continuity loss)**로 미세 조정된다. 이 손실은 $t_0$에서 Coarse 모델과 Fine 모델의 속도장 차이를 최소화한다. Fine 모델의 궤적은 ODE 역변환 (역방향 적분) 을 통해 목표 상태 $x_{t_0}$를 생성하는 데 사용되어, Coarse 모델이 Fine 모델의 경로에 정렬하도록 학습한다.

3. 주요 기여

1. LFM 안정성의 경험적 증명: 저자들은 독립적으로 훈련된 LFM 모델이 데이터셋 부분 표본 추출, 아키텍처 변경, 그리고 다양한 무작위 초기화 하에서 매우 일관된 수송 궤적으로 수렴함을 보인다. 이 안정성은 FM 목적 함수의 이론적 속성과 연결된다.
2. 실용적 가지치기 알고리즘: FM 에 적합한 세 가지 가지치기 기준 (기울기, 손실, 클러스터링) 을 도입한다. 그들은 공격적인 가지치기 (예: 데이터의 50-95% 제거) 가 성능을 저하시키지 않으며, 제한된 컴퓨팅 예산 하에서 수렴 속도와 생성 품질을 오히려 향상시킬 수 있음을 입증한다.
3. Coarse-to-Fine 추론 가속화: 궤적을 분할하는 2 모델 생성 전략을 제안한다. 이 접근법은 안정성 가정이 모델 아키텍처 전반에 걸쳐 성립한다면 생성 품질을 유지하면서 상당한 추론 속도 향상 (약 2.15 배) 을 달성한다.
4. 균형 가지치기를 통한 편향 완화: 저자들은 균형 잡힌 클러스터링 가지치기 ($C_b$) 가 명시적인 레이블 인식 균형 조정 없이도 훈련 세트에서 의미 클러스터의 표현을 균등화함으로써 생성 데이터의 편향 (예: 성별, 피부색) 을 완화할 수 있음을 보인다.

4. 실험 결과

이 방법들은 CelebA-HQ, FFHQ, ImageNet-1K에서 평가되었다.

• 가지치기 하의 안정성: CelebA-HQ 에서 50% 로 가지치기된 불연속 부분 집합으로 훈련된 모델들은 동일한 노이즈 시드 $x_0$에서 시작할 때 시각적으로 유사한 이미지를 생성했다. 가지치기된 모델과 가지치기되지 않은 모델 간의 ArcFace 코사인 유사도는 0.79 에서 0.83 사이였으며, 이는 무작위 쌍 (0.37) 보다 훨씬 높았다.
• 성능 지표:
  • CelebA-HQ: 균형 잡힌 클러스터링 ($C_b$) 을 사용하여 데이터의 50% 를 가지치기하면 가지치기되지 않은 경우의 Fréchet Inception Distance (FID) 가 24.24 에서 22.80 으로 약간 개선되었다.
  • ImageNet: 데이터의 75% 를 가지치기 ($pr=0.75$) 한 결과 전체 데이터셋에 비해 더 빠른 수렴을 보였으며 FID 를 유지하거나 약간 개선했다. 95% 가지치기에서도 모델은 초기에 더 빠르게 수렴했으나 일정 횟수의 반복 이후 성능이 저하되었다.
• 교란에 대한 견고성:
  • 아키텍처: 모델 용량 증가 (DiT-S/2 에서 DiT-XL/2 로) 는 궤적 변화를 최소화했다 (유사도 $s=0.81$). U-Net 아키텍처로 전환하면 유사도가 감소 ($s=0.55$) 했지만 거친 속성은 유지되었다.
  • 잠재 공간: 안정성은 잠재 공간 파라미터화 변경 (예: VAE 시드 변경 또는 가역 변환 적용) 에 민감하여 안정성이 깨졌다.
  • 목적 함수: 흐름 매칭에서 점수 기반 확산으로 전환하면 궤적이 완전히 변경되어 안정성이 FM 목적 함수에 특화됨을 나타냈다.
• Coarse-to-Fine 효율성: CelebA-HQ 에서 $t_0=0.7$인 C2F 접근법은 FID 를 유지하면서 추론 시간을 이미지당 93.95ms(Fine 만) 에서 43.53ms(약 2.15 배 속도 향상) 로 줄였다.
• 공정성: 균형 잡힌 가지치기 ($C_b$) 는 레이블 무관임에도 성별, 나이, 피부색과 같은 속성에 대한 균일 분포에 대한 KL 발산을 가지치기되지 않은 모델보다 더 효과적으로 감소시켰다.

5. 중요성 및 주장

저자들은 관찰된 안정성이 고정된 잠재 표현 하에서 직교 흐름 매칭 패러다임의 고유한 특성이지 특정 학습 알고리즘의 인공물이 아니라고 주장한다.

• 효율성: 주요 중요성은 품질을 희생하지 않고 학습 시간과 데이터 요구 사항을 줄일 수 있다는 점이다. 저자들은 ImageNet 과 같은 대규모 데이터셋의 경우 가지치기가 수렴을 가속화할 수 있음을, CelebA-HQ 와 같은 소규모 데이터셋의 경우 생성 품질을 향상시킬 수 있음을 입증한다.
• 추론 속도 향상: C2F 접근법은 생성의 초기 거친 단계를 경량 모델에 위임함으로써 추론 비용을 줄이는 원칙적인 방법을 제공한다. 궤적의 안정성에 의존하여 고용량 모델이 최종 세부 사항만 정제하면 된다는 점을 보장한다.
• 편향 제어: 이 연구는 특히 균형 잡힌 클러스터링을 통한 데이터셋 가지치기가 복잡한 레이블 인식 재가중 없이도 생성 데이터의 분포 불균형을 제어하고 완화하는 메커니즘으로 작용할 수 있음을 시사하며, 더 공정한 생성 모델을 향한 길을 제시한다.

이 논문은 안정성이 데이터와 모델 용량 변경에 대해 견고하지만 잠재 공간 구조와 생성 목적 함수 자체의 변경에는 민감하다고 결론지었다. 이 통찰은 실무자가 고품질 생성을 유지하면서 더 자원 효율적인 파이프라인을 설계할 수 있게 한다.