기술 요약: Stable Velocity: Flow Matching에 대한 분산 관점의 접근

1. 문제 정의

Flow Matching (FM) 및 Conditional Flow Matching (CFM)은 명시적으로 전방 SDE나 PF-ODE를 시뮬레이션하지 않고도 확률 흐름(probability flow)을 학습함으로써 생성 모델링을 위한 우아한 프레임워크를 제공하지만, 이들은 근본적인 한계인 학습 타겟의 높은 분산(high variance) 문제에 직면해 있습니다. CFM에서 학습 목적 함수는 실제 주변 속도장(marginal velocity field) $v_t(x_t | x_0)$에 대한 단일 샘플 몬테카를로 추정치에 의존하여 주변 속도 $v_t(x_t)$를 근사합니다.

저자들은 이 단일 샘플 추정치가 특히 주변 분포가 사전 분포(prior)와 가까운 시점(고분산 영역)에서 상당한 분산을 보인다는 점을 지적합니다. 이러한 높은 분산은 다음과 같은 문제를 야기합니다:

1. 최적화 불안정성: 노이즈가 섞인 그래디언트가 수렴을 늦춥니다.
2. 추론 비효전성: 특정 영역에서 흐름의 결정론적(deterministic) 특성을 활용하지 못해 불필요하게 작은 적분 단계를 사용하게 됩니다.
3. 비효과적인 보조 감독(Auxiliary Supervision): 기존의 표현 정렬 방법(예: REPA)은 감독을 균등하게 적용하며, 고분산(노이즈가 많은) 영역에서는 시맨틱 정렬이 부적절(ill-posed)할 수 있다는 사실을 인지하지 못합니다.

2. 방법론

본 논문은 변수의 분산을 명시적으로 규명하고 그에 따라 학습 및 샘플링 전략을 조정하는 Stable Velocity 프레임워크를 제안합니다.

2.1 분산 분석

저자들은 다양한 데이터 분포(GMM, CIFAR-10, ImageNet)에 걸쳐 CFM 타겟의 분산 $V_{CFM}(t)$를 분석합니다. 이를 통해 두 가지 영역 구조를 밝혀냈습니다:

• 고분산 영역 ($t \in [\xi, 1]$): 사전 분포(노이즈) 근처에서는 사후 분포 $p(x_0|x_t)$가 다봉형(multimodal)을 띠므로, 조건부 속도가 크게 변동합니다.
• 저분산 영역 ($t \in [0, \xi]$): 데이터 분포 근처에서는 사후 분포가 단일 참조 샘플에 집중되어, 조건부 속도 $v_t(x_t | x_0)$가 실제 주변 속도 $v_t(x_t)$와 거의 일치하게 됩니다.

2.2 학습: Stable Velocity Matching (StableVM)

편향(bias)을 도입하지 않으면서 분산을 줄이기 위해, 저자들은 **Stable Velocity Matching (StableVM)**을 제안합니다.

• 메커니즘: 단일 참조 샘플 $x_0$를 사용하는 대신, $n$개의 참조 샘플 집합 $\{x_0^i\}_{i=1}^n$을 사용하여 합성된 조건부 경로를 구축합니다.
• 타겟: 학습 타겟은 참조 배치에 대한 조건부 속도의 자기 정규화된 중요도 가중 평균(self-normalized importance-weighted average)입니다:
  $$\hat{v}_{StableVM}(x_t) = \frac{\sum_{k=1}^n p_t(x_t | x_0^k) v_t(x_t | x_0^k)}{\sum_{j=1}^n p_t(x_t | x_0^j)}$$
• 특성: 이론적으로 CFM과 동일한 전역 최솟값(global minimizer)을 공유하는 **불편 추정량(unbiased estimator)**임이 증명되었으며, 동시에 학습 타겟의 분산을 $O(1/n)$ 비율로 엄격하게 감소시킵니다.
• 클래스 조건부 확장: 클래스 조건부 생성의 희소성을 처리하기 위해 메모리 뱅크를 사용하여 클래스당 다양한 참조 샘플을 유지함으로써, 배치 내 클래스 빈도가 낮더라도 혼합(mixture)이 효과적으로 유지되도록 합니다.

2.3 학습: Variance-Aware Representation Alignment (VA-REPA)

두 가지 영역에 대한 통찰을 바탕으로, 저자들은 보조 시맨틱 감독을 최적화하기 위한 VA-REPA를 제안합니다.

• 통찰: 표현 정렬(예: 잠재 특징을 사전 학습된 인코더와 정렬하는 것)은 시맨틱 정보가 보존되는 저분산 영역에서만 유효합니다. 고분산 영역에서 시맨틱 신호는 노이즈가 많고 부적절합니다.
• 전략: 가중치 함수 $w(t)$를 정렬 손실에 적용하여, 감독을 저분산 영역($t < \xi$)으로 제한합니다. 이는 모델이 노이즈가 많고 정보가 없는 신호로부터 학습하는 것을 방지합니다.

2.4 추론: Stable Velocity Sampling (StableVS)

