ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

기술 요약: ShuffleFlow

문제 정의
베이지안 역영상(Bayesian inverse imaging)은 노이즈가 포함된 불완전한 측정값 $y = A(x) + n$으로부터 잠재 이미지 $x$를 복원하는 것을 목표로 하며, 이는 사후 분포 $p(x|y)$를 추정하는 과정이다. 변분 추론(Variational Inference, VI)은 최적화 후 효율적인 샘플 생성을 가능하게 하는 원칙적인 사후 추론 접근법을 제공하지만, 대규모 영상 재구성 작업에 적용할 경우 확장성 문제에 직면한다. 흐름 기반 네트워크(flow-based networks)에 의존하는 기존의 VI 방식은 일반적으로 전체 이미지 분포를 직접 모델링하며, 이로 인해 계산 비용이 이미지 차원(예: $N \times N$ 픽셀에 대해 $O(N^4)$)에 따라 매우 나쁘게 증가한다. 이는 최적화 과정에서 감당하기 어려운 메모리 및 시간 예산을 초아한다. 네이티브 패칭(naive patching)이나 픽셀화(pixelating) 전략과 같은 기존의 확장성 개선 시도들은 종종 이미지 품질이나 사후 분포의 표현력을 희생시킨다. 반대로, 확산 모델(diffusion models)은 강력하지만, 학습을 위해 대규모 데이터셋이 필요하며 반복적인 샘플링 특성으로 인해 고해상도, 고샘플 수의 사후 분포를 생성할 때 높은 계산 비용이 발생한다.

방법론: ShuffleFlow
저자들은 사후 분포 모델링 작업을 좌표 샘플링, 특징 인코딩, 사후 추정의 세 가지 모듈형 구성 요소로 분해하도록 설계된 확장 가능한 VI 프레임워크인 ShuffleFlow를 제안한다.

1. 좌표 샘플링 (Pixel-Unshuffling): 전체 $N \times N$ 이미지를 모델링하거나 독립적인 패치를 사용하는 대신, ShuffleFlow는 다운샘플링 계수 $S$를 사용하는 픽셀 언셔플링(pixel-unshuffling) 연산을 통해 이미지를 $S^2$개의 상관관계가 있는 서브 이미지 스택으로 분할한다. 각 서브 이미지는 $(N/S) \times (N/S)$의 차원을 갖는다. 이는 문제의 차원을 $N^2$에서 $N^2/S^2$로 줄여준다.
2. 특징 인코딩 (Neural Field): 흐름 네트워크가 전체 이미지 재구성 작업의 부담을 지지 않으면서 공간 구조를 포착할 수 있도록, 이 프레임워크는 뉴럴 필드(Neural Field, NF) 인코더(구체적으로 좌표 기반 MLP인 SIREN)를 채택한다. 샘플링 위치의 2D 좌표 오프셋은 NF에 입력되어 각 서브 이미지에 대한 특징 임베딩($h_i$)을 생성한다.
3. 사후 추정 (Conditional Normalizing Flow): 공유된 조건부 정규화 흐름(Conditional Normalizing Flow, CNF)이 서브 이미지 스택의 결합 분포를 모델링한다. CNF는 각 서브 이미지의 위치에 대응하는 특징 임베딩($h_i$)에 의해 조건화된다. 결정적으로, 이 프레임워크는 모든 서브 이미지에 대해 단일 기본 잠재 변수 $z_0$를 공유하는 드 피네티(De Finetti) 영감을 받은 접근 방식을 채택한다. 이를 통해 결합 분포를 다음과 같이 분해할 수 있다:
   $$q(x) = \left( \prod_{i=1}^{S^2} q(x_i | z_0; c_i) \right) q(z_0)$$
   여기서 $z_0$는 표준 정규 분포에서 추출된다. 이러한 공유 메커니즘은 서브 이미지 간의 상관관계를 형성하여 픽셀 독립성과 관련된 노이즈 아티팩트를 방지하는 동시에, 흐름 네트워크의 크기를 $O(N^4)$에서 $O(N^4/S^4)$로 크게 줄여준다.

