An invertible generative model for forward and inverse problems

기술 요약: 순방향 및 역방향 문제를 위한 가역적 생성 모델

문제 정의
순방향 연산자 $K$에 대하여 $Ku = f$를 만족하는$u$를 찾는 것으로 정식화되는 역문제(Inverse problems)는 흔히 부적정(ill-posed)합니다. 변분법(예: 데이터 충실도와 정규화를 함께 최소화하는 방식)은 재구성(reconstruction)을 위한 프레임워크를 제공하지만, 계산 비용이 많이 들 수 있습니다. 반대로, 데이터 기반 접근 방식은 직접적인 역 매핑을 학습하지만, 본질적인 부적정성과 불확실성 정량화의 결여로 인해 어려움을 겪는 경우가 많습니다. 베이지안 프레임워크는 사후 분포를 특징짓는 엄격한 해결책을 제공하지만, 고차원 사후 분포로부터 샘플링하는 것은 계산적으로 매우 까다롭습니다(예: MCMC를 통한 방식).

현재의 생성 모델링 접근 방식은 시뮬레이션(우도 $\mu_{F|U=u}$로부터 샘 샘플링)과 추론(사후 분포 $\mu_{U|F=f}$로부터 샘플링)을 별도의 조건부 모델을 사용하는 별개의 작업으로 취급하는 경우가 많습니다. 이러한 분리는 순방향 과정과 역방향 과정 사이의 일관성을 강제하지 못하며, 두 개의 별도 모델을 훈련해야 하므로 계산 비용을 증가시킵니다.

방법론: 가역적 시뮬레이터 (The Reversible Simulator)
저자들은 단일 가역 맵인 가역적 시뮬레이터($S: \mathbb{R}^{n+m} \to \mathbb{R}^{n+m}$)를 중심으로 하는 통합 프레임워크를 제안합니다. 이 맵은 적용 방향에 따라 시뮬레이션과 추론을 모두 수행하도록 설계되었습니다:

1. 순방향 방향 (시뮬레이션): 잠재 변수 $Y$와 파라미터 $u$가 주어지면, $S$는 $(u, Y)$를 $(X, F)$로 매핑하며, 여기서 $F$는 우도 $\mu_{F|U=u}$를 따릅니다.
2. 역방향 방향 (추론): 잠재 변수 $X$와 관측값 $f$가 주어지면, 역 맵 $R = S^{-1}$은 $(X, f)$를 $(U, Y)$로 매핑하며, 여기서 $U$는 사후 분포 $\mu_{U|F=f}$를 따릅니다.

삼각 변환 맵(Triangular Transport Maps)을 통한 구성
본 논문은 두 개의 삼각 변환 맵(Knothe-Rosenblatt rearrangement)을 결합하여 $S$를 명시적으로 구성하는 방법을 제시합니다:

• 하삼각 맵 ($\tilde{F}$): 참조 측도를 결합 분포 $\mu_{U,F}$로 밀어내어(push-forward), $u$에 조건화된 우도의 조건부 샘플링을 용이하게 합니다.
• 상삼각 맵 ($\hat{F}$): 마찬가지로 $f$에 조건화된 사후 분포의 조건부 샘플링을 용이하게 합니다.

가역적 시뮬레이터 $S$는 이 맵들과 그 역함수들의 합성으로 구성됩니다. 구체적으로, $S = U \circ L$이며, 여기서 $L$은 하삼각 구조를 사용하여 $u$에 조건화하고, $U$는 상삼각 구조를 사용하여 $f$에 조건화합니다. 이러한 구성은 $S$가 사전 분포와 잠재 노이즈의 곱($\mu_U \otimes \mu_Y$)과 잠재 노이즈와 주변 데이터의 곱($\mu_X \otimes \mu_F$) 사이의 유효한 변환 맵임을 보장합니다.

변분 정식화 및 훈련
일반적인 가역 신경망 아키텍처를 사용할 수 있도록(단순히 명시적인 삼각 구조에만 의존하지 않고), 저자들은 변분 정식화를 도출합니다. 이들은 거리 척도 $M$(예: 와서스테인 거리)을 기반으로 한 네 가지 손실 함수를 정의합니다:

1. $J_1(S)$: 사전-노이즈 곱의 밀어내기(push-forward)와 결합 분포 사이의 불일치를 측정합니다(우도 성분을 통해).
2. $J_2(R)$: 사후 분포 성분에 대한 불일치를 측정합니다.
3. $J_3(S)$ 및 $J_4(R)$: 순방향 및 역방향 방향에서 결합 분포의 밀어내기 사이의 불일치를 측정합니다.