저자들은 미세 조정 없이 샘플링을 가속화하기 위해 저분산 영역을 활용합니다.

• 통찰: 저분산 영역에서 흐름 역학(flow dynamics)은 효과적으로 결정론적이며, 궤적은 일정한 속도를 가진 직선에 근사합니다.
• 메커니즘: StableVS는 이 영역에서의 역방향 SDE 및 PF-ODE에 대한 폐쇄형(closed-form) 단순화를 도출합니다. 선형 보간법(linear interpolants)의 경우, 업데이트 규칙은 다음과 같이 단순화됩니다:
  $$x_\tau = x_t + (\tau - t)v_t(x_t)$$
• 구현: 샘플러는 고분산 영역에서는 표준 솔버(예: Euler, DPM-Solver++)를 사용하지만, 저분산 영역에서는 폐쇄형 StableVS 업데이트로 전환합니다. 이를 통해 샘플 품질을 저하시키지 않으면서도 큰 적분 단계(더 적은 총 단계 수)를 사용할 수 있습니다.

3. 주요 기여

1. 분산 이론적 관점: 본 논문은 확률적 보간법(stochastic interpolants)에서 차원에 의존하는 두 가지 영역 구조(고분산 vs 저분산)를 밝혀내어, 이것이 학습 역학 및 추론 동작을 지배함을 이론적으로 제시합니다.
2. StableVM: 여러 참조 샘플을 결합하여 최적화를 안정화함으로써, 표준 CFM 및 기존의 분산 감소 기술(STF 등)보다 우수한 수렴 속도와 샘플 품질을 보이는 불편 추정 학습 목적 함수를 제안합니다.
3. VA-REPA: 저분산 영역에서 선택적으로 보조 감독을 적용하여 학습 효율성과 안정성을 높이는 영역 인식형 표현 정렬을 제안합니다.
4. StableVS: 저분산 영역의 결정론적 특성을 활용하여 미세 조정 없이 SOTA 모델에서 2배 이상의 샘플링 가속을 달고 실현합니다.

4. 실험 결과

본 프레임워크는 ImageNet 256×256 및 다양한 대규모 텍스트-이미지(T2I) 및 텍스트-비디오(T2V) 모델(SD3.5, Flux, Qwen-Image, Wan2.2)에서 검증되었습니다.

• 학습 효율성:
  • StableVM + VA-REPA는 모델 규모(SiT-B에서 SiT-XL까지)와 학습 기간 전반에 걸쳐 바닐라 REPA 및 기타 베이스라인(REG, iREPA)보다 일관되게 우수한 성능을 보였습니다.
  • ImageNet에서 이 방법은 엔드 투 엔드 미세 조정을 요구하는 베이스라인(예: REPA-E)에 비해 훨씬 적은 학습 에포크만으로도 경쟁력 있는 FID 및 IS 점수를 달랐습니다.
  • 절제 연구(Ablation studies)를 통해 정렬을 저분산 영역으로 제한하는 것이 균등하게 적용하는 것보다 더 나은 FID를 얻음을 확인했습니다.
• 샘플링 속도:
  • StableVS는 SD3.5, Flux, Wan2.2에서 2배 이상의 추론 가속을 달성했습니다(예: 30단계 베이스라인을 약 17~20단계로 매칭).
  • 결정적으로, 이러한 가속화는 샘플 품질을 저하시키지 않았습니다. 참조 지표(PSNR, SSIM, LPIPS) 및 생성 지표(GenEval, T2V-CompBench)는 단계 수를 줄인 표준 솔버와 비슷하거나 더 나은 결과를 보였습니다.
  • 정성적 결과에서 StableVS의 출력은 30단계 베이스라인과 시각적으로 구별할 수 없을 정도로 유사했으나, 단순한 단계 축소는 눈에 띄는 아티팩트를 발생시켰습니다.

5. 의의 및 주장

본 논문은 생성 궤적의 분산 구조를 명시적으로 모델링하는 것이 더 효율적인 생성 모델을 설계하기 위한 원칙적인 토대를 제공한다고 주장합니다.

• 통합 프레임워크: 하나의 이론적 관점 아래 학습에서의 분산 감소와 샘в링에서의 영역 인식 가속을 통합합니다.
• 실질적 영향: 제안된 방법들은 모듈식이며 기존 아키텍처(예: SiT, DiT) 및 솔버와 호환되어, 구조적 변경이나 비용이 많이 드는 미세 조정 없이도 즉각적인 학습 안정성 및 추론 속도 향상을 제공합니다.
• 이론적 통찰: 본 연구는 확산/흐름 궤적을 균일한 과정으로 보는 관점에 도전하며, 서로 다른 영역(고분산 vs 저분산)에는 서로 다른 최적화 및 샘플링 전략이 필요함을 입증합니다.

저자들은 자신들의 접근 방식이 흐름 매칭 목적 함수에서 발생하는 참조 샘플의 확률성(stochasticity)에 의한 특정 분산원을 해결함으로써, 차원 축소나 개선된 솔버와 같은 기존 기술을 보완한다고 강조합니다.