기술 요약: 드리프트 흐름 매칭 (DFM)

문제 정의

Flow Matching 및 Diffusion 모델과 같은 반복적 생성 모델은 추론 단계 수를 늘림으로써 생성 품질, 제어 가능성, 분포 커버리지를 향상시키는 강력한 "테스트 시간 확장 (test-time scaling)" 능력을 입증해 왔습니다. 그러나 이러한 방법은 여러 번의 추론 단계를 필요로 하므로 계산 비용이 많이 듭니다. 반면, 드리프트 모델 (Drift Models) 은 소스 분포에서 타겟 분포로의 직접적인 푸시포워드 맵을 학습함으로써 효율적인 1 단계 생성을 제공합니다. 하지만 이들의 직접 생성 패러다임은 테스트 시간에 추론 단계를 늘려 결과를 개선할 수 있는 자연스러운 메커니즘이 부재하여, 다양한 품질 - 효율성 요구 사항에 대한 유연성과 적응성이 제한됩니다.

방법론: 드리프트 흐름 매칭 (DFM)

저자들은 1 단계 드리프트 모델의 효율성과 다단계 Flow Matching 의 확장성을 통합하는 드리프트 흐름 매칭 (Drift Flow Matching, DFM) 프레임워크를 제안합니다. DFM 은 Flow Matching 확률 경로 상에서 임의의 현재 시간 $t$에서 임의의 미래 시간 $r > t$로 샘플을 이동시킬 수 있는 **이중 시간 수송 모델 (two-time transport model)**을 학습함으로써 이를 달성합니다.

핵심 메커니즘

1. 이중 시간 주변 분포 구성:
   표준 Flow Matching 이 점별 조건부 속도 목표를 회귀하는 것과 달리, DFM 은 주변 분포 쌍 $(p_t, p_r)$에서 작동합니다. 이는 소스 ($X_0$) 와 타겟 ($X_1$) 분포의 독립적인 결합을 기반으로 $X^Z_t = \alpha(t)X^Z_0 + \beta(t)X^Z_1$ 및 $X^Z_r = \alpha(r)X^Z_0 + \beta(r)X^Z_1$로 정의된 조건부 경로를 설정합니다. 이러한 구성은 임의의 시간 단계에서 주변 분포 $p_t$와 $p_r$을 유도합니다.

2. 평균 속도 파라미터화:
   DFM 은 현재 분포 $p_t$를 미래 분포 $p_r$로 수송하기 위해 평균 속도장 $u_\theta(x_t, t, r)$을 파라미터화합니다. 수송 맵은 다음과 같이 정의됩니다:
   $$T^\theta_{t,r}(x_t) = x_t + (r - t)u_\theta(x_t, t, r)$$
   이 공식은 모델이 임의의 단계 크기 $(r-t)$를 처리할 수 있게 합니다. $r \to t$로 극한에 도달하면 평균 속도는 표준 Flow Matching 의 순간 주변 속도를 회복합니다.

3. 드리프트 기반 감독:
   속도 벡터를 회귀하는 대신, DFM 은 $T^\theta_{t,r}$에 의해 생성된 예측 분포 $q^\theta_{t,r}$를 실제 타겟 주변 분포 $p_r$로 수송하기 위해 드리프트 기반 목적 함수를 사용합니다.

   • 드리프트 장: 드리프트 장 $V_{q^\theta_{t,r}, p_r}$은 타겟 샘플 ($V^+_{p_r}$) 로 향하는 커널 가중 인력과 현재 모델 분포 ($V^-_{q^\theta_{t,r}}$) 로부터의 반발적 자기 수정 항을 사용하여 구성됩니다.
   • 손실 함수: 학습 목적 함수는 스톱-그레이디언트 (stop-gradient) 공식을 사용하여 드리프트 장의 크기를 최소화합니다:
     $$L_{DFM}(\theta) = \frac{1}{2} \mathbb{E}_{(t,r), X_t \sim p_t} \left[ \| \hat{X}^\theta_r - \text{sg}(\hat{X}^\theta_r + V_{q^\theta_{t,r}, p_r}(\hat{X}^\theta_r)) \|^2 \right]$$
   • 그룹 추정: 핵심적으로, 드리프트 장과 손실은 각 샘플링된 시간 쌍 $(t, r)$에 대해 **그룹 단위 (group-wise)**로 계산되어야 합니다. 서로 다른 시간 쌍의 샘플을 혼합하면 각 쌍마다 타겟 주변 분포 $p_r$이 다르기 때문에 드리프트 추정이 왜곡됩니다.

