기술적 요약: 역문제 해결을 위한 일관성 정규화 그래디언트 흐름

1. 문제 정의

본 논문은 $y = A(x_0) + n$ 형태의 선형 역문제 해결 과제를 다룹니다. 여기서 $x_0$는 신호, $A$는 알려진 순방향 연산자, $n$은 가우시안 잡음입니다. 목표는 사후 분포 $p(x_0|y) \propto p(y|x_0)p(x_0)$를 사용하여 $y$로부터 $x_0$를 복원하기 위해 베이지안 추론을 수행하는 것입니다.

비전 - 언어 잠재 확산 모델 (LDMs) 은 강력한 생성적 사전 분포 $p(x_0)$를 제공하지만, 기존 LDM 기반 솔버는 두 가지 주요 한계에 직면해 있습니다:

1. 계산 비용: 이들은 일반적으로 많은 수의 신경 함수 평가 (NFEs) 와 대규모 사전 학습된 구성 요소 (특히 VAE 디코더) 를 통한 역전파를 필요로 하여 높은 메모리 사용량과 느린 추론을 초래합니다.
2. 프롬프트 민감성: 최첨단 방법들은 종종 생성을 조건 짓기 위해 텍스트 프롬프트에 의존합니다. 프롬프트가 장면과 완벽하게 정렬되지 않으면 재구성 품질이 저하됩니다. 프롬프트 최적화가 존재하지만, 현재 방법들 (예: LATINO-PRO) 은 샘플링과 최적화 단계를 번갈아 수행하여 수렴 속도가 느리고 NFE 수가 높습니다 (종종 50 이상).

2. 방법론

저자들은 잠재 일관성 모델 (LCM) 의 잠재 공간에서 사후 샘플링과 프롬프트 최적화를 동시에 수행하는 통합 프레임워크인 **일관성 정규화 와세르슈타인 그래디언트 흐름 (CWGF)**을 제안합니다.

2.1 이론적 프레임워크

이 방법은 역문제를 프롬프트 임베딩 $C$의 곱공간과 잠재 공간 $P_2(Z)$의 확률 측도 간의 최적화 문제로 공식화합니다. 목적 함수 $F[c, \mu]$는 다음을 결합합니다:

• 정규화 ($R$): 잠재 분포 $\mu$가 프롬프트 $c$에 조건부인 LCM 사전 분포 $p_c$와 정렬되도록 유도하는 쿨백 - 라이블러 (KL) 발산 $D_{KL}(\mu \| p_c)$.
• 가능도 ($L$): 샘플이 관측된 데이터와 일치하도록 보장하는 음의 로그 가능도 $-E_\mu[\log p(y|z)]$.

최적화는 연속 시간 그래디언트 흐름에 의해 주도됩니다:

• 유클리드 그래디언트: 주변 가능도 (MMLE) 를 최대화하기 위해 프롬프트 임베딩 $c$를 업데이트합니다.
• 와세르슈타인 -2 그래디언트: 목적 함수를 최소화하기 위해 확률 측도 $\mu$를 업데이트합니다.

2.2 이산화 및 구현

연속 흐름을 계산적으로 실행 가능하게 만들기 위해, 저자들은 리 - 트로터 분할과 특정 근사를 사용하여 이산 알고리즘을 유도합니다:

1. 사전 하흐 (잠재 공간):

   • 계산 불가능한 스코어 함수를 직접 계산하는 대신, 이 방법은 일관성 모델 (CM) 속성을 활용합니다.
   • "흐름 - 맵 불일치"를 사용하여 KL 발산의 그래디언트를 근사합니다. 이는 CM 이 깨끗한 잠재 $g_\theta(z_t, t, c)$를 예측한 값과 현재 입자 분포의 평균에 대한 커널 추정치를 비교하는 과정을 포함합니다.
   • 이 단계는 입자 붕괴를 방지하기 위해 입자들을 국소 바리센터에서 밀어내는 "반발" 항을 포함합니다.

2. 프롬프트 흐름 (임베딩 공간):

   • 프롬프트 $c$에 대한 그래디언트는 피셔의 항등식과 디노이징 스코어 매칭을 사용하여 유도됩니다.
   • 중요한 점은 이 그래디언트가 전체 역 SDE/ODE 나 VAE 디코더를 통해 미분할 필요 없이 LCM 의 스코어 출력을 프록시로 사용하여 추정된다는 것입니다.

