기술 요약: 자신감 있게 흐름에 따르기 (FMwC)

문제 제기

생성 모델, 특히 흐름 매칭 (Flow Matching, FM) 모델은 언어, 비전, 과학 분야에서 최첨단 성능을 달성했습니다. 그러나 이들은 (할루시네이션된 인용문, 해부학적으로 부정확한 도표, 또는 열역학적으로 불안정한 결정체와 같은) 출력물을 생성할 때 신뢰도에 대한 내재적 표시를 제공하지 않습니다. 딥 앙상블 (Deep Ensembles) 이나 MC-Dropout 과 같은 기존 불확실성 정량화 솔루션은 각 샘플당 $k$개의 별도 모델이나 $k$개의 확률적 궤적을 실행하는 데 의존합니다. 이 접근법은 $k\times$의 계산 비용을 초래하여 대규모 배포 시 샘플별 신뢰도 추정을 과도하게 비싸게 만듭니다. 또한, 가능도 기반 방법들은 심층 생성 모델에서 분포 내 입력과 분포 외 입력을 구별하는 데 신뢰할 수 없는 것으로 입증되었습니다. 다루어지는 핵심 질문은 흐름 매칭 모델이 추론 비용을 증가시키거나 가능도 접근을 요구하지 않고도 생성된 출력물에 대한 자체 불확실성을 정량화할 수 있는지 여부입니다.

방법론: 자신감 있는 흐름 매칭 (FMwC)

FMwC 는 표준 결정론적 흐름 매칭과 동일한 계산 비용으로 흐름 매칭 모델에 샘플별, 궤적 수준의 신뢰도 점수를 부여하는 단일 통과 (single-pass) 접근법입니다. 이 방법은 변분 적응형 드롭아웃 (Variational Adaptive Dropout, VAD) 을 상미분 방정식 (ODE) 궤적을 통한 폐형식 모멘트 전파와 통합합니다.

1. 변분 공식화

FMwC 는 단일 결정론적 속도장 $v_{t,\theta}(x_t)$를 학습하는 대신, 네트워크 가중치 $\omega$를 변분 사후분포 $q_\psi(\omega)$에서 추출된 확률 변수로 취급합니다. 모델은 증거 하한 (Evidence Lower Bound, ELBO) 을 최대화합니다:
$$\text{ELBO}(\psi) = \mathbb{E}_{\omega \sim q_\psi} [\log \mathcal{L}(D | \omega)] - \text{KL}[q_\psi(\omega) \parallel \pi(\omega)]$$
여기서 가능도 항은 표준 조건부 흐름 매칭 회귀 손실에 해당합니다. 중요하게도, 사후분포 하의 기대 로그 가능도의 기울기는 가중치 샘플에서 평가된 표준 회귀 손실의 기울기와 일치하므로, 시뮬레이션 없는 훈련이 가능합니다.

2. VAD 를 통한 상태 의존적 사후분포

불확실성이 흐름의 국소 기하학에 민감하도록 보장하기 위해, FMwC 는 변분 적응형 드롭아웃 (VAD) 을 사용합니다. 이 프레임워크에서 층 $l$의 가중치에 대한 사후분포는 가우시안 $q_\psi(\omega_l | x, t) = \mathcal{N}(\theta_l, \alpha_l(x, t)^2 \theta_l^2)$이며, 여기서 드롭아웃 스케일 $\alpha_l(x, t)$은 작은 추론 네트워크 $E_\gamma$에 의해 예측됩니다. 이는 사후분포의 확산이 기하학적 모호성 영역 (속도장을 맞추기 어려운 영역) 에서는 넓어지고 다른 곳에서는 좁아지도록 하여, 오직 ELBO 를 통해 학습되도록 합니다.

3. 분석적 분산 전파

FMwC 의 핵심 혁신은 신경망 레이어와 ODE 통합 궤적을 모두 통해 속도장의 분산을 단일 순방향 통과로 분석적으로 전파하는 것입니다:

• 네트워크를 통한 전파: 로컬 재파라미터화 트릭을 사용하여 활성화의 평균 ($\mu_l$) 과 분산 ($\sigma_l^2$) 을 레이어별로 전파합니다. 선형 레이어의 경우 분산 재귀는 폐형식입니다. 비선형성의 경우, 모멘트는 몬테카를로 샘플링의 필요성을 피하기 위해 가우스 - 헤르미트 구적법이나 테일러 선형화를 사용하여 분석적으로 계산됩니다.
• ODE 를 통한 전파: 각 오일러 단계에서 네트워크는 사후평균 속도 $\mu_t$와 그 분산 $\sigma_t^2$를 출력합니다. 적분기는 $\mu_t$를 사용하여 위치를 진행시키는 반면, $\sigma_t^2$는 기록됩니다. 이는 단일 스칼라가 아닌 $\{\sigma_t^2\}_{t=0}^T$의 분산 궤적을 생성합니다.

4. 신뢰도 판독

분산 궤적은 하류 작업을 위한 스칼라 신뢰도 점수로 집계될 수 있습니다:

