An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

기술적 요약: 무한 차원 선형 역문제에서 조건부 점수의 무조건부 표현

문제 제기
점수 기반 확산 모델 (SDM) 은 베이지안 역문제에서 사후 분포를 샘플링하는 강력한 도구가 되었으며, 기존 수동 설계 정규화 기법보다 우수한 유연한 사전 분포를 제공합니다. 그러나 기존 사후 샘플링 방법들은 계산 효율성과 확장성 사이에서 치명적인 트레이드오프에 직면해 있습니다.

1. 보정 기반 방법(예: 가능도 기울기를 점수에 추가하거나 데이터 일관성을 강제하는 방법) 은 샘플링 과정에서 전진 연산자 $A$의 반복적인 평가를 요구합니다. 이는 대규모 고차원 문제에서 금전적으로 prohibitive 한 계산 병목 현상을 초래합니다.
2. 직접 조건부 학습 방법은 측정치에 조건부인 점수 함수 ($s(x, t; \mu_y)$) 를 학습합니다. 이는 샘플링 중 전진 평가를 피하지만, 신경망의 입력 차원을 ($x$에서 $(x, y)$로) 급격히 증가시켜, 특히 무한 차원 설정에서 조건부 점수를 정확하게 근사하기 위해 광범위한 학습 데이터와 복잡한 아키텍처를 요구합니다.

본 논문은 전진 모델 평가의 높은 온라인 비용이나 직접 조건부 학습의 높은 샘플 복잡성을 치르지 않고도, 무한 차원 선형 역문제에서 효율적이고 확장 가능한 사후 샘플링을 수행하는 과제를 다룹니다.

방법론: 조건부 점수의 무조건부 표현 (UCoS)
저자들은 UCoS 를 제안하며, 이는 조건부 부여의 계산 부하를 온라인 샘플링 단계에서 오프라인 학습 단계로 이동시키는 프레임워크입니다. 핵심 통찰은 가우시안 노이즈를 가진 선형 역문제에서 조건부 점수 함수는 수정된 작업 의존적 확산 과정의 무조건부 점수 함수로부터 아핀 변환을 통해 정확하게 유도될 수 있다는 것입니다.

• 이론적 기반:

  • 이 방법은 이산화 불변성을 보장하는 분리 가능 힐베르트 공간 $\mathcal{H}$에서 작동합니다.
  • $s(x, t; \mu_y)$를 근사하기 위해 신경망을 학습하는 대신, 저자들은 사전 분포와 전진 연산자 $A$에 의존하는 가우시안 항의 혼합물인 분포를 갖는 새로운 확산 과정 $\tilde{X}_t$를 정의합니다.
  • 그들은 (정리 3.7) 진정한 조건부 점수 $s(x, t; \mu_y)$가 이 새로운 과정의 무조건부 점수 $\tilde{s}$와 다음 항등식을 통해 관련됨을 증명합니다:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    여기서 $r$은 무조건부 점수의 변환된 버전이고, $\lambda(t)$는 시간 의존 스칼라이며, $\xi_t(x, y)$는 측정치 $y$와 수반 연산자 $A^*$를 포함하는 아핀 변환입니다.
  • 결정적으로, 항 $A^*\Gamma^{-1}y$ (여기서 $\Gamma$는 노이즈 공분산) 는 샘플링 시작 전에 한 번 계산할 수 있습니다. 이후 샘플링 과정은 학습된 네트워크 $r_\theta$의 평가와 간단한 벡터 덧셈만 요구하며, 역확산 단계 중 $A$나 $A^*$의 추가 적용이 필요 없습니다.

• 구현:

  • 오프라인 단계: 신경망 (구체적으로 푸리에 신경 연산자, FNO) 이 함수 $r$을 근사하도록 학습됩니다. 학습은 수정된 과정 $\tilde{X}_t$에 대한 디노이징 점수 매칭 목적 함수를 활용합니다. 이 단계는 $A$와 $A^*$에 대한 접근이 필요하지만 특정 측정 데이터 $y$는 필요하지 않습니다.
  • 온라인 단계: 변환된 점수를 사용하여 역시간 SDE 를 풀어서 사후 샘플이 생성됩니다. 이 과정은 "행렬 프리"이며, 사전 계산된 시프트 $A^*\Gamma^{-1}y$와 네트워크 평가만 요구합니다.

