Constrained Particle Seeking: Solving Diffusion Inverse Problems with Just Forward Passes

이 논문은 역전파와 같은 기울기 정보가 불필요한 '제약된 입자 탐색 (CPS)'이라는 새로운 무기울기 (gradient-free) 접근법을 제안하여, 전방향 통과 (forward passes) 만으로 확산 모델을 활용한 역문제 해결을 가능하게 하고 기존 무기울기 방법들보다 우수한 성능을 달성함을 보여줍니다.

핵심

🕵️‍♂️ 배경: "상자 속의 보물 찾기" (역문제)

우리가 살면서 겪는 많은 문제는 **'상자 속의 보물 찾기'**와 비슷합니다.

• 상황: 상자에 보물 (원본 이미지나 데이터) 이 들어있는데, 우리는 상자를 직접 열 수 없습니다. 대신 상자를 흔들어서 나는 소리나 진동 (관측 데이터) 만 들을 수 있습니다.
• 목표: 이 소리만 듣고 상자 안에 정확히 어떤 보물이 들어있는지 추론하는 것입니다.
• 문제: 같은 소리라도 보물이 여러 가지일 수 있습니다. (예: "동글동글한 소리"는 공일 수도 있고, 구슬일 수도 있음). 그래서 정답을 찾기 매우 어렵습니다.

기존의 AI(확산 모델) 는 이 문제를 해결하기 위해 **"상자 옆에 있는 전문가 (미분 가능한 함수)"**의 도움을 받았습니다. 전문가가 "소리가 이쪽으로 갈수록 보물에 가까워요"라고 방향을 알려주면, AI 가 그 방향으로 걸어가는 것입니다.

하지만 현실에서는 전문가가 없거나, 전문가가 너무 바빠서 방향을 알려줄 수 없는 경우 (예: 복잡한 물리 시뮬레이션, 블랙홀 이미지 등) 가 많습니다. 이때는 AI 가 길을 잃고 헤매게 됩니다.

---

💡 새로운 해결책: CPS (제약된 입자 찾기)

이 논문은 전문가의 도움 없이도, 혼자서도 가장 좋은 답을 찾아내는 방법을 제안합니다. 이를 **'CPS'**라고 부릅니다.

1. 기존 방식의 한계: "일등만 뽑고 나머지는 버리기"

기존의 무작위 방법 (SCG 등) 은 다음과 같이 작동했습니다.

• 상황: 보물을 찾기 위해 100 명의 탐정 (입자) 을 보냅니다.
• 작동: 100 명이 각자 다른 길을 가보더니, "가장 보물과 비슷한 소리"를 낸 단 한 명만 남기고 나머지는 모두 "실패자"로 쳐서 버립니다.
• 문제: 버려진 99 명도 "아, 저쪽은 소리가 안 들리네"라는 귀중한 정보를 가지고 있었지만, 그 정보를 무시하고 버려버린 것입니다. 비효율적입니다.

2. CPS 의 혁신: "모든 탐정의 정보를 합쳐서 지도 만들기"

CPS 는 완전히 다른 방식을 취합니다.

• 전략: 100 명의 탐정을 모두 보내고, 그들의 말을 모두 듣습니다.
  • "A 는 저쪽이 좋다고 했어."
  • "B 는 반대 방향이 나쁘다고 했어."
  • "C 는 약간 비틀거렸는데, 그걸로 보물 위치를 추정할 수 있겠어."
• 수학적인 비유 (국소 선형화): CPS 는 이 100 명의 탐정들이 남긴 정보를 모아, **"지금 이 순간의 보물 지도 (대리 모델)"**를 그립니다. 이 지도는 전문가가 그려준 지도는 아니지만, 탐정들의 집단 지성으로 만든 매우 정확한 지도입니다.
• 제약 조건 (Constrained): 그리고 이 지도를 볼 때, **"우리가 아는 보물 상자의 규칙 (확률 분포)"**을 잊지 않습니다. "보물은 보통 저기 있을 거야"라는 사전 지식을 바탕으로, 지도에서 가장 그럴듯한 한 곳을 선택합니다.

