기술 요약: Beckmann 수송 모델 (Beckmann Transport Models)

문제 정의

Flow matching 및 확산 모델(diffusion models)을 통한 표준 생성 모델링은 기저 분포 $\mu_0$에서 타겟 분포 $\mu_1$로 샘플을 수송하기 위해 시간 의존적 속도장(drift)에 의존한다. 이는 효과적이지만, 추론 시 고정된 구간 $[0, 1]$ 동안 시간 의존적 상미분 방정식(ODE)을 적분해야 하므로 계산 비용이 많이 든다.

Wang과 Du [2025]가 "평형 매칭(Equilibrium Matching, EM)"으로 제안한 대안적 접근 방식은 드리프트를 시간에 독립적인(autonomous) 형태로 제한한다. 이는 아키텍처와 동역학계를 단순화한다. 그러나 기존의 EM 정식화는 시간 독립적 드리프트를 가진 자율 흐름(autonomous flow)이 $\mu_0$를 $\mu_1$로 정확하게 수송할 수 있다는 엄밀한 증명이 부족했다. 또한, 원래 EM 논문에서 제안된 손실 함수는 수송 과정에서 편향을 유발하여, 특히 타겟 분포가 저차원 매니폴드 위에 지지(support)되어 있을 때 정확한 질량 할당(mass allocation)을 수행하지 못하는 것으로 나타났다.

방법론

본 논문은 시간 독립적 속도장 $b: \mathbb{R}^d \to \mathbb{R}^d$를 사용하여 flow matching을 구현하는 프레임워크인 **Beckmann 수송 모델(BTM)**을 도입한다. 핵심적인 이론적 통찰은 타겟 분포 $\mu_1$이 특이(singular)할 경우(즉, 저차원 매니폴드 $M_1$ 위에 지지될 경우), 시간 독립적 드리프트가 유효한 수송 맵을 정의할 수 있다는 것이다.

1. 자율 흐름 및 발산 조건 (Autonomous Flow and Divergence Condition)

저자들은 양의 가중치 함수 $\nu$를 포함하는 특정 발산 방정식을 만족하는 경우, 시간 독립적 드리프트 $b$가 유효한 수송을 정의함을 입증한다:
$$\nabla \cdot (\nu b) = \mu_0 - \mu_1$$
여기서 $\mu_0$는 소스(source) 역할을 하고 $\mu_1$은 싱크(sink) 역할을 한다. 가중치 $\nu$는 표준 flow matching 설정에서의 시간 평균 보간 분포(time-averaged interpolant distribution)를 나타낸다.

• 수렴성: 이 조건 하에서, $\mu_0$로부터 시작하는 자율 ODE $\dot{X}_t = b(X_t)$의 궤적은 $\mu_1$의 지지 집합 $M_1$으로 수렴한다.
• 수송 맵: 히팅 타임(hitting time) $\tau(x_0)$ ($M_1$에서의 드리프트 거동에 따라 유한하거나 무한할 수 있음)은 $\mu_0$를 $\mu_1$로 밀어내는(push forward) 맵 $T(x_0) = \lim_{t \to \tau} X_t(x_0)$를 정의한다 ($T_\sharp \mu_0 = \mu_1$).

2. 평형 매칭의 교정 (Correction of Equilibrium Matching)

본 논문은 기존의 평형 매칭 손실이 정확한 수송에 필요한 발산 조건을 엄격히 만족하지 못하는 회귀 목적 함수를 최소화한다는 점을 지적한다. 보간물(interpolant)의 스케줄을 타겟에 대한 별도의 스케일링 인자가 아닌 보간물 자체의 일부로 취급함으로써, 저자들은 학습된 드리프트가 발산 조건을 반드시 만족하도록 보장하는 일관된 손실(consistent loss) (식 5)을 도출한다. 이 교정은 원래 EM 방식에서 관찰되었던 질량 할당 편향을 제거한다.

3. 보존 방정식을 통한 원스텝 맵 (One-Step Map via Conservation Equation)

핵심 기여 중 하나는 수송 맵 $T$를 ODE 적분의 결과가 아니라, 흐름 궤적을 따르는 **보존 방정식(conservation equation)**의 유일한 해로 규정하는 것이다:
$$b(x) \cdot \nabla T(x) = 0 \quad (x \notin M_1)$$
단, 경계 조건은 $x \in M_1$일 때 $T(x) = x$이다.
이 방정식은 맵 $T$의 값이 $b$의 흐름선을 따라 일정함을 의미한다. 이러한 구조적 통찰은 ODE를 풀지 않고도 $T$를 직접 학습할 수 있게 한다.

4. 직접 학습 목적 함수 (Direct Learning Objective)

샘플로부터 $T$를 직접 학습하기 위해, 저자들은 흐름 맵을 지배하는 오일러 방정식(Eulerian equation)을 이산화한다. 전방 오일러 단계(forward-Euler step)를 사용하여, 네트워크 $T_\theta$가 흐름 시퀀스의 다음 상태를 예측하도록 훈련하는 회귀 목적 함수(식 15)를 도출한다:
$$\mathcal{L}(T) = \mathbb{E}_{t, x_0, x_1} \left[ \left| T(I_t) - \text{sg}\left( T(I_t) + \dot{I}_t \cdot \nabla T(I_t) \right) \right|^2 \right] + \lambda \mathbb{E}_{x_1} \left[ |T(x_1) - x_1|^2 \right]$$
여기서 $\text{sg}$는 stop-gradient 연산자를 나타낸다. 이를 통해 원스텝 생성($x_1 = T(x_0)$) 또는 반복적 정교화($x_1 = T^k(x_0)$)를 수치적 ODE 적분 없이 수행할 수 있다.