주요 기여

1. 무한 차원에서의 정확한 항등식: 본 논문은 무한 차원 힐베르트 공간에서 아핀 변환을 통해 조건부 점수를 작업 의존적 무조건부 점수와 연결하는 엄밀한 이론적 항등식 (정리 3.7) 을 수립합니다. 이는 특정 유한 차원 이산화를 넘어 원리를 일반화합니다.
2. 수렴 분석: 저자들은 생성 과정에 대한 정량적 오차 한계 (정리 4.1) 를 제공하며, 학습 손실, 수치적 이산화, 초기화에서 기인하는 오차 기여를 명시적으로 특징화합니다. 이는 방법이 진정한 사후 측도로 수렴함을 정당화합니다.
3. 확장성과 효율성: 조건부 부여를 샘플링 루프에서 분리함으로써, UCoS 는 전진 평가가 없는 무조건부 방법의 계산 효율성을 달성하면서도 조건부 방법의 정확성을 유지합니다. 이는 결합 공간 $(x, y)$에서 네트워크를 학습할 때 발생하는 차원의 저주를 피합니다.

실험 결과
저자들은 이미지 인페인팅, 희소 뷰 컴퓨터 단층촬영 (CT), 이미지 디블러링이라는 세 가지 역문제에서 UCoS 를 검증합니다.

• 인페인팅 (가우시안 사전): 사후 분포가 가우시안이고 점수가 폐형식으로 계산될 수 있는 설정에서, UCoS 는 정확한 조건부 방법 및 기타 베이스라인과 비교 가능한 성능을 보여주어 이론적 항등식을 확인합니다.
• CT 영상 및 디블러링 (학습된 사전): 각각 LIDC-IDRI 및 CelebA 데이터셋을 사용하여, UCoS 는 SDE ALD, 확산 사후 샘플링 (DPS), 투영 기반 방법 (Proj), 그리고 직접 조건부 학습과 비교됩니다.
  • 정확도: UCoS 는 저복잡도 아키텍처 (예: 레이어당 32 개 노드) 를 사용할 때 무조건부 방법과 비교 가능한 또는 더 나은 편향과 표준 편차를 가진 샘플을 생성하며, 조건부 방법보다 현저히 우수합니다.
  • 효율성: UCoS 는 조건부 방법과 동일한 샘플링 속도를 기록합니다 (둘 다 전진 평가를 피함) 하지만, 반복적인 전진 통과가 필요한 보정 기반 방법 (DPS, Proj, SDE ALD) 보다 현저히 우수합니다.
  • 강건성: 결합 분포를 근사하는 복잡성 증가로 인해 네트워크 파라미터가 줄어들면 조건부 방법의 성능이 눈에 띄게 저하되는 반면, UCoS 는 더 단순한 아키텍처로도 고품질을 유지합니다.

의의 및 주장
본 논문은 UCoS 가 "양쪽 세계의 장점을 모두 가진" 솔루션을 제공한다고 주장합니다. 즉, 샘플링 중 전진 모델 평가의 계산 병목 현상 (보정 기반 방법의 한계) 을 제거하면서도, 직접 조건부 학습의 높은 샘플 복잡성과 아키텍처 요구 사항을 피합니다.

저자들은 그들의 접근 방식이 이산화 불변임을 강조하며, 이는 고차원 유한 차원 근사에서는 성능 보장이 종종 실패하는 무한 차원 역문제에 적합하다고 합니다. 그들은 이 방법이 작업 의존적이며, 전진 연산자 $A$가 변경될 경우 재학습이 필요함을 인정합니다. 또한 정확한 대응은 전진 모델의 선형성과 가능도의 가우시안성 (완전한 제곱 단계 허용) 에 의존합니다. 이 작업은 이론적 공식화에서 근사 오차 없이 조건부 복잡성을 학습 단계로 이동시키는 새로운 종류의 확산 기반 샘플러에 대한 개념 증명으로 제시됩니다.