🌟 핵심 비유:

마치 100 명의 요리사가 각자 요리를 해보고 맛을 본 뒤, 그 맛 정보를 모두 모아 최고의 레시피를 한 번에 만들어내는 것과 같습니다. 기존 방식은 "가장 맛있는 요리사 한 명만 남기고 나머지는 쫓아내는" 방식이었다면, CPS 는 "모두의 맛을 분석해 더 완벽한 요리를 만드는" 방식입니다.

---

🔄 추가 기술: "다시 시작하기 (Restart)"

혹시 처음에 보낸 탐정들이 엉뚱한 방향 (나쁜 초기 잡음) 으로 출발했다면 어떨까요?
CPS 는 "다시 시작하기" 전략을 사용합니다.

• "아, 지금 방향이 틀린 것 같아. 잠시 멈추고, 조금 뒤로 물러서서 다시 출발해 보자."
• 이렇게 중간에 방향을 수정하며 누적된 실수를 고쳐주니, 최종 결과물이 훨씬 정확해집니다.

---

🏆 결과: 왜 이것이 중요한가요?

이 논문은 CPS 가 다음과 같은 분야에서 기존 방법들을 압도한다고 말합니다.

1. 이미지 복원: 흐릿하거나 찢어진 사진을 선명하게 만드는 작업에서, 전문가 (미분 정보) 가 없어도 기존 방법보다 훨씬 좋은 결과를 냅니다.
2. 과학적 발견:
   • 블랙홀 이미지: 우주에서 멀리 떨어진 블랙홀의 이미지를 복원할 때, 복잡한 물리 법칙 때문에 전문가의 도움을 받기 어렵습니다. CPS 는 이 어려운 문제도 잘 해결했습니다.
   • 유체 역학: 바람이나 물의 흐름을 예측할 때도 마찬가지입니다.

📝 한 줄 요약

"기존에는 '가장 잘한 사람 한 명'만 뽑아서 버리는 방식이었다면, CPS 는 '모든 사람의 의견'을 모아 가장 똑똑한 답을 찾아내는 방식입니다. 전문가의 도움 없이도, 집단 지성과 규칙을 이용해 복잡한 문제를 해결하는 새로운 길입니다."

이 방법은 앞으로 의료, 기상 예보, 우주 탐사 등 데이터가 부족하거나 복잡한 시뮬레이션이 필요한 분야에서 큰 역할을 할 것으로 기대됩니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

역문제 (Inverse Problem) 의 본질:
역문제는 관측된 데이터 $y$로부터 원본 신호 $x$를 복원하는 과정으로, $y = H(x) + \eta$와 같이 모델링됩니다. 여기서 $H(\cdot)$는 관측 연산자, $\eta$는 잡음입니다. 이 문제는 $m < d$ (관측값의 차원이 원본보다 작음) 인 경우 본질적으로 잘못 설정된 (ill-posed) 문제이며, 여러 가지 해가 존재할 수 있습니다.

기존 방법의 한계:

• 기울기 기반 방법 (Gradient-based): 확산 모델 (Diffusion Models) 을 사전 지식 (Prior) 으로 활용하여 역문제를 해결하는 기존 방법들 (예: DPS, RED-diff 등) 은 관측 과정의 기울기 (Gradient) 정보를 필요로 합니다.
• 실제 적용의 어려움: 많은 실제 시나리오 (예: 유체 역학 시뮬레이션, 블랙홀 영상화 등) 에서 관측 모델 $H(\cdot)$은 비선형적이거나 계산 비용이 매우 높으며, 미분이 불가능하거나 (Black-box) 기울기 계산이 극도로 어려운 경우가 많습니다.
• 기울기 없는 방법 (Gradient-free) 의 비효율성: 기울기 정보가 없을 때 사용하는 기존 방법들 (예: SCG, EnKG) 은 다음과 같은 문제를 가집니다.
  • SCG (Symbolic Constrained Guidance): 매 단계에서 여러 후보 입자 (Particles) 를 생성한 후, 관측값과 가장 잘 맞는 하나만 선택하고 나머지는 폐기합니다. 이는 많은 계산 자원을 낭비하며, 폐기된 입자들도 유용한 정보를 담고 있을 수 있다는 점을 간과합니다.
  • EnKG (Ensemble Kalman Guidance): 수천 개의 입자를 유지해야 하므로 고차원 문제에서 계산 오버헤드가 큽니다.
  • 성능 저하: 기울기 기반 방법에 비해 복원 품질이 현저히 낮습니다.