최적화 및 추론
CNF와 NF는 변분 분포와 베이지안 사후 분포 사이의 쿨백-라이블러(KL) 발산을 최소화함으로써 공동으로 학습된다. 손실 함수는 데이터 가능도(likelihood), 사전 확률(Total Variation 또는 가중치가 고정된 학습된 스코어 기반 사전 확률), 그리고 다양성을 장려하기 위한 음의 엔트로피 항의 세 가지 항으로 구성된다. 추론 시에는 NF가 샘플링 좌표에 대한 임베딩을 생성하고, CNF는 이러한 임베딩과 공유된 기본 변수에 조건화된 샘플을 생성하며, 생성된 샘플은 픽셀 셔플링(pixel-shuffling)을 통해 전체 이미지로 재조립된다.

주요 기여

1. 문제 분해: 저자들은 대규모 이미지 사후 모델링을 좌표 샘플링, 부분 이미지 표현, 경량화된 불확실성 정량화로 분해하는 새로운 모듈형 VI 프레임워크를 도입했다. 이 구조는 특정 네트워크 아키텍처에 구애받지 않으며, 뉴럴 필드 및 생성적 사전 확률(generative priors)의 발전을 통합할 수 있게 한다.
2. 확장성 및 효율성: 픽셀 언셔플링과 공유 잠재 변수를 활용함으로써, ShuffleFlow는 모델 크기와 계산 자원을 획기적으로 줄인다. 이론적 및 경험적 분석에 따르면, 이 방법은 픽셀 단위 추론만큼 효율적으로 확장되며, $256 \times 256$ 이미지($S=8$일 때)에 대해 전체 이미지 흐름 모델과 비교하여 흐름 네트워크 크기를 4096배 감소시킨다.
3. 고해상도 사후 분포 생성: 이 방법은 고해상도 샘플을 빠르게 생성하는 능력을 보여준다. 샘플링 시간이 샘플 수에 따라 선형적으로 증가하는 확산 샘플러와 달리, ShuffleFlow는 짧은 학습 단계 이후 수천 개의 샘플을 효율적으로 생성할 수 있다.

실험 결과
저자들은 선형(모션 디블러링, 압축 센싱 MRI) 및 비선형(푸리에 위상 복원) 역문제 모두에서 ShuffleFlow를 평가했다.

• 선형 문제: ShuffleFlow는 이미지 품질(PSSR, SSIM, LPIPS) 및 불확실성 보정(NLL, ECE) 측면에서 기존의 흐름 기반 VI 방식(DPI) 및 네이티브 패칭 전략을 능가하며 VI 방법들 중 최고 수준(SOTA)의 성능을 달야냈다. 이미지 품질 면에서는 선도적인 확산 샘플러(예: DAPS)와 대등한 성능을 보이면서도 계산 오버헤드는 현저히 낮았다.
• 비선형 문제: (직립 이미지 대 회전 이미지와 같은) 이봉형(bimodal) 사후 분포를 보이는 도전적인 푸리에 위상 복원 작업에서, ShuffleFlow는 이봉성을 성공적으로 포착했다. 반면, CF-NeRF 및 네이티브 패칭과 같은 방법들은 이러한 복잡한 전역적 상관관계를 모델링하는 데 실패했다.
• 효율성: 10,000개의 사후 샘플이 필요한 시나리오에서, ShuffleFlow는 16분 이내에 고해상도 이봉형 사후 분포를 생성한다. 동일한 시간 동안 확산 샘플러(DAPS, DPS)는 훨씬 적은 수의 샘플을 생성하며 두 번째 모드를 명확히 특징짓지 못하며, 유사한 해상도를 달성하기 위해 약 20배 더 많은 시간을 필요로 한다.

의의 및 주장
본 논문은 ShuffleFlow가 사후 정확도를 희생하지 않으면서 과학적 이미징을 위한 변분 추론의 핵심적인 확장성 병목 현상을 해결한다고 주장한다. 이 방법은 특히 확산 모델을 학습시킬 수 없는 데이터 부족 환경과 복잡한 고해상도 사후 분포(예: 비선형 역문제에서 이봉 분포 식별)의 빠른 탐색이 필요한 응용 분야에서 VI를 확산 모델의 중요한 대안으로 포지셔닝한다. 저자들은 이 방법이 매우 효율적이지만, 현재 학습 불안정성으로 인해 비선형 설정에서 다소 과도한 확신(overconfidence)을 보이는 경향이 있으며, 이를 향후 연구 과제로 식별하였다. 이 프레임워크는 사후 복잡성이 광범한 샘플 생성을 정당화하는지 판단할 수 있는 탐색적 도구로 제시된다.