기술 요약: 중점 생성 모델 (MGM)

문제 제기

확산 (Diffusion) 및 흐름 (Flow) 모델은 이미지, 비디오, 오디오 합성 분야에서 생성형 AI 의 핵심이 되었습니다. 그러나 실제 배포를 위한 주요 병목 현상은 추론 속도입니다. 표준 확산 모델은 고품질 샘플을 생성하기 위해 수백에서 수천 개의 함수 평가 (NFE) 가 필요합니다. 더 나은 수치 솔버 개발, 사전 훈련된 모델을 한 단계 생성기로 증류 (distill) 하거나, 확산의 구조적 속성을 활용하여 처음부터 한 단계 모델을 훈련하는 등 생성 속도를 높이는 여러 접근법이 존재하지만, 여전히 사전 훈련된 "교사" 모델에 의존하지 않고 데이터로부터 직접 한 단계 생성기를 훈련하기 위한 원칙적이고 효과적인 방법이 필요합니다.

방법론

본 논문은 선형 보간을 통한 흐름 매칭 (Flow Matching) 의 특정 대칭성 속성에 기반한 한 단계 생성 모델을 훈련하기 위한 프레임워크인 **중점 생성 모델 (Midpoint Generative Models, MGM)**을 소개합니다.

핵심 통찰: 중점 대칭성

저자들은 흐름 매칭에서 두 개의 동일한 끝점 분포 ($p_0 = p_1$) 를 선형 경로 $X_t = (1-t)X_0 + tX_1$를 사용하여 보간할 때, 속도장 $v_t(x)$가 정확히 중점 시간 $t=1/2$에서 소멸함을 관찰했습니다. 구체적으로, $v_{1/2}(x) = \mathbb{E}[X_1 - X_0 | X_{1/2} = x] = 0$입니다. 반대로, $p_0 \neq p_1$인 경우 이 장은 일반적으로 0 이 아닙니다. 이 속성을 통해 중점에서의 속도장 노름을 분포 간의 불일치 측도로 사용할 수 있습니다.

중점 발산 (Midpoint Divergence)

이를 공식화하기 위해 저자들은 중점 발산을 정의합니다:
$$D_{\text{mid}}(p_0, p_1) := \mathbb{E}_{X_{1/2}} \left[ \| \mathbb{E}_{X_0, X_1}[X_1 - X_0 | X_{1/2}] \|_2^2 \right]$$
이 발산은 **결정적 (definite)**임을 증명합니다: 즉, 음수가 아니며 $p_0 = p_1$일 때만 0 이 됩니다.

중점을 넘어선 확장

이 개념을 모든 시간 단계 $t \in [0, 1]$로 단순하게 확장하는 것은 실패합니다. 내재된 방향 편향 (과거 끝점을 미래 끝점보다 재구성하기가 더 쉽기 때문) 으로 인해 $p_0 = p_1$일지라도 중점이 아닌 시간에서는 속도장이 0 이 아니기 때문입니다. 이를 해결하기 위해 MGM 은 **무작위 뒤집힌 보간 (randomly flipped interpolations)**을 도입합니다:

1. 베르누이 변수 $B \sim \text{Bernoulli}(1/2)$가 보간이 전진 ($X_t$) 또는 후진 ($X_{1-t}$) 하는지 여부를 결정합니다.
2. 관측치는 뒤집힌 변수 $\tilde{X}_t$이지만, 뒤집기 $B$는 관측되지 않습니다.
3. 이 대칭화 과정을 통해 임의의 $t$에 대해 $p_0 = p_1$이면 조건부 변위장 $\tilde{v}_t(x) = \mathbb{E}[X_1 - X_0 | \tilde{X}_t = x]$가 소멸하도록 보장합니다.

이로 인해 $t \in [0, 1/2]$에 걸쳐 이 대칭화된 변위장의 제곱 노름을 적분하여 정의된 **시간 적분 중점 발산 (Time-integrated Midpoint Divergence)**이 도출됩니다.

확률적 보간으로의 일반화

이 프레임워크는 $X_t = I_t(X_0, X_1) + \sigma_t \epsilon$ 형태의 대칭 확률적 보간으로 더 일반화됩니다. 여기서 보간 함수 $I_t$와 노이즈 스케줄 $\sigma_t$는 대칭 조건 ($I_t(x_0, x_1) = I_{1-t}(x_1, x_0)$ 및 $\sigma_t = \sigma_{1-t}$) 을 만족합니다. 이는 유효한 불일치 측도로 남는 일반화된 중점 발산을 제공합니다.

변분 훈련 목적 함수