2. 제안 방법: Constrained Particle Seeking (CPS)

저자들은 기울기 (Gradient) 없이 오직 순전파 (Forward Pass) 만을 사용하여 역문제를 해결하는 새로운 프레임워크인 CPS를 제안합니다.

핵심 아이디어:
기존의 '수동적 선택 (Passive Selection)' 방식에서 벗어나, 모든 후보 입자의 정보를 활용하여 능동적으로 최적의 입자를 탐색 (Active Seeking) 하는 방식입니다. 이를 위해 역문제를 제약 최적화 (Constrained Optimization) 문제로 재정의합니다.

세부 알고리즘:

1. 관측 과정의 국소 대리 모델 (Local Surrogate) 구축:

   • 각 시간 단계 $t$에서 무조건부 전이 커널 $p(x_t|x_{t+1})$로부터 $n$개의 후보 입자 $x_t^1, \dots, x_t^n$을 샘플링합니다.
   • 각 입자에 대해 $H(\hat{x}_0|t)$ (Tweedie's formula 등을 통해 추정된 원본) 를 계산합니다.
   • 이 샘플들을 사용하여 관측 함수 $H$를 선형 대리 모델 (Linear Surrogate) $H(x) \approx Ax + b$로 근사합니다. 여기서 $A$와 $b$는 최소제곱법으로 추정되며, 이는 모든 후보 입자의 정보를 통합한 가중 선형 조합입니다.

2. 제약 최적화 문제 설정:

   • 목적 함수: 관측값 $y$와 대리 모델 예측값 $(Ax_t + b)$ 간의 오차 최소화 ($\min \|y - (Ax_t + b)\|^2$).
   • 제약 조건: 최적화 변수 $x_t$는 무조건부 사전 분포의 고밀도 영역 (High-density region) 내에 있어야 합니다. 확산 모델의 특성상 고차원 가우시안 분포의 질량은 반경 $\sigma_t\sqrt{d}$인 초구 (Hypersphere) $S^{d-1}$에 집중되어 있으므로, $x_t$가 이 초구 위에 있도록 제약합니다.
   • 수식:
     $$\min_{x_t} \|y - (Ax_t + b)\|^2 \quad \text{s.t.} \quad x_t \in S^{d-1}(\mu_t, \sigma_t\sqrt{d})$$

3. 해의 도출 (Asymptotic Solution):

   • 라그랑주 승수법을 적용하여 최적해를 구합니다. 확산 과정의 후반부에서 $\sigma_t$가 매우 작다는 점을 활용하여, 라그랑주 승수 $\lambda$가 매우 크다고 가정하고 점근적 해 (Asymptotic solution) 를 유도합니다.
   • 최종 해 $x_t^*$는 구의 중심 $\mu_t$에서 목적 함수의 기울기 방향 ($A^\top(y - \bar{H})$) 으로 고정된 거리만큼 이동한 위치로 계산됩니다. 이는 모든 후보 입자의 정보를 합성하여 생성된 최적의 단일 입자입니다.

4. 재시작 전략 (Restart Strategy):

   • 누적 오류를 보정하고 초기 잡음의 불확실성을 줄이기 위해, 샘플링 과정 중 특정 단계에서 생성된 입자를 다시 노이즈화 (Re-noise) 하여 이전 단계로 되돌리는 'Restart' 전략을 적용합니다. 이는 CPS 의 견고성을 크게 향상시킵니다.

3. 주요 기여 (Key Contributions)

1. 기울기 없는 최적화 프레임워크: 관측 모델의 기울기 정보가 전혀 없는 (Black-box) 상황에서도 고품질의 역문제 해결이 가능합니다.
2. 정보 활용의 패러다임 전환: 기존 방법들이 '나쁜' 입자를 폐기하는 방식에서 벗어나, 모든 후보 입자의 정보를 선형 대리 모델 구축에 활용하여 더 정확한 탐색 방향을 유도합니다.
3. 효율성과 정확성의 균형: EnKG 와 같이 수천 개의 입자를 유지할 필요 없이, 소수의 입자 (예: 64 개) 만으로도 기울기 기반 방법과 견줄 만한 성능을 달성합니다.
4. 광범위한 적용 가능성: 이미지 복원뿐만 아니라, 비선형성이 강한 과학적 역문제 (블랙홀 영상화, 유체 데이터 동화) 에서도 우수한 성능을 입증했습니다.

4. 실험 결과 (Results)

저자들은 FFHQ 데이터셋 기반의 이미지 역문제와 두 가지 과학적 역문제 (블랙홀 영상화, 유체 데이터 동화) 에서 CPS 를 평가했습니다.

• 이미지 역문제 (Inpainting, Super-resolution, Deblurring, JPEG Restoration):
  • CPS 는 기존 기울기 없는 방법 (SCG, DPG, EnKG) 보다 현저히 우수한 성능을 보였습니다.
  • 특히 비미분 가능한 JPEG 복원 작업에서 기존 방법들이 실패하거나 낮은 성능을 보인 반면, CPS 는 기울기 기반 방법 (DPS, RED-diff 등) 과 비슷하거나 더 나은 PSNR/SSIM/LPIPS 점수를 기록했습니다.
• 블랙홀 영상화 (Black Hole Imaging):
  • 매우 비선형적이고 희소한 관측 데이터를 다루는 이 작업에서 CPS 는 SCG 보다 높은 PSNR 을 기록하며 Ground Truth 와 가장 유사한 이미지를 복원했습니다.
• 유체 데이터 동화 (Fluid Data Assimilation):
  • Navier-Stokes 방정식을 푸는 수치 해석기 (Black-box) 를 사용하는 환경에서, CPS 는 희소하고 잡음이 많은 관측 데이터로부터 초기 유동장을 성공적으로 복원했습니다.
• 입자 효율성 (Particle Efficiency):
  • 입자 수를 8 개에서 128 개까지 변화시킨 실험에서, CPS 는 적은 수의 입자 (8 개) 만으로도 안정적이고 우수한 성능을 유지했습니다. 반면 EnKG 는 입자 수가 많아야 성능이 나갔고, SCG 는 이미지 문제에서 실패했습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 모델을 활용한 역문제 해결 분야에서 중요한 전환점을 제시합니다.

• 실용성: 실제 과학 및 공학 분야에서는 관측 모델의 기울기를 구하는 것이 불가능하거나 비용이 많이 드는 경우가 많습니다. CPS 는 이러한 실제적인 제약 조건 (Black-box, 비선형성) 하에서도 고품질 해를 제공할 수 있는 강력한 대안을 제시합니다.
• 계산 효율성: 불필요한 입자 폐기나 대규모 앙상블 유지 없이, 적은 수의 순전파 (Forward Pass) 만으로 최적의 해를 탐색함으로써 계산 자원을 효율적으로 사용합니다.
• 이론적 통찰: 역문제를 제약 최적화 문제로 재정의하고, 국소 선형화 (Local Linearization) 와 기하학적 제약 (초구 제약) 을 결합한 접근법은 향후 다른 생성 모델 기반 역문제 해결에도 영감을 줄 수 있습니다.

결론적으로, CPS 는 기울기 정보가 부재한 환경에서도 가장 효율적이고 강력한 역문제 해결사로 자리 잡을 잠재력을 가지고 있습니다.