[논문 리뷰] Rectified Flow: A Marginal Preserving Approach to Optimal Transport

Qiang Liu|arXiv (Cornell University)|2022. 09. 29.

Stochastic Gradient Optimization Techniques인용 수 10

한 줄 요약

이 논문은 입력 분포 $\pi_0$ 와 $\pi_1$ 를 유지하면서 사용자가 지정한 볼록 비용 함수 $c$ 를 반복적으로 최소화하는 최적 운반(OT)을 위한 새로운 마진 유지형 흐름 기반 방법을 소개한다. 직선 보정 흐름을 단일 비용 함수를 향해 적응시킴으로써, 신경 ODE에 대한 비제약 회귀를 통한 단조성 내부 최적화 과정을 수립하여, 외부에서 제약 조건을 강제하지 않더라도 최적의 결합에 수렴함을 보장한다.

ABSTRACT

We present a flow-based approach to the optimal transport (OT) problem between two continuous distributions $π_0,π_1$ on $\mathbb{R}^d$, of minimizing a transport cost $\mathbb{E}[c(X_1-X_0)]$ in the set of couplings $(X_0,X_1)$ whose marginal distributions on $X_0,X_1$ equals $π_0,π_1$, respectively, where $c$ is a cost function. Our method iteratively constructs a sequence of neural ordinary differentiable equations (ODE), each learned by solving a simple unconstrained regression problem, which monotonically reduce the transport cost while automatically preserving the marginal constraints. This yields a monotonic interior approach that traverses inside the set of valid couplings to decrease the transport cost, which distinguishes itself from most existing approaches that enforce the coupling constraints from the outside. The main idea of the method draws from rectified flow, a recent approach that simultaneously decreases the whole family of transport costs induced by convex functions $c$ (and is hence multi-objective in nature), but is not tailored to minimize a specific transport cost. Our method is a single-object variant of rectified flow that guarantees to solve the OT problem for a fixed, user-specified convex cost function $c$.

연구 동기 및 목표

• 이중성 또는 정규화에 의존하지 않고 연속 최적 운반에서 무한차원의 마진 제약 조건을 이행하는 데 도전하는 것.
• 운반 비용을 감소시키기 위해 가용 결합 집합 내에서 이동하는 단조성 내부 최적화 방법을 개발하는 것.
• 원래 다중 목표 방법이었던 보정 흐름을 고정된 사용자 지정 비용 함수 $c$ 를 갖는 단일 목표 OT에 적응시키는 것.
• 고차원 연속 설정에서 데이터 기반 최적 운반을 가능하게 하기 위해 단순하고 비제약적인 회귀 문제의 시퀀스를 사용하는 것.

제안 방법

• 이 방법은 시간 $t \in [0,1]$ 동안 결합 $(X_0, X_1)$ 이 진화하는 신경 평범 미분방정식(ODE)의 시퀀스를 구성하며, 이때 마진 법칙 $\mathrm{Law}(X_0) = \pi_0$ 와 $\mathrm{Law}(X_1) = \pi_1$ 가 유지된다.
• 각 단계에서 속도장 $v^X_t$ 는 선형 보간 경로 $X_t = tX_1 + (1-t)X_0$ 에 따라 $X_1 - X_0$ 의 이동을 예측하기 위해 최소 제곱 회귀를 통해 학습된다.
• 결과적으로 생성된 ODE 흐름 $\mathrm{d}Z_t = v^X_t(Z_t)\mathrm{d}t$ 는 변환 전반에 걸쳐 마진 분포가 그대로 유지됨을 보장한다.
• 이 방법은 특정 비용 함수 $c$ 를 향해 목표로 하는 보정 흐름의 변종으로, 다중 목표 문제를 $c$-보정 흐름 절차를 통해 단일 목표 문제로 변환한다.
• 최적화 과정은 반복적으로 수행되며, 각 단계에서 비제약 회귀 문제를 해결함으로써 $c$-운반 비용이 안정적이고 단조롭게 감소한다.
• 이론적 분석은 이 방법이 OT 문제의 KKT 조건을 만족하며, 최적의 매핑이 볼록 쌍대 함수의 기울기로 나타남을 확인한다.

실험 결과

연구 질문

• RQ1 외부 제약 강제를 피하면서 무한차원의 마진 제약 조건을 만족하는 연속 최적 운반을 위한 단조성 내부 접근법을 구성할 수 있는가?
• RQ2 보정 흐름을 다수의 볼록 비용을 동시에 최소화하는 데서 비롯된 것에서, 특정 사용자 정의 볼록 비용 함수 $c$ 를 최소화하도록 어떻게 적응시킬 수 있는가?
• RQ3 반복적인 $c$-보정 흐름 절차의 수렴 행동과 최적성 갭은 어떠한가?
• RQ4 실제로 반복적으로 $c$-보정 흐름을 적용할 때 오차 누적이 어떻게 완화될 수 있는가?
• RQ5 보정 흐름은 다중 목표 최적 운반의 어떤 변종을 해결하는가? 그리고 고정된 $c$ 를 갖는 단일 목표 OT는 언제 더 바람직한가?

주요 결과

• 이 방법은 비제약 회귀 문제의 시퀀스를 통해 $c$-운반 비용을 단조롭게 감소시키며, 최적의 결합에 수렴함을 보장한다.
• 마진 제약 조건은 ODE 흐름을 통해 자동으로 유지되어, 라그랑주 이중성 또는 엔트로피 정규화의 필요성이 없어진다.
• 이론적 분석은 이 방법이 OT 문제의 KKT 조건을 만족하며, 최적의 매핑 $v_t = \nabla c^*(\nabla \lambda_t)$ 와 $\dot{\lambda}_t + c^*(\nabla \lambda_t) = 0$ 이 성립함을 확인한다.
• 이 방법은 안정성과 제약 조건 이행이 열악한 민맥스 및 정규화 기반 접근법에 대한 실용적인 대안을 제공한다.
• 반복적 적용에서의 오차 누적 문제는 최종 $Z_1^k$ 를 대상 $\pi_1$ 과 일치시키기 위해 순열 또는 이질성 최소화를 사용함으로써 완화될 수 있다.
• 이 방법은 고차원 연속 설정에서 데이터 기반 OT 를 가능하게 하며, 마진 충실도 유지를 위한 선형 시간 조정 단계 적용 가능성이 있다.