기술 요약: 임의의 보조 경로를 활용한 흐름 매칭 (AuxPath-FM)

문제 제기

스코어 기반 확산 모델 및 흐름 매칭 (FM) 프레임워크를 포함한 연속 시간 생성 모델은 복잡한 데이터 분포를 모델링하는 강력한 도구로 자리 잡았습니다. 조건부 생성 모델링에서 측면 정보 (예: 클래스 레이블) 를 통합하는 것은 제어 가능한 합성을 위해 필수적입니다. 그러나 기존 접근법들은 주로 신경망의 조건부 설정을 통해 이러한 정보를 도입할 뿐, 근본적인 확률 경로는 변경하지 않습니다. 최근 이론적 연구는 흐름 매칭이 고유한 다양체 구조에 적응함을 시사하지만, 경로의 확률적 성분이 궤적 수준의 제어에 미치는 영향은 아직 충분히 탐구되지 않았습니다.

또한, 확률적 보간자 (stochastic interpolants) 및 슈뢰딩거 브리지 (Schrödinger Bridge) 방법과 같은 기존 경로 기반 공식들은 일반적으로 확률적 보간을 위해 가우시안 노이즈나 브라운 역학에 의존합니다. 이러한 제한은 구조화된 안내를 위한 확률 경로 설계의 유연성을 제한합니다. 최근 연구들이 비가우시안 교란 또는 이산 상태 공간 흐름을 탐구하기는 했으나, 이러한 방법들은 종종 특정 데이터 유형에 국한되거나 단일 일반 프레임워크 하에서 연속 및 이산 사전 분포를 통합하지 못합니다. 따라서 네트워크 수준의 변조에만 의존하는 것이 아니라 생성 궤적이 다양한 보조 신호에 의해 직접적으로 형성될 수 있는 원칙적인 방법이 필요합니다.

방법론: AuxPath-FM

저자들은 **임의의 보조 경로를 활용한 흐름 매칭 (AuxPath-FM)**을 소개합니다. 이는 확률 경로에 임의의 분포에서 추출된 보조 확률 변수 $\eta$를 직접 통합함으로써 조건부 흐름 매칭을 일반화한 프레임워크입니다.

1. 임의의 보조 확률 경로

보조 항을 가우시안 노이즈로 제한하는 표준 보간 방식과 달리, AuxPath-FM 은 다음과 같은 형태의 궤적을 구성합니다:
$$X_t = a(t)X_1 + b(t)X_0 + c(t)\eta$$
여기서:

• $X_0 \sim p_0$ (기본 분포) 및 $X_1 \sim p_1$ (목표 데이터 분포).
• $\eta \sim p_\eta$는 임의의 분포 (예: 가우시안, 균일, 라플라스, 라데마허, 또는 이산 분포) 에서 추출된 보조 변수입니다.
• $a(t), b(t), c(t)$는 $p_0$에서 $p_1$로 샘플을 운반하도록 보장하는 경계 조건 ($a(0)=0, b(0)=1, c(0)=0$ 및 $a(1)=1, b(1)=0, c(1)=0$) 을 만족하는 시간 의존 스칼라 함수입니다.

2. 이론적 일관성

본 논문은 이 구성이 연속 방정식을 보존함을 입증합니다. 정리 1은 조건부 벡터장 $u_t(x | x_0, \eta)$이 연속 방정식을 만족할 경우, 조건부 장들을 집계하여 얻은 주변 벡터장 $u_t(x)$가 주변 확률 경로 $p_t(x)$를 생성함을 보여줍니다. 이는 보조 변수를 통해 복잡한 경로를 설계할 수 있는 원칙적인 근거를 제공합니다.

3. 훈련 목적 함수 및 동등성

$(X_0, \eta)$에 대한 적분으로 인해 주변 흐름 매칭 목적 함수를 직접 최적화하는 것은 계산적으로 불가능합니다. 대신 저자들은 계산 가능한 경로에 정의된 조건부 목적 함수를 최적화합니다:
$$L_{\text{AuxPath}} = \mathbb{E}_{t, X_0, X_1, \eta} \left[ \| v_\theta(X_t, t) - \dot{X}_t \|^2 \right]$$
여기서 $\dot{X}_t = \dot{a}(t)X_1 + \dot{b}(t)X_0 + \dot{c}(t)\eta$입니다. 정리 2는 이 조건부 목적 함수가 네트워크 매개변수 $\theta$와 무관한 상수까지 주변 흐름 매칭 목적 함수와 동등함을 증명합니다. 따라서 조건부 손실을 최적화하는 것은 계산 불가능한 주변 목적 함수와 동일한 최적 해를 산출합니다.

4. 의미적 보조 변수 및 안내

조건부 생성을 가능하게 하기 위해, 프레임워크는 $\eta = F_\phi(Y)$를 정의합니다. 여기서 $Y$는 클래스 레이블이며, $F_\phi$는 클래스별 중심 ($F_\phi(Y) \approx \mathbb{E}[X_1 | Y]$) 을 근사하도록 프로토타입 학습을 통해 훈련된 경량 신경망입니다.