주요 기여

1. 이론적 정당화: 자율 흐름의 드리프트가 특이 타겟(singular targets)으로의 유효한 수송을 정의함을 증명하여, 손실 교정을 통해 Equilibrium Matching을 복구한다.
2. 일반적 프레임워크: 발산 방정식 $\nabla \cdot (\nu b) = \mu_0 - \mu_1$을 만족하는 임의의 쌍 $(\nu, b)$가 유효한 자율 수송을 생성함을 보여준다. 이는 Flow Matching(여기서 $\nu$는 시간 평균 보간물)과 훈련이 필요 없는 Poisson Flow 생성 모델($\nu \equiv 1$)을 통합한다.
3. 보존 방정식: 수송 맵을 특징짓는 고유한 보존 방정식($b \cdot \nabla T = 0$)을 도출하여, Beckmann 수송 문제의 플럭스(flux) 제약에 대한 동역학적 해석을 제공한다.
4. 직접 맵 학습: 흐름의 오일러 방정식에서 유도된 잔차 손실(residual loss)을 사용하여 샘플로부터 수송 맵 $T$를 직접 학습하는 방법을 제안하며, 이를 통해 추론 시 ODE 적분을 우회한다.
5. 편향 교정: 기존의 Equilibrium Matching 손실이 질량 할당에 편향을 유발함을 밝히고, 추가 비용 없이 이 문제를 해결하는 일관된 보간물을 제공한다.

실험 결과

저자들은 2D 원자 타겟(atomic targets) 및 이미지 벤치마크를 통해 이론을 검증한다:

• 2D 원자 타겟: 교정된 BTM 손실은 모드 전반에 걸쳐 거의 완벽한 질량 할당(MAE $\approx$ 0.005)을 달성한 반면, 기존의 Equilibrium Matching 손실은 잘못된 분지 구조(basin structures)로 인해 상당한 편향(MAE $\approx$ 0.102)을 보였다.
• 반복적 수렴: 고차원에 임베딩된 나선형 타겟(spiral target)에서, 부분적으로 훈련된 맵 $T$는 확산된 샘플을 생성하지만, 맵을 반복 적용하면(예: $T^3$) 생성이 날카로워지며 자율 흐름 궤적에 수렴한다. 완전히 훈련된 $T$는 단 한 번의 단계로 흐름과 일치한다.
• ImageNet 256x256:
  • 교정: Equilibrium Matching 아키텍처(EqM-XL/2)에 일관된 손실을 적용했을 때, 아키텍처나 훈련 예산을 변경하지 않고도 FID가 1.90에서 1.87로 개선되었다.
  • 원스텝 생성: SD-VAE 잠재 공간(latent space)에서 직접 맵을 훈련하여 단 한 번의 순전파(forward pass)만으로 FID 17.58을 달달성했다(classifier-free guidance 미사용). 이는 iCT(FID 30.10)와 같은 다른 원스텝 방법보다 우수하며, 가이던스를 필요로 하는 방법들과 경쟁할 만한 수준이지만, 다단계 확산 모델보다는 여전히 낮다.

의의 및 주장

본 논문은 폐쇄형 모델(Poisson Flow)과 최근의 평형 기반 방법들을 하나의 이론적 체계 아래 복구하는 통합 프레임워크를 제공한다고 주장한다. 자율 흐름을 Beckmann 수송 문제와 연결함으로써, 저자들은 Benamou–Brenier의 Monge–Kantorovich 최적 수송 공식과는 구별되는, 플럭스 제약에 대한 동역학적 해석을 제공한다.

그 의의는 다음과 같다:

• 이론적 엄밀성: 자율 흐름이 특이 매니폴드로 정확하게 수송할 수 있다는 누락된 증명을 제공한다.
• 실용적 효율성: 추론 시 ODE 적분이 필요 없는 원스텝 생성을 가능하게 한다.
• 기존 방법의 교정: 기존 Equilibrium Matching 손실 함수의 근본적인 편향을 식별하고 수정한다.
• 설계 유연성: 가중치 $\nu$를 설계 노브(design knob)로 도입하여, 서로 다른 생성 모델(예: FM vs. Poisson Flow) 간의 트레이드오프를 조절하고 타겟 매니폴드 근처에서의 드리프트 거동을 제어할 수 있게 한다.

저자들은 ImageNet에 대한 현재의 원스텝 BTM 성능이 최첨단 다단계 확산 모델에 비해 "보통(modest)" 수준임을 인정하며, FID 17.58이 이 접근 방식의 실현 가능성을 입증한다고 언급했다. 또한, BTM에 classifier-free guidance를 확장하는 것과 훈련 시간과 추론 반복 횟수 사이의 트레이드오프를 최적화하는 것을 향후 연구 과제로 제시했다.