저자들은 이 손실 함수들을 최소화하는 것(특히 $J_1$과 $J_3$)이 학습된 맵이 가역적 시뮬레이터의 속성을 만족하도록 보장함을 증명합니다. 이를 통해 모델은 명시적인 밀도에 대한 지식 없이도 쌍으로 된 샘플 $\{(u_i, f_i)\}_{i=1}^N$로부터 직접 훈련될 수 있습니다.

주요 기여

1. 통합 프레임워크: 순방향 및 역방향 문제 모두를 위한 단일 가역 맵으로서 시뮬레이션과 추론을 통합하는 "가역적 시뮬레이터"를 도입했습니다.
2. 명시적 구성: 두 개의 삼각 노멀라이징 플로우를 결합하여 이러한 시뮬레이터의 명시적인 사례를 구축했습니다.
3. 이론적 분석: 가역적 시뮬레이터의 존재성과 비유일성을 조사하여, 측도 보존 변환과의 합성을 통해 무수히 많은 그러한 맵이 존재함을 보여주었습니다.
4. 훈련 목적 함수: 표준 생성 모델 아키텍처를 사용하여 쌍으로 된 데이터로부터 직접 훈련할 수 있게 하는 변분 정식화를 도출했습니다.
5. 수치적 검증: 분석적으로 다루기 쉬운 선형 가우시안 문제, 비선형 합성 예시, 그리고 MNIST 데이터셋을 이용한 정형화된 이미지 인페인팅(inpainting) 작업에서 프레임워크를 입증했습니다.

결과

• 선형 가우시안 케이스: 선형 역문제와 가우시안 사전/우도가 있는 경우, 구성된 가역적 시뮬레이터는 순방향 모델이 결정론적일 때(즉, 노이즈 분산이 0에 가까워질 때)도 잘 정의됨을 보여주었습니다. 이는 표준 삼각 맵이 퇴화할 수 있는 것과 대조적입니다.
• 비선형 예시: 구성 방식이 비선형 의존성을 성공적으로 처리하며, 하삼각, 상삼각, 그리고 결합된 가역 맵의 서로 다른 변환 플랜(transport plans)을 시각화합니다.
• 이미지 인페인팅 (MNIST): 가역적 시뮬레이터에 적응된 GLOW 기반 아키텍처를 사용하여, 모델은 다음을 성공적으로 학습했습니다:
  • 원래 숫자의 구조를 보존하면서도 순방향 연산자(노이즈 추가 및 마스킹)를 정확하게 반영하는 우도 샘플 생성.
  • 누락된 부분의 숫자를 정확하게 재구성(인페인팅)하고, 마스킹된 영역에 불확실성을 정확하게 국지화하는 사후 샘플 생성.
  • 모델은 OOD(Out-of-distribution) 데이터(문자 'A' 및 'F')에 대해서도 일반화 능력을 보여주었으나, 이 문자들에 대한 우도 샘플은 학습된 순방향 연산자가 훈련 분포(숫자)의 특정 픽셀 값에 영향을 받았음을 나타내며, 이는 순방향 연산자가 미지의 클래스에 대해 완벽하게 학습되지 않았음을 시사합니다.

의의 및 주장
본 논문은 가역적 시뮬레이터가 순방향 및 역방향 문제 모두를 위한 조건부 생성 모델링에 대한 원칙적이고 통합된 접근 방식을 제공한다고 주장합니다. 단일 가역 맵을 사용함으로써, 이 프레임워크는 별도의 모델을 훈련하는 계산 부담을 줄이고 시뮬레이션과 추론 사이의 일관성을 보장합니다. 저자들은 이 구성이 표준 삼각 플로우에 비해 학습된 매핑의 조건화(conditioning)를 개선할 수 있다고 강조합니다. 저자들은 이 연구를 일반적인 가역 신경망과 아마티즈드 베이지안 추론(amortized Bayesian inference)을 위한 노멀라이징 플로우 사이의 가교 역할을 한다고 보며, 이러한 매핑을 데이터로부터 직접 학습하기 위한 이론적 토대를 제공합니다. 또한, 현재의 유일성 결과가 확립되었지만, 가능한 모든 가역적 시뮬레이터를 더 상세히 규명하는 것이 훈련을 위한 최적성 기준을 도출하는 데 기여할 수 있다고 언급했습니다. 마지막으로, 디퓨전 모델(diffusion models)과 같은 다른 생성 모델로의 확장 및 실험 설계에 대한 응용 가능성을 제시했습니다.