기술 요약: 드리프팅(Drifting)을 통한 생성 모델링

문제 정의

생성 모델링은 근본적으로 사전 분포 $p_{prior}$(예: 가우시안 노이즈)의 푸시포워드(pushforward)가 데이터 분포 $p_{data}$와 일치하도록 하는 매핑 $f$를 학습하는 작업, 즉 $f\#p_{prior} \approx p_{data}$를 수행하는 과제입니다. 확산 모델(Diffusion Models)이나 플로우 매칭(Flow Matching)과 같은 기존 패러다임은 복잡한 푸시포워드 맵을 실행 가능한 변환의 체인으로 분해하여, **추론 시간(inference time)**에 반복적으로 샘플을 정제함으로써 이를 달성합니다. 그러나 이 방식은 샘플당 여러 번의 신경망 평가(NFE)를 필요로 하므로 높은 계산 비용이 발생합니다.

단일 단계(one-step) 생성기를 만들기 위한 기존의 시도들은 사전 학습된 다단계 모델을 증류(distilling)하거나, 훈련 중에 확률 미분 방정식(SDE) 또는 상미분 방정식(ODE)의 궤적을 근사하는 데 의존해 왔습니다. 저자들은 반복적인 과정을 추론이 아닌 **훈련 시간(training time)**으로 근본적으로 전환하여, SDE/ODE 정식화에 의존하지 않고도 단 한 번의 신경 함수 평가(1-NFE)로 고품질 생성을 가능하게 하는 새로운 패러다임을 제안합니다.

방법론: 드리프팅 모델 (Drifting Models)

이 논문의 핵심 기여는 생성 모델링을 훈련 과정 자체 동안 발생하는 푸시포워드 분포의 진화로 정식화하는 **드리프팅 모델(Drifting Model)**입니다.

1. 훈련 시간의 진화 (Training-Time Evolution)

저자들은 훈련을 단순히 정적인 손실(loss)을 최소화하기 위해 가중치를 업데이트하는 과정으로 보는 대신, 반복적인 최적화(예: SGD)에 의해 생성된 모델 시퀀스 $\{f_i\}$를 일련의 푸시포워드 분포 $\{q_i\}$로 간주합니다. 각 반복마다, 샘로 $x_i = f_i(\epsilon)$는 파라미터 업데이트로 인해 $x_{i+1} = f_{i+1}(\epsilon)$로 암묵적으로 "드리프트(drift)"됩니다. 목표는 이 드리프트가 분포 $q$를 자연스럽게 $p_{data}$로 유도하는 메커니즘을 정의하는 것입니다.

2. 드리프팅 필드 (The Drifting Field)

이 진화를 제어하기 위해, 저자들은 현재 생성된 분포 $q$와 타겟 데이터 분포 $p$를 기반으로 샘플 $x$에 대한 변위 벡터 $\Delta x$를 계산하는 함수인 드리프팅 필드 $V_{p,q}(x)$를 도입합니다.

• 평형 조건 (Equilibrium Condition): 이 필드는 $p = q$일 때만 $V_{p,q}(x) = 0$이 되도록 설계되었습니다. 분포가 일치하면 샘플의 드리프트가 멈추고 평형 상태에 도달합니다.
• 반대칭성 (Anti-Symmetry): 필드는 $V_{p,q}(x) = -V_{q,p}(x)$를 만족합니다. 이 성질은 분포가 서로 바뀌면 드리프트의 방향이 반전됨을 보장하며, $p=q$일 때 $V=0$이 됨을 보장합니다.

3. 인스턴스화: 인력과 척력 (Attraction and Repulsion)

본 논문에서 사용된 구체적인 드리프팅 필드는 평균 이동(mean-shift) 알고리즘에서 영감을 얻었으며, $p$의 양의 샘플과 $q$의 음의 샘플을 사용하여 다음과 같이 구성됩니다:
$$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
여기서:

• $V^+_p(x)$는 $x$를 데이터 분포 $p$ 쪽으로 끌어당기는 **인력(attractive force)**입니다.
• $V^-_q(x)$는 $x$를 생성된 분포 $q$로부터 밀어내는 **척력(repulsive force)**입니다.
• 이러한 힘들은 커널 함수 $k(x, y)$($\ell_2$ 거리에 기반한 지수 커널)를 사용하여 계산되며, 안정성을 위해 소프트맥스(softmax) 연산을 통해 정규화됩니다.

4. 훈련 목적 함수 (Training Objective)

훈련 목적 함수는 평형 상태에서의 고정점 관계인 $f(\epsilon) = f(\epsilon) + V_{p,q}(f(\epsilon))$로부터 유도됩니다. 손실 함수는 현재의 예측값과 "동결된 타겟"(예측값에 드리프트를 더한 값) 사이의 거리를 최소화합니다:
$$L = \mathbb{E}_{\epsilon} \left[ \| f_\theta(\epsilon) - \text{stopgrad}(f_\theta(\epsilon) + V_{p,q}(f_\theta(\epsilon))) \|^2 \right]$$
이 정식화는 최적화 도구가 샘플을 드리프트된 버전으로 이동시킴으로써 분포를 진화시키도록 합니다. 결정적으로, 그래디언트는 분포 $q$ 자체를 통과하여 역전파되지 않으므로, 분포를 미분하는 복잡함을 피할 수 있습니다.

