기술 요약: 플럭스 매칭을 활용한 생성 모델링

문제 제기

이미지 합성, 단백질 설계, 로봇 공학과 같은 고차원 영역에서 현재 최첨단 생성 모델링은 주로 **스코어 기반 모델 (score-based models)**에 크게 의존합니다. 이러한 모델들은 일반적으로 **스코어 매칭 (score matching)**을 통해 Stein 스코어 함수, $\nabla \log p_{\text{data}}(x)$를 학습합니다. 일단 학습되면, 이 스코어는 타겟 분포 $p_{\text{data}}$로부터 샘플링하기 위해 확산 과정 (예: Langevin 역학) 을 구동하는 데 사용됩니다.

그러나 스코어 함수는 $p_{\text{data}}$를 생성할 수 있는 유일한 벡터 장이 아닙니다. **Fokker-Planck 방정식 (FPE)**에 따르면, 조건 $\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$을 만족하는 임의의 벡터 장 $f_\theta$는 $p_{\text{data}}$를 정상 분포로 갖게 됩니다. 이는 동일한 타겟 분포를 생성하지만 서로 다른 역학 (예: 다른 혼합 속도, 순환 패턴, 또는 지향적 의존성) 을 유도할 수 있는 무한한 가족의 생성 벡터 장 (generative vector fields) (비보존 장) 이 존재함을 의미합니다.

기존의 스코어 매칭 목적 함수는 모델이 이 공간의 단일 점 ($\nabla \log p_{\text{data}}$) 으로 수렴하도록 강제하며, 그 편차가 타겟 분포를 보존하더라도 그 편차를 처벌합니다. 이는 유용한 귀납적 편향, 구조적 사전 지식, 또는 더 빠른 혼합이나 인과적 의존성과 같은 동적 속성을 생성 장에 직접 인코딩하는 능력을 제한합니다.

방법론: 플럭스 매칭

저자들은 **플럭스 매칭 (Flux Matching)**을 제안합니다. 이는 모델이 점별적으로 스코어와 일치하도록 강제하는 대신, 타겟 분포를 생성하는 장의 가족 내의 임의의 벡터 장을 학습함으로써 스코어 매칭을 일반화하는 새로운 패러다임입니다.

핵심 목적 함수

모델 $f_\theta$와 스코어 $\nabla \log p_{\text{data}}$ 사이의 Fisher 발산을 최소화하는 대신, 플럭스 매칭은 **확률 플럭스 (probability flux)**의 발산을 최소화합니다. 장이 생성적이기 위한 조건은 다음과 같습니다:
$$\nabla \cdot (p_{\text{data}} f_\theta) = \nabla \cdot (p_{\text{data}} \nabla \log p_{\text{data}})$$
이 조건은 점별 등식보다 약하며, 무한히 많은 유효한 벡터 장을 포함하는 "영공간 (nullspace)"을 허용합니다.

이론적 유도

이 조건을 최적화하기 위해, 저자들은 **투사된 Fisher 발산 (Projected Fisher Divergence)**을 정의합니다:
$$\tilde{J}(\theta) := \mathbb{E}_{x \sim p_{\text{data}}} \left[ \| \Pi_{\text{flux}} f_\theta(x) - \nabla \log p_{\text{data}}(x) \|^2 \right]$$
여기서 $\Pi_{\text{flux}}$는 벡터 장을 동일한 플럭스 발산을 공유하는 기울기 장의 공간으로 투사합니다. $\tilde{J}(\theta)$는 이론적으로 타당하지만 고차원에서는 계산이 불가능합니다.

저자들은 다음과 같이 계산 가능한 플럭스 매칭 손실 (Flux Matching Loss) ($L_{\text{flux}}$) 을 유도합니다:

1. 잔여 장 $u_\theta = f_\theta - \nabla \log p_{\text{data}}$와 Langevin-Stein 연산자 $r_\theta = \nabla \cdot u_\theta + u_\theta \cdot \nabla \log p_{\text{data}}$를 정의합니다.
2. 확산 시뮬레이션 $d x_t = \nabla \log p_{\text{data}}(x_t) dt + \sqrt{2} dW_t$를 사용하여 $r_\theta$를 데이터 분포 전체에 전파합니다.
3. 부분적분과 stop-gradient 트릭을 적용하여 목적 함수를 값별 미분을 통해 역전파하는 것을 피하는 형태로 변환합니다.

최종 손실은 다음과 같습니다:
$$L_{\text{flux}}(\theta) := -\mathbb{E}_{t \sim q, x_0 \sim p_{\text{data}}, x_t | x_0} \left[ \frac{1}{q(t)} u_\theta(x_0)^\top \text{sg}\left( \frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t) \right) \right]$$
여기서 $\text{sg}$는 stop-gradient 연산자를 나타냅니다. 저자들은 $\nabla_\theta \tilde{J}(\theta) = 2 \nabla_\theta L_{\text{flux}}(\theta)$임을 증명합니다.

노이즈 어닐링 및 실제 구현

이 방법은 노이즈 수준 $\sigma$에 대한 장의 연속체를 학습함으로써 노이즈 어닐링 (noise-annealed) 설정 (확산 모델과 유사) 으로 확장됩니다.