이 공식은 **궤적 수준 무조건부 안내 (CFG)**를 가능하게 합니다. 조건부 및 무조건부 속도장은 다음과 같이 표현됩니다:
$$v(x, t, y) = v_\theta(x, t) + \dot{c}(t)F_\phi(y)$$
$$v(x, t, \emptyset) = v_\theta(x, t) + \dot{c}(t)F_\phi(\emptyset)$$
표준 CFG 를 적용하면 전역 드리프트로 작용하는 안내 항이 생성됩니다. 중요한 점은 이 접근법이 안내를 백본 모델에서 분리한다는 것입니다. 기존 CFG 가 백본 $v_\theta$의 두 번의 평가를 요구하는 반면, AuxPath-FM 은 $v_\theta$의 단 한 번의 평가만 요구하며 경량 모듈 $F_\phi$를 통해 안내를 적용합니다.

주요 기여

1. 확률 경로의 일반화: 본 논문은 확률 경로의 보조 성분이 임의의 분포를 따를 수 있도록 흐름 매칭을 일반화하여, 확률적 보간자 및 슈뢰딩거 브리지 방법에서 발견된 가우시안 노이즈 제한을 넘어섭니다.
2. 이론적 기반: 주변 흐름이 연속 방정식과 일관되며 조건부 훈련 목적 함수가 주변 목적 함수와 동등함을 엄밀하게 이론적으로 증명합니다.
3. 궤적 수준 제어: 이 프레임워크는 네트워크 조건부 설정에만 의존하는 것이 아니라 구조화된 의미 정보 (예: 클래스 레이블) 를 생성 궤적 자체에 직접 인코딩할 수 있게 합니다.
4. 효율적인 무조건부 안내: 저자들은 백본 네트워크의 이중 평가를 피함으로써 추론 계산 비용을 약 절반으로 줄이는 궤적 수준의 CFG 공식을 제안합니다.
5. 통합적 관점: 이 방법은 이산 및 두꺼운 꼬리 (heavy-tailed) 사전 분포를 포함한 다양한 보간 방식을 단일 프레임워크 하에서 통합하여, 구조화된 연속 시간 생성 모델링을 위한 유연한 관점을 제공합니다.

실험 결과

저자들은 MNIST, CIFAR-10, ImageNet-1k 데이터셋에서 AuxPath-FM 을 평가했습니다.

• 조건부 생성 성능: MNIST 및 CIFAR-10 에서 표준 조건부 흐름 매칭 (CFM) 에 학습된 의미 프로토타입을 사용한 보조 경로 안내를 결합하면, 기준 CFM 이나 보조 안내 단독에 비해 의미 정확도 (분류 정확도) 와 생성 품질 (FID, sFID) 이 일관되게 향상되었습니다.
• 분포 선택에 대한 강건성: Ring-64 데이터셋 및 ImageNet-1k 에 대한 실험은 가우시안, 균일, 라플라스, 라데마허 등 다양한 보조 분포에 걸쳐 프레임워크가 안정적인 생성 성능을 유지함을 보여주었습니다. 레이블 안내 보조 변수가 가장 좋은 성능을 달성하여 개선된 제어 가능한 다중 모드 생성을 나타냈습니다.
• 안내 스케일 효과: 궤적 수준 CFG 에서 안내 스케일 $w$를 증가시키면 경쟁력 있는 생성 품질을 유지하면서 의미 정렬이 개선되었습니다.
• 효율성: 궤적 수준 CFG 공식은 안내의 효과를 유지하면서 표준 CFG 대비 추론 비용 (GFLOPs 단위) 을 약 50% 감소시키는 데 성공했습니다.
• 모델 적응: 이 프레임워크는 추가 비용과 아키텍처 변경 없이 사전 훈련된 무조건부 흐름 모델을 조건부 생성기로 성공적으로 변환했습니다.

중요성 및 주장

본 논문은 AuxPath-FM 이 확률 경로 설계에 대한 원칙적이고 일반적인 기반을 제공한다고 주장합니다. 보조 변수가 임의의 분포를 따를 수 있도록 함으로써, 이 프레임워크는 이론적 일반성과 실용적 유연성을 모두 제공합니다.

저자들은 그들의 방법이 조건부 생성의 패러다임을 순수한 네트워크 수준 변조에서 궤적 수준 설계로 전환한다고 강조합니다. 이는 안내 신호를 생성 역학에 직접 임베딩할 수 있게 합니다. 이 연구의 중요성은 다음과 같은 능력에 있습니다:

• 기존 보간 방식을 특수한 경우로 통합.
• 다양한 사전 분포를 통해 생성 흐름에 다양한 기하학적 속성 부여.
• 무조건부 안내를 위한 계산적으로 효율적인 메커니즘 제공.
• 기존 무조건부 모델을 조건부 모델로 업그레이드하는 간단하면서도 효과적인 방법 제공.

본 논문은 이 접근법이 구조화된 연속 시간 생성 모델링을 위한 통합적 관점을 제공하며, 의미 보조 CFG 가 우수한 생성 품질을 유지하면서 계산 효율성을 크게 향상시킨다는 광범위한 실험을 통해 검증되었다고 결론지었습니다.