5. 특징 공간 및 가이던스 (Feature Space and Guidance)

• 특징 공간 (Feature Space): 고차원 데이터를 처리하기 위해, 드리프팅 손실은 사전 학습된 자기지도 학습 인코더(예: ResNet 스타일의 MAE)를 사용하는 특징 공간에서 계산됩니다. 이는 의미적으로 유사한 샘플들이 서로 가깝게 유지되도록 하여 더 풍부한 그래디언트 신호를 제공합니다.
• 분류기 없는 가이던스 (Classifier-Free Guidance, CFG): 이 방법은 훈련 중에 무조건적(unconditional) 데이터 샘플을 음의 샘플 분포에 혼합함으로써 CFG를 자연스럽게 지원하며, 이를 통해 재학습 없이 제어 가능한 생성을 가능하게 합니다.

주요 기여

1. 새로운 패러다임: 푸시포워드 분포가 훈련 시간 동안 진화하는 프레임워크인 "드리프팅 모델"을 도입하여, 추론 시의 반복적 과정을 제거합니다.
2. 드리프팅 필드: 데이터에 대한 인력과 현재 생성물에 대한 척력을 기반으로, 생성 분포와 데이터 분포가 일치할 때 평형을 이루는 새로운 손실 메커니즘을 제안합니다.
3. 최첨단 단일 단계 생성: 단일 패스 네트워크가 다단계 확산 모델 및 플로우 기반 모델과 대등하거나 더 우수한 성능을 달-성할 수 있음을 입증합니다.
4. 모드 붕괴에 대한 강건성: 장난감 실험(toy experiments)을 통해 이 방법이 모드 붕괴에 강하다는 것을 보여줍니다. 만약 생성된 분포가 하나의 모드로 붕괴하더라도, 데이터 분포의 다른 모드들로부터 오는 인력이 계속해서 드리프트를 유도하여 분포가 진화하도록 밀어붙입니다.

실험 결과

저자들은 ImageNet 256×256의 잠재 공간(latent space)과 픽셀 공간(pixel space) 모두에서 드리프팅 모델을 평가합니다.

• 잠재 공간 (Latent Space): 표준 SD-VAE 토크나이저를 사용하여, 모델은 1-NFE FID 1.54(L/2 모델 기준)를 달성했습니다. 이는 이전의 모든 단일 단계 방법들(예: iMeanFlow, AdvFlow)보다 뛰어난 성능이며, 다단계 확산 모델(예: 250 스텝의 SiT-XL/2, FID 2.06)과도 경쟁할 만한 수준입니다.
• 픽셀 공간 (Pixel Space): 잠재 토크나이저 없이도, 모델은 1-NFE FID 1.61을 달성하여 이전의 픽셀 공간 단일 단계 방법들(예: StyleGAN-XL, FID 2.30)을 크게 앞질렀으며, 다단계 픽셀 공간 모델들과 경쟁했습니다.
• 로보틱스 제어 (Robotics Control): Diffusion Policy 설정에서, 단일 단계 드리프팅 모델은 다양한 로봇 조작 작업에서 다단계 확산 정책(100 NFE)과 대등하거나 더 나은 성능을 보여주었습니다.

의의 및 주장

본 논문은 고품질의 효율적인 생성 모델링을 위한 새로운 기회를 열었다고 주장합니다. 반복적인 신경망 훈련을 분포 진화의 메커니즘으로 재정의함으로써, 이 연구는 생성 작업을 위한 SDE/ODE 정식화의 필요성에 도전합니다.

저자들은 자신들의 방법이 다음과 같음을 강조합니다:

• 자연스럽게 **단일 단계 추론(1-NFE)**을 허용하여, 배포 시 계산 비용을 크게 줄입니다.
• 사전 학습된 다단계 모델으로부터의 증류에 의존하지 않고 처음부터 훈련합니다.
• 확산 및 플ow matching과 개념적으로 구별되는 대안을 제공하며, 추론 시의 업데이트 규칙($x_{i+1} = x_i + \Delta x_i$)에 집중합니다.

저자들은 이론적 보장에 대해서는 겸허한 태도를 유지하며, $p=q \implies V=0$은 증명되었으나, 그 역($V \to 0 \implies p \to q$)은 경험적으로 그리고 특정 비퇴화 가정 하에서 성립하지만 임의의 벡터 필드에 대해 보장되는 것은 아니라고 언급했습니다. 그들은 이 연구를 분포 진화 메커니즘의 다른 구현을 탐구하기 위한 출발점으로 보고 있습니다.