• 스코어 근사: $\nabla \log p_{\text{data}}$는 알 수 없으므로, 저자들은 미니배치에 대한 비모수 커널 밀도 추정 (KDE) 근사를 사용하여 스코어를 추정합니다.
• 분산 감소: 경로별 항 $\frac{\partial x_t}{\partial x_0}^\top \nabla_{x_t} r_\theta(x_t)$를 추정하기 위해, 저자들은 단일 체인에 의존하는 대신 가우시안 전이 확률로 가중치를 둔 미니배치의 모든 샘플을 사용하는 교차 체인 추정기를 사용합니다.

특정 장 선택

플럭스 매칭은 장의 가족을 학습합니다. 원하는 속성을 가진 특정 장을 선택하기 위해 저자들은 두 가지 전략을 제안합니다:

1. 애플리케이션별 손실: 플럭스 매칭 목적 함수에 혼합 속도나 구조적 제약과 같은 보조 손실을 추가합니다.
2. 모델 파라미터화: $f_\theta$에 직접 특정 구조를 강제하는 아키텍처 제약 (예: 인과적 어텐션 마스크) 을 적용합니다.

주요 기여

1. 플럭스 매칭 패러다임: 확률 플럭스의 발산을 매칭함으로써 스코어 이상의 벡터 장을 학습하는 생성 모델링 프레임워크를 도입했습니다.
2. 효율적인 손실 유도: 비보존 벡터 장을 허용하면서도 Fisher 발산 기하학을 보존하는 확장 가능한 플럭스 매칭 손실을 유도했습니다.
3. 노이즈 어닐링 확장: 복잡한 고차원 이미지 데이터셋에 적용할 수 있도록 노이즈 어닐링 분포로 방법을 확장했습니다.
4. 검증된 유연성: 플럭스 매칭이 표준 스코어 매칭 하에서는 불가능한 애플리케이션을 가능하게 함을 보여주는 실증적 증거:
   • 더 빠른 혼합: 샘플링 수렴을 가속화하는 벡터 장 학습.
   • 해석 가능성: 블랙박스 신경 근사 없이 구조화된 생물학적 의미 모델 (예: RNA velocity) 적합.
   • 구조화된 역학: 지향적 의존성 (예: 시간적 자기회귀) 을 생성 장에 직접 임베딩.

결과

이 논문은 다섯 가지 설정에서 플럭스 매칭을 평가합니다:

1. 제어 가능한 장 (Toy Data): 2D 가우시안 혼합 모델에서 플럭스 매칭은 타겟 분포를 보존하지만 역학 (예: 회전, 순환) 이 다른 장에 대해 0 손실을 부여하는 반면, 스코어 매칭은 이러한 편차를 처벌합니다.
2. 해석 가능한 장 (RNA Velocity): 단일 세포 유전체학에 적용하여 플럭스 매칭은 경사 하강법을 통해 scVelo 동적 모델 파라미터를 적합합니다. 다섯 개의 데이터셋에서 일관성과 경계 간 정확도 측면에서 표준 EM 기반 scVelo 적합보다 우월하며, 이 이득이 모델 표현력 증가가 아닌 최적화 절차에서 비롯됨을 보여줍니다.
3. 제한 없는 생성 (이미지): CIFAR-10 및 CelebA 에서 플럭스 매칭은 독립적인 목적 함수로 작용합니다. 경쟁력 있는 FID 점수를 달성하지만, 현재 Denoising Score Matching (DSM) 에 비해 훈련 중 34 배 느리고 23 배 더 많은 메모리를 사용합니다. 그러나 고차원 분포로 확장 가능합니다.
4. 빠른 혼합: 혼합 프록시를 최적화함으로써 플럭스 매칭은 기존 플럭스 매칭이나 DSM 보다 훨씬 적은 샘플링 단계로 합리적인 FID 점수에 도달하는 장을 학습합니다.
5. 구조 임베딩: 궤적 생성 작업 (비선형 스프링) 에서 플럭스 매칭은 시간적 자기회귀를 강제하기 위해 인과적 어텐션 마스크를 성공적으로 통합합니다. 반면, 동일한 마스크를 DSM 에 적용하면 스코어 함수의 대칭적 야코비안이 지향적 의존성과 충돌하여 성능이 저하됩니다.

중요성 및 주장

저자들은 플럭스 매칭이 벡터 장 자체를 고정된 목표가 아닌 디자인 선택으로 취급함으로써 생성 모델링에 "새로운 차원"을 열었다고 주장합니다.

• 샘플링을 넘어: 주요 가치는 표준 이미지 생성 (여기서 DSM 이 매우 최적화됨) 을 DSM 으로 대체하는 것이 아니라 유용한 속성을 가진 역학을 가능하게 하는 데 있습니다.
• 유연성: 이는 실무자가 귀납적 편향, 구조적 사전 지식, 그리고 메커니즘적 제약을 생성 과정에 직접 부과할 수 있게 합니다.
• 애플리케이션: 이 논문은 이 유연성으로 인해 가능해진 구체적인 사용 사례를 강조합니다: 비가역적 역학을 통한 더 빠른 샘플링, RNA velocity 와 같은 해석 가능한 과학적 모델, 그리고 지정된 인과적 또는 시간적 구조를 가진 생성 모델.

이 논문은 플럭스 매칭이 현재 더 높은 계산 비용을 수반하지만, 비보존적이고 구조화된 생성 장을 학습할 수 있는 능력은 출력의 분포만큼이나 생성의 역학이 중요한 애플리케이션에 강력한 도구를 제공한다고 결론지었습니다.