Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

기술 요약: 비로그-오목 샘플링을 위한 분산 감소 기법 및 역문제에 대한 응용

문제 정의
비정규화 밀도(unnormalized densities)를 가진 고차원 비로그-오목 분포로부터 샘플링하는 것은 머신러닝, 특히 포텐셜 함수의 정확한 그래디언트를 사용할 수 없는 상황에서 근본적인 과제입니다. 이러한 시나리오에서는 그래디언트를 추정치를 통해 근사해야 하며, 이는 고정된 반복당 계산 예산 하에서 높은 분산을 나타내는 경우가 많습니다. 비볼록 최적화(non-convex optimization)에서 분산 감소 기술(예: 모멘텀을 사용한 SGD, STORM, PAGE)이 수렴 성능을 개선해 왔으나, 비로그-오목 분포로부터의 샘플링에 대한 적용은 여전히 미개척 분야로 남아 있습니다. 또한, 점수 기반 생성 모델(SGM)을 사용하는 역문제에 대한 기존의 사후 확률 샘플링 방법들은 일반적으로 풀 배치(full-batch) 환경에서 작동하며, 이는 3D 및 동적 영상과 같은 대규모 데이터셋에 대한 적용 가능성을 제한합니다.

방법론
저자들은 비볼록 최적화의 분산 감소 기술을 샘플링 영역으로 적응시키는 통합 프레임워크를 제안합니다. 이들은 두 가지 통합 추정기를 도입합니다:

1. 평균-립시츠 분산 감소 랑제빈 확산 (ML-VRLD): 이 추정기는 SGD와 모멘텀을 포함한 그 스토캐스틱 대응물을 통합합니다. 이는 정확한 그래디언트의 립시츠 연속성(가정 1)을 가정하며, 추정기가 미니 배치 그래디언트와 모멘텀 기반 업데이트의 가중 평균이 되는 통합된 처리를 제공합니다.
2. 샘플-립시츠 분산 감소 랑제뱅 확산 (SL-VRLD): 이 추정기는 PAGE와 STORM을 통합합니다. 이는 스토캐스틱 그래디언트의 평균 제곱 의미에서의 립시츠 연속성(가정 4)을 가정합니다. 이는 현재와 과거의 스토캐스틱 그래디언트 차이에 기반한 보정 항을 포함하며, 이는 샘플-립시츠 조건 하에서 추정 정확도를 향상시킵니다.

두 방법 모두 $\epsilon$-상대적 피셔 정보(FI) 수렴을 달哮하도록 설계되었으며, 이는 표준 SGLD의 $O(Ld\sigma^2/\epsilon)$ 요구 사항과 대조적으로 반복당 $O(1)$의 스토캐스틱 그래디언트 계산을 수행합니다.

이 프레임워크는 SGM 사전 확률(prior)을 사용하는 어닐드 랑제빈 몬테카를로(LMC) 샘플링 과정에 이 추정기들을 통합함으로써 역문제(예: MRI 및 CT 재구성)로 확장됩니다. 사후 확률 샘플링 정식화는 반복(iterates)이 평활화된 사전 확률에서 실제 사후 확률로 안내되도록 하는 어닐링 및 평활화 스케줄을 갖는 일련의 중간 타겟 분포를 활용합니다. 이론적 보증은 유계된 SGM 추정 오차 가정 하에 확립되었습니다.

주요 기여

• 통합 분석: 본 논문은 비로그-오목 분포로부터의 샘플링을 위한 분산 감소 추정기(특히 ML-VRLD 및 SL-VRLD)에 대한 최초의 통합 분석을 제공합니다.
• 개선된 복잡도:
  • ML-VRLD와 SL-VRLD 모두 반복당 $O(1)$의 그래디언트 계산으로 $\epsilon$-상대적 FI 수렴을 달성하며, 이는 SGLD의 반복당 복잡도를 크게 개선합니다.
  • SL-VRLD는 총 계산 복잡도를 $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD)에서 $O(L^2d^2\sigma^2/\epsilon^3)$로 더욱 감소시킵니다.
• 수렴 보증: 저자들은 푸앵카레 부등식(Poincaré inequality) 가정 하에 제곱 총 변동(TV) 거리에서의 비점근적 수렴율을 확립하고, 두 알고리즘 모두에 대해 타겟 분포로의 약수렴(weak convergence)을 증명합니다.
• 역문제 확장: 프레임워크는 SGM 사전 확률를 가진 사후 확률 샘플링으로 확장되며, 유계된 SGM 추정 오차 하에서 동일한 수렴 보증을 확립합니다. 이는 기존의 풀 배치 사후 확률 샘플링 방법들의 한계를 해결합니다.

결과

• 합성 실험: 가우시안 혼합 모델(GMM) 및 압축 센싱 역문제에 대한 합성 실험에서, 제안된 방법들은 SGLD에 비해 낮은 상대적 FI 값으로 더 빠르게 수렴함을 보여줍니다. 고정된 배치 크기를 사용하는 SGLD는 하위 최적의 FI 수준으로 수렴하는 반면, ML-VRLD와 SL-VRLD는 동일한 고정 배치 크기로 타겟 정확도에 도달합니다. SL-VRLD는 ML-VRLD보다 일관되게 타겟 분포에 더 빠르게 도달합니다.
• 통계적 검증: 알려진 그라운드 트루스가 있는 사후 확률 샘플링 작업에서, SGLD는 유의미하게 팽창된 분산을 생성했습니다. ML-VRLD와 SL-VRLD는 이 분산을 크게 감소시켰으며, SL-VRLD가 그라운드 트루스 사후 확률 분산에 가장 근접한 결과를 얻었습니다.
• 영상 적용:
  • MRI 재구성: 가속 계수가 8 및 12인 fastMRI 데이터셋에서, ML-VRLD와 SL-VRLD는 베이스라인(DPS, PnP-DM, SGLD 포함)에 비해 더 선명한 해부학적 세부 사항과 적은 아티팩트를 생성했습니다. SL-VRLD가 가장 높은 PSNR 및 SSIM 점수를 달성했습니다.
  • 희소 각도 CT 재구성: 180 및 360 투영 뷰를 가진 2DeteCT 데이터셋에서, 제안된 방법들은 베이스라인보다 미세한 구조적 세부 사항을 보존하는 데 있어 더 우수한 성능을 보였으며, 이는 고노이즈 및 언더샘플링 환경에서 분산 감소의 이점을 입증합니다.

의의
본 논문은 $O(1)$ 그래디언트 계산을 사용하는 분산 감소 추정기를 사용하여 비로그-오목 샘플링에 대한 최초의 비점근적 수렴 보증을 제공한다고 주장합니다. SGM 사전 확률를 가진 역문제로 이 기술들을 확장함으로써, 본 연구는 풀 배치 방법이 계산적으로 불가능한 대규모 영상 작업에 대한 원칙적인 사후 확률 샘플링을 가능하게 합니다. 실증적 결과는 분산 감소가 고정된 메모리 및 계산 제약 조건 하에서 샘플 품질과 통계적 정확도를 일관되적으로 개선함을 검증하며, 고차원 베이지안 추론을 위한 확장 가능한 솔루션을 제공합니다.