재시작 기반 후방 샘플링을 이용한 확산 모델 역문제 해결

초록

역문제는 불완전하거나 잡음이 섞인 측정값으로부터 근본적인 신호나 상태를 추정하는 과학·공학의 핵심 과제이다. 최근 확산 모델을 강력한 암시적 사전분포로 활용하는 접근법이 주목받고 있는데, 이는 복잡한 데이터 분포를 포착할 수 있기 때문이다. 그러나 기존 확산 기반 역문제 방법은 사후분포에 대한 강한 근사에 의존하거나, 스코어 네트워크를 통한 그래디언트 역전파에 높은 계산 비용이 소요되며, 선형 측정 모델에만 적용 가능한 제한점이 있다. 본 연구에서는 사전 학습된 확산 모델을 활용해 선형·비선형 역문제를 모두 해결할 수 있는 일반적이고 효율적인 프레임워크인 Restart for Posterior Sampling (RePS)을 제안한다. RePS는 무조건적 확산에서 샘플 품질을 향상시키는 것으로 알려진 재시작 기반 샘플링 아이디어를 사후 추론으로 확장한다. 본 방법은 어떠한 미분 가능 측정 모델에도 적용 가능한 조건부 ODE를 사용하고, 샘플링 과정에서 누적되는 근사 오차를 수축시키는 단순화된 재시작 전략을 도입한다. 또한 스코어 네트워크에 대한 역전파를 회피함으로써 계산 비용을 크게 절감한다. 다양한 역문제(선형·비선형)를 대상으로 한 실험에서 RePS는 기존 확산 기반 베이스라인에 비해 빠른 수렴과 우수한 복원 품질을 입증한다.

상세 요약

본 논문은 확산 모델을 사전 확률로 활용하는 최신 역문제 해결 방법들의 한계를 명확히 짚고, 이를 극복하기 위한 새로운 프레임워크인 Restart for Posterior Sampling (RePS)를 제시한다. 기존 접근법은 크게 두 가지 문제점에 봉착한다. 첫째, 사후분포를 근사하기 위해 Langevin dynamics나 score‑based MCMC와 같은 복잡한 샘플링 절차를 사용하면서, 스코어 네트워크에 대한 반복적인 역전파가 필요해 GPU 메모리와 연산량이 급증한다. 둘째, 대부분의 방법이 선형 측정 모델에 국한되거나, 비선형 모델을 다루기 위해 추가적인 변분 추정이나 제약식 도입 등 복잡한 수식적 변형을 요구한다. 이러한 제약은 실제 공학·의료 분야에서 흔히 마주치는 비선형 측정(예: 광학 회절, 전자 현미경) 상황에 적용하기 어렵게 만든다.

RePS는 두 가지 핵심 아이디어로 이러한 문제를 해결한다. 첫째, “재시작”(restart) 메커니즘을 조건부 ODE 흐름에 통합한다. 기존 무조건적 확산에서 재시작은 샘플이 특정 시간 구간에서 발산하거나 품질이 저하될 때, 해당 구간을 초기화하고 다시 진행함으로써 전체 경로의 안정성을 높인다. 논문은 이를 사후 샘플링에 그대로 적용해, 측정 모델에 의해 부과된 데이터 일관성 제약이 ODE 흐름을 왜곡시킬 때마다 재시작을 수행한다. 이때 재시작은 현재 상태를 일정 비율로 수축시키는 간단한 선형 변환으로 구현되어, 누적된 근사 오차를 효과적으로 감소시킨다.

둘째, RePS는 스코어 네트워크에 대한 역전파를 전혀 수행하지 않는다. 대신, 측정 모델의 미분가능성을 이용해 조건부 ODE의 오른쪽 항을 직접 계산한다. 이는 기존 방법이 필요로 했던 “gradient through score” 단계—즉, ∇θ sθ(x, t)와 같은 파라미터에 대한 미분—를 완전히 생략한다. 결과적으로 메모리 사용량이 크게 줄어들고, 배치 크기를 늘리거나 고해상도 이미지에 적용하는 것이 현실적으로 가능해진다.

실험 부분에서도 저자는 선형 역문제(압축 센싱, MRI 재구성)와 비선형 역문제(광학 전송 함수 역추정, 비선형 전자 현미경 이미지 복원)를 포괄적으로 평가한다. 정량적 지표인 PSNR·SSIM뿐 아니라 수렴 속도(반복 횟수 대비 재구성 오차 감소)에서도 RePS가 기존 Diffusion Posterior Sampling, Score‑Based Inverse Problems 등과 비교해 현저히 우수함을 보인다. 특히 비선형 설정에서 RePS는 측정 모델의 복잡도에 비례해 연산량이 급증하지 않는 점이 큰 장점으로 부각된다.

하지만 몇 가지 한계도 존재한다. 첫째, 재시작 전략의 하이퍼파라미터(재시작 시점, 수축 비율)가 문제마다 민감하게 작용할 가능성이 있다. 논문에서는 경험적으로 몇 가지 값을 제시했지만, 자동 튜닝 메커니즘이 없으면 실무 적용 시 추가적인 실험이 필요할 것이다. 둘째, 조건부 ODE를 풀기 위해서는 고정밀 수치 적분이 요구되는데, 이는 여전히 연산 비용을 완전히 없애지는 못한다. 다만 역전파를 배제함으로써 전체 파이프라인의 효율성은 크게 개선된 점은 분명하다.

향후 연구 방향으로는 (1) 재시작 파라미터를 베이지안 최적화나 강화학습을 통해 자동 조정하는 방법, (2) 비미분 가능 측정 모델(예: 양자화, 이산화)에도 적용 가능한 확장, (3) 다중 모달리티(예: 이미지·텍스트 동시 복원) 상황에서의 공동 사후 샘플링 프레임워크 구축 등이 제시될 수 있다. 전반적으로 RePS는 확산 모델을 실제 엔지니어링 문제에 적용하기 위한 중요한 전진이며, 계산 효율성과 일반성을 동시에 만족시키는 드문 사례로 평가할 수 있다.