[논문 리뷰] Likelihood Training of Schr\\"odinger Bridge using Forward-Backward SDEs Theory

Tianrong Chen, Guan-Horng Liu외 1명|arXiv (Cornell University)|2021. 10. 21.

Generative Adversarial Networks and Image Synthesis인용 수 11

한 줄 요약

이 논문은 전진-역행 확률미분방정식(FBSDE) 이론을 활용하여 슈뢰딩거 브리지(SB) 모델의 새로운 가능도 훈련 프레임워크인 SB-FBSDE를 소개한다. SB에 대한 정확한 로그가능도 표현식을 유도함으로써 SGM 가능도 목표를 일반화함으로써, 원칙적인 현대적 훈련 기법을 가능하게 하여, 평가 시 랑주뱅 보정을 필요로 하지 않고 MNIST, CelebA 및 CIFAR10에서 최고 성능의 FID 점수를 달성한다.

ABSTRACT

Schr\\"odinger Bridge (SB) is an entropy-regularized optimal transport problem that has received increasing attention in deep generative modeling for its mathematical flexibility compared to the Scored-based Generative Model (SGM). However, it remains unclear whether the optimization principle of SB relates to the modern training of deep generative models, which often rely on constructing log-likelihood objectives.This raises questions on the suitability of SB models as a principled alternative for generative applications. In this work, we present a novel computational framework for likelihood training of SB models grounded on Forward-Backward Stochastic Differential Equations Theory - a mathematical methodology appeared in stochastic optimal control that transforms the optimality condition of SB into a set of SDEs. Crucially, these SDEs can be used to construct the likelihood objectives for SB that, surprisingly, generalizes the ones for SGM as special cases. This leads to a new optimization principle that inherits the same SB optimality yet without losing applications of modern generative training techniques, and we show that the resulting training algorithm achieves comparable results on generating realistic images on MNIST, CelebA, and CIFAR10. Our code is available at https://github.com/ghliu/SB-FBSDE.

연구 동기 및 목표

• 슈뢰딩거 브리지(SB) 최적화와 현대 딥 러닝 생성 모델링, 특히 로그가능도 목표에 기반한 스코어 기반 생성 모델(SGM) 간의 이론적 원칙적인 연결 고리를 구축하기 위해.
• SB 모델이 SGM의 이론적으로 타당한 대안이 될 수 있는지에 대한 모호성을 제거하기 위해, SB 훈련을 가능도 프레임워크에 기반하게 하기 위해.
• 이전의 수단적 또는 IPF 기반 최적화의 한계를 극복하기 위해, FBSDE 이론을 활용한 가능도 기반 SB 모델 훈련을 위한 계산 프레임워크를 개발하기 위해.
• 선형 확산 과정에서의 SGM 가능도 목표를 완전히 비선형 확산 과정으로 일반화함으로써, SB와 SGM을 동일한 가능도 원칙 아래 통합하기 위해.
• 이 가능도 프레임워크를 통해 훈련된 SB 모델이 표준 벤치마크에서 경쟁력 있거나 우수한 이미지 생성 성능을 달성할 수 있음을 입증하기 위해.

제안 방법

• 스토케스틱 최적 제어 이론을 활용하여 PDE 기반의 최적성 조건을 다룰 수 있는 전진-역행 SDE(FBSDE) 시스템으로 변환함으로써, SB의 최적성 조건을 FBSDE로 유도한다.
• FBSDE 시스템을 푸름으로써 SB에 대한 정확한 로그가능도 표현식을 구성하며, 이는 선형 확산일 경우 SGM 가능도 목표의 특수한 경우로 일반화된다.
• 이전의 평균 매칭 회귀 접근 방식을 대체하기 위해, 산란 기반 목표함수(예: KL 산란)를 사용하여 전진 및 역행 과정을 훈련한다.
• 표준 딥 러닝 기법을 활용하여 전진 및 역행 SDE를 공동 최적화할 수 있도록 하며, 특히 SGM 유사 랑주뱅 보정기를 활용한 개선도 가능하다.
• 유도된 가능도 목표를 기반으로 전진 및 역행 정책을 번갈아 최적화함으로써 안정적이고 확장 가능한 훈련을 가능하게 한다.
• 이 방법은 PyTorch 기반 코드베이스로 구현되었으며, U-Net 및 NCSN++와 같은 표준 아키텍처를 사용하여 MNIST, CelebA 및 CIFAR10에서 평가되었다.

실험 결과

연구 질문

• RQ1 슈뢰딩거 브리지가 SGM 가능도 프레임워크를 일반화하는 가능도 목표를 통해 훈련될 수 있는가?
• RQ2 SB 최적성의 FBSDE 기반 공식화는 현대 딥 러닝 생성 모델링 기법과 어떤 식으로 원칙적인 연결 고리를 형성하는가?
• RQ3 FBSDE 기반의 SB 가능도 훈련은 이전의 평균 매칭 회귀 기반 SB 방법에 비해 샘플 품질과 다양성 측면에서 향상된 성능을 보이는가?
• RQ4 유도된 가능도 목표는 랑주뱅 보정기와 같은 현대적 훈련 개선 기법을 SB 모델에 통합하는 데 사용될 수 있는가?
• RQ5 표준 이미지 생성 벤치마크인 MNIST, CelebA 및 CIFAR10에서 가능도 기반 훈련된 SB 모델의 정량적 성능은 어떠한가?

주요 결과

• SB-FBSDE 프레임워크는 SGM 가능도 목표를 특수한 경우로 일반화하는 슈뢰딩거 브리지에 대한 정확한 로그가능도 표현식을 도출함으로써, 두 모델 클래스 간의 근본적인 이론적 연결 고리를 확립한다.
• CIFAR10에서 이 방법은 10,000개의 생성 샘플을 사용하여 Fréchet Inception Distance(FID) 11.85를 달성하였으며, 이는 이전의 최적 운반 모델을 크게 능가한다.
• SGM 회귀 후 50,000단계에서 5,000회 반복의 산란 기반 훈련만으로도 FID가 33.68에서 13.35로 감소하여 빠른 성능 향상을 보였다.
• 정성적 결과는 특히 CelebA의 얼굴 특징 및 배경 변화 측면에서 이전의 SB 방법에 비해 더 높은 이미지 다양성과 시각적 품질을 보였다.
• FID 평가 시 랑주뱅 보정기를 사용하지 않더라도 MNIST, CelebA 및 CIFAR10에서 최고 성능을 달성하여 내재된 뛰어난 샘플 품질을 입증하였다.
• 프레임워크는 랑주뱅 보정기와 같은 현대적 훈련 기법의 사용을 가능하게 하여, 아키텍처 및 데이터셋 크기에 관계없이 일관된 향상 효과를 보였다.