3. 가능도 하흐 (데이터 일관성):

   • 메모리 집약적인 VAE 디코더를 통한 역전파를 피하기 위해, 저자들은 VAE 인코더 $q_\phi(z|x)$를 활용합니다.
   • 유도된 근사 (명제 2) 를 사용하여, 가능도 드리프트는 인코더의 출력과 잠재 및 관측에 따른 픽셀 공간의 가우시안 사후 분포를 사용하여 계산됩니다. 이를 통해 디코더를 미분하지 않고도 데이터 일관성 업데이트가 가능합니다.

알고리즘 요약:
알고리즘은 $K$회 (예: 16 단계) 반복됩니다. 각 반복에서 다음을 수행합니다:

1. 확산 시간 단계를 샘플링하고 입자에 잡음을 추가합니다.
2. KL 발산에 대한 경사 하강을 통해 프롬프트 임베딩 $c$를 업데이트합니다.
3. 분할 단계를 통해 잠재 입자 $\mu$를 업데이트합니다: 먼저 CM 흐름 맵을 사용하여 사전 분포 쪽으로 이동한 다음, 인코더 기반 드리프트를 사용하여 가능도 쪽으로 이동합니다.

3. 주요 기여

• 통합 그래디언트 흐름: 단일 흐름에서 텍스트 프롬프트와 사후 분포를 동시에 최적화하는 새로운 유클리드 - 와세르슈타인 그래디언트 흐름 프레임워크로, 샘플링과 최적화 루프를 번갈아 수행할 필요가 없습니다.
• 역전파 없는 추론: 사전 학습된 인코더를 활용하여 가능도 업데이트 중 오토인코더 (VAE) 구성 요소를 통한 역전파를 피함으로써 메모리 사용량을 크게 줄입니다.
• 효율성: 일관성 모델 (LCMs) 을 활용함으로써, 확산 기반 베이스라인이 수백 번 요구하는 것에 비해 매우 적은 NFE(특히 16 회) 로 고품질 재구성을 달성합니다.
• 강건성: 공동 최적화를 통해 추론 과정에서 정렬되지 않거나 "적대적인" 프롬프트를 수정하여 다른 방법들이 달성하지 못하는 의미론적 정렬을 회복할 수 있습니다.

4. 실험 결과

이 방법은 FFHQ-512 및 ImageNet-512 데이터셋을 사용하여 표준 선형 역문제 (가우시안 블러 제거, 모션 블러 제거, $\times8$ 초해상도) 에서 평가되었습니다.

• 성능: CWGF 는 FID, PSNR, LPIPS 측면에서 최첨단 (SOTA) 성능을 달성합니다.
  • FFHQ-512 가우시안 블러 제거에서 CWGF 는 단 16 NFE로 FID 21.95와 PSNR 29.42를 달성했습니다.
  • 이는 LATINO-PRO(65 NFE, FID 32.48) 와 P2L(400 NFE, FID 45.12) 보다 우수합니다.
• 프롬프트 강건성: 초기 프롬프트가 잘못 초기화된 경우 (예: 얼굴 재구성을 위한 "고양이 사진") 실험에서 CWGF 는 프롬프트를 최적화하여 올바른 의미론을 성공적으로 복원한 반면, 프롬프트 최적화가 없거나 고정된 프롬프트를 사용하는 방법들은 그럴듯한 고해상도 특징을 생성하지 못했습니다.
• 계산 효율성: CWGF 는 경쟁자들 (65~1000 범위) 에 비해 훨씬 적은 NFE(16 회) 를 요구하며, 디코더 역전파가 필요한 방법들 (예: LATINO-PRO) 보다 낮은 메모리 사용량을 보여줍니다.

5. 의의 및 주장

본 논문은 프롬프트 튜닝과 사후 샘플링을 통합함으로써 생성적 사전 분포를 이용한 역문제 해결에서 중요한 진전을 이루었다고 주장합니다. 주요 의의는 다음과 같습니다:

1. 계산 비용 절감: 기존 LDM 기반 솔버가 요구하는 계산 자원 (NFE 및 메모리) 의 일부로 SOTA 결과를 달성합니다.
2. 실용성: 값비싼 디코더 역전파를 제거함으로써 소비자 등급 GPU 에서 고품질 이미지 복원을 가능하게 합니다.
3. 이론적 엄밀성: 경험적 베이지안 추론과 거리 공간의 그래디언트 흐름 간의 간극을 메우는 (강한 로그 오목성 가정 하에) 증명 가능한 수렴 프레임워크를 제공합니다.

저자들은 프롬프트 임베딩에 대한 유클리드 기하학 가정은 그 내재적 구조를 완전히 포착하지 못할 수 있으며, 현재 방법은 일관성 모델 사전 분포로 제한되지만, 향후 작업으로 다른 few-step 모델로의 확장이 제안된다는 한계를 지적합니다.