진짜 조건부 변위장을 알 수 없기 때문에 발산을 직접 최소화하는 것은 불가능합니다. 저자들은 변분 공식화를 유도합니다:
$$D(p_\theta, p_{\text{data}}) = \max_{f} \int_0^{1/2} \mathbb{E} \left[ 2\langle f(\tilde{X}_t), X_1 - X_0 \rangle - \|f(\tilde{X}_t)\|_2^2 \right] dt$$
여기서 $f$는 학습 가능한 벡터 값 함수 (변위장 또는 비평가) 입니다. 이는 미니맥스 목적 함수를 도출합니다:

1. $f_\psi$ (비평가) 업데이트: $f_\psi$를 회귀하여 대칭화된 변위장 $X_1 - X_0$를 근사합니다.
2. $G_\theta$ (생성기) 업데이트: 고정된 $f_\psi$를 사용하여 변분 목적 함수를 최소화합니다.

생성기 $G_\theta$는 잠재 사전 분포 $z$를 $X_0$로 매핑한 후, 이를 실제 데이터 $X_1$와 짝지어 손실을 계산합니다.

주요 기여

1. 중점 발산: 선형 보간 흐름 매칭의 중점 대칭성에서 유도된 새로운 분포 불일치 도입. 저자들은 이 발산이 분포가 같을 때만 소멸하는 결정적 성질을 증명합니다.
2. 시간 뒤집기를 통한 대칭화: 무작위 시간 뒤집기를 통해 발산을 모든 시간 단계로 확장하여 방향 편향을 제거하고 시간 적분 목적 함수를 가능하게 함.
3. 일반화된 중점 발산: 개념을 대칭 확률적 보간으로 일반화하여 다양한 노이즈 스케줄을 위한 유연한 프레임워크 제공.
4. MGM 알고리즘: 한 단계 생성기를 위한 실용적인 훈련 프레임워크:
   • 사전 훈련된 확산/흐름 교사 모델 없이 데이터로부터 직접 훈련.
   • 학습된 변위장을 사용하여 발산을 근사.
   • 기존 한 단계 방법들과 경쟁력 있는 성능 달성.

실험 결과

저자들은 무조건부 CIFAR-10 (32x32) 생성에서 MGM 을 평가하여 1-NFE 샘플링에 대한 Frechet Inception Distance (FID) 를 보고합니다.

• 성능: MGM 은 2.27의 FID 를 달성하여, 교사 모델 훈련 신호를 사용하지 않는 방법들 (즉, 사전 훈련된 모델을 증류하지 않고 처음부터 훈련된 방법들) 중 가장 좋은 결과를 기록했습니다. DiffRatio-DiJS (2.39) 및 iCT-deep (2.51) 와 같은 다른 교사 없는 접근법들을 능가했습니다.
• 애블레이션 연구:
  • 시간 적분: 중점 ($t=1/2$) 만 사용하는 것은 유효하지만 노이즈가 많은 샘플을 생성합니다 (비평가 노이즈로 FID 약 13.55). 반면, 시간에 걸쳐 적분하면 품질이 크게 향상됩니다 (FID 2.27).
  • 무작위 뒤집기: 단순한 뒤집지 않은 시간 적분 목적 함수는 방향 편향으로 인해 실패합니다 (FID > 100), 이는 무작위 뒤집기의 필요성을 확인시켜 줍니다.
  • 확률성: 보간체에 작은 양의 노이즈 ($\sigma > 0$) 가 안정적인 훈련에 필수적입니다. 결정적 브리지는 성능이 저조합니다.
  • 비평가 확률성: 생성기와 비평가의 노이즈 수준을 분리 (비평가에만 노이즈 사용) 하는 것이 최상의 결과를 냅니다.

중요성 및 주장

본 논문은 MGM 이 생성 모델링에 대한 효과적이고 이론적으로 근거 있는 접근법을 제공한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 이론적 기반: 대칭 확률적 보간을 단순한 생성 경로가 아닌, 결정적 발산을 통해 끝점 분포를 비교하는 기반으로 재해석함.
• 훈련 패러다임: 사전 훈련된 모델을 증류하는 것의 복잡성과 잠재적 한계를 피하고, 데이터로부터 직접 한 단계 생성기를 훈련할 수 있는 원칙적인 방법 제공.
• 경쟁력: 이 특정 대칭성에 기반한 방법이 교사 없는 한 단계 생성기 중 최첨단 결과를 달성할 수 있음을 보여줌으로써, 고품질 고속 생성을 위해 증류가 필수적이라는 통념에 도전함.

저자들은 한계점을 인정하며, 미니맥스 훈련이 GAN 과 유사한 최적화 불안정성을 초래할 수 있음을 지적하고, 성능은 노이즈 스케줄 및 보간체 유형과 같은 설계 선택에 의존한다고 언급합니다. 그들은 모든 교사 기반 증류 방법 (예: CTM 은 1.87 FID 달성) 을 능가한다고 주장하지는 않지만, 처음부터 훈련하는 데 있어 MGM 이 선도적인 대안임을 위치시킵니다.