• 종점 신뢰도: 최종 단계 ($t=1$) 의 분산.
• 적분 신뢰도: 궤적을 따라 속도 표준편차의 합.
• 시간적 비율: 수렴하는 샘플 (낮은 후기 분산) 과 모호하게 남는 샘플을 구별하는 데 도움이 되는 후기 창 분산과 초기 창 분산의 비율.
• 학습된 헤드: 전체 분산 시퀀스를 작업별 점수로 매핑하는 작은 지도 학습 헤드.

주요 기여

1. 단일 통과 신뢰도 추정기: FMwC 는 단일 결정론적 순방향 통과 비용으로 흐름 매칭에 대한 보정된 샘플별 신뢰도 점수를 제공하는 최초의 방법이며, 표준 FM 의 추론 예산과 일치합니다. 이는 학습된 입력 의존적 노이즈를 ODE 를 통해 분석적으로 전파함으로써 이를 달성합니다.
2. 훈련 없는 추론 개선: 신뢰도 점수는 재훈련 없이 실용적인 응용을 가능하게 합니다:
   • 필터링: 낮은 신뢰도 출력을 폐기함으로써 샘플 품질을 단조롭게 향상시킵니다.
   • 편집: 모델이 구조에 결속하는 정확한 시간 단계 ($t^*$, 분산 피크) 를 식별하여 $t^*$에서 노이즈를 주입함으로써 표적화되고 제약 조건을 준수하는 편집 (예: 결정체 내 화학적 교체) 을 가능하게 합니다.
   • 적응적 스텝핑: 높은 분산 (모호성) 영역에 ODE 계산 단계를 집중시켜 균일한 스텝핑이 실패하는 낮은 단계 예산에서도 고품질 샘플을 복원합니다.
3. 기하학적 해석 가능성: 이 논문은 FMwC 신뢰도 점수가 학습된 속도장의 국소 발산과 상관관계가 있음을 입증합니다. 이는 생성 과정에 대한 창을 제공하여 인식론적 불확실성을 흐름의 기하학적 불안정성 (분기점) 과 연결합니다.

실험 결과

저자들은 FMwC 를 2D 밀도 추정 (Checkerboard), 이미지 생성 (MNIST), 그리고 새로운 무기 결정체 생성 (Riemannian Flow Matching) 의 세 가지 모달리티에 걸쳐 평가했습니다.

• 샘플 품질: FMwC 는 표준 FM 에 비해 샘플 품질을 유지하거나 향상시킵니다. Checkerboard 와 MNIST 에서는 각각 잘못된 배치율과 FID 점수가 낮아졌으며, 결정체에서는 동등한 품질을 유지했습니다.
• 추론 비용: FMwC 는 더 큰 백본 (예: MNIST 의 경우 약 1.003$\times$, 결정체의 경우 약 1.08$\times$) 에서 결정론적 기준선에 비해 무시할 만한 오버헤드를 발생시킵니다. 반면, 앙상블 및 MC-Dropout 기준선은 $5\times$의 계산이 필요합니다.
• 필터링 성능: FMwC 점수에 기반한 필터링은 $k=5$ 앙상블 및 MC-Dropout 과 정밀도 - 재현율 곡선 아래 면적 (AUPRC) 점수가 비슷하지만 단일 궤적 비용으로 달성했습니다. 예를 들어, Checkerboard 데이터셋에서 시간적 비율 점수는 0.55 의 AUPRC 를 달성하여 1/5 의 계산 비용으로 앙상블 기준선 (0.55) 과 일치했습니다.
• 편집 성공률: 분산 피크 ($t^*$) 에서 노이즈를 주입하여 결정체를 편집할 때, FMwC 는 무작위 시간 단계 선택에 비해 훨씬 높은 성공률 (예: 화학적 교체의 경우 39% 대 22%) 을 달성하여 고분산 영역이 중요한 결정 지점에 해당함을 입증했습니다.
• 발산과의 상관관계: 분산 궤적 $\sigma_t^2$는 속도장의 발산 크기 ($|\nabla \cdot v_{t,\theta}|$) 와 강한 상관관계를 보였습니다. 결정체 데이터셋에서 적분 상관관계는 스피어만 $\rho$가 0.96 에 달하여 신뢰도 점수의 기하학적 해석을 검증했습니다.

중요성과 주장

이 논문은 FMwC 가 샘플 생성기를 신뢰성 측정치와 짝지어진 샘플 생성기로 변환한다고 주장합니다. 가중치에 대한 입력 의존적 사후분포를 학습하고 분산을 분석적으로 전파함으로써, FMwC 는 과학 및 안전-중요 응용 분야에서 신뢰도 인식 생성을 실용적인 기본값으로 만듭니다. 저자들은 신뢰도 점수가 앙상블에 의해 측정되는 모델 간 불일치와 구별되는 단일 훈련된 흐름의 고유한 기하학적 속성임을 강조합니다. 이는 기존 불확실성 정량화 방법의 과도한 비용 없이 생성 과정의 "외과적" 안내를 가능하게 하여 불안정성의 순간을 표적화합니다. 이 연구는 학습된 속도장의 기하학을 이해하고 활용하는 것이 더 신뢰할 수 있고, 해석 가능하며, 제어 가능한 생성 모델로 이어질 수 있음을 시사합니다.