주요 기여

• 통합 프레임워크: DFM 은 효율적인 1 단계 드리프트 모델과 반복적 Flow Matching 방법 간의 간극을 메워, 직접 생성과 반복적 정제를 모두 지원하는 단일 패러다임을 창출합니다.
• 테스트 시간 확장: 이 프레임워크는 드리프트 기반 방법에게 테스트 시간 확장 능력을 부여합니다. 사용자는 재학습 없이 추론 단계 수 (NFE) 를 조정하여 샘플링 속도와 생성 품질 간의 트레이드오프를 실현할 수 있습니다.
• 새로운 학습 원리: 본 논문은 Flow Matching 에서 표준적으로 사용되는 점별 속도 회귀에서 드리프트 감독을 통한 주변 분포 수송으로 전환하는 새로운 학습 원리를 제시합니다. 이는 다른 1 단계 생성 접근법에서 사용되는 자기 증류 (self-distillation) 방법과 관련된 불안정성을 피합니다.
• 이론적 보장: 저자들은 DFM 이 엔드포인트 생성을 Wasserstein 거리 ($W_2$) 로 제어되는 단거리 주변 수송 문제로 분해하며, 무한소 단계 극한에서 표준 Flow Matching 을 회복함을 보여주는 이론적 분석을 제공합니다.

실험 결과

합성 데이터, MNIST, FFHQ, ImageNet-1k, 그리고 로봇 제어 작업에 대한 광범위한 실험은 DFM 의 효과성을 입증합니다:

• 합성 데이터: 시각화는 DFM 이 단일 단계에서 고품질 샘플을 생성할 수 있음을 보여주며, NFE 가 증가함에 따라 분포 커버리지와 샘플 품질이 추가로 향상됨을 확인합니다.
• 조건부 생성 (MNIST 및 FFHQ): DFM 은 1 단계 추론으로 경쟁력 있는 성능을 달성하여 (Drift Model 베이스라인과 일치) NFE 를 2, 5, 또는 10 으로 증가시킬 때 EMD, FID 등의 지표를 일관되게 개선합니다.
• 대규모 이미지 합성 (ImageNet-1k): ImageNet-256x256 데이터셋에서 DFM 은 NFE=1 일 때 1.52 의 경쟁력 있는 FID 를 달성합니다. NFE 를 10 으로 늘리면 FID 가 1.31 로 추가 감소하여, 반복적 Flow Matching 방법과 비교 가능한 효과적인 테스트 시간 확장을 입증합니다.
• 로봇 제어: 로봇 조작 작업 (예: Lift, Can, PushT) 에서 DFM 은 Drift Policy 의 1 단계 추론 능력을 유지하면서 증가된 추론 단계로 명확한 성능 향상을 보이며, 특정 작업에서 1 단계 Drift Policy 와 다단계 Diffusion Policy 베이스라인 모두를 능가합니다.

중요성과 주장

본 논문은 DFM 이 직접 생성과 반복적 정제가 더 이상 상호 배타적인 설계 선택이 아닌 단일 프레임워크 내에서 볼 수 있는 보다 유연한 생성 패러다임을 확립한다고 주장합니다. 드리프트 생성 모델링과 흐름 기반 반복적 생성을 연결함으로써, DFM 은 생성 과정이 다양한 품질 - 효율성 요구 사항에 맞춰 샘플링 계산을 적응하도록 합니다. 저자들은 이 접근법이 이전 드리프트 모델의 주요 한계를 해결하면서 강력한 1 단계 생성을 달성함과 동시에 반복적 방법의 확장성 혜택을 유지하는 새로운 방법을 제공한다고 강조합니다.