기술 요약: 이산 랑제빈 영감 사후 샘플링 (∆LPS)

문제 정의

본 논문은 이산 상태 공간에서의 역문제 해결이라는 과제를 다룹니다. 역문제는 간접적이고 희소하며 잡음이 포함된 측정치 $y$ (모델링: $y = A(x) + n$) 로부터 알려지지 않은 신호 $x$ 를 복원하는 것을 의미합니다. 연속 확산 모델은 사후 샘플링을 통해 이러한 문제를 해결하는 표준이 되었으나, 그 이산 대응 모델들은 아직 충분히 탐구되지 않았습니다.

기존 이산 사후 샘플링 방법들은 세 가지 주요 한계를 겪고 있습니다:

1. 연속 완화: G2D2 와 같은 방법은 그래디언트 안내를 사용하기 위해 범주형 변수를 연속 공간으로 완화합니다. 이는 추론 기하학을 변경하며, 투사된 이산 토큰이 일관되게 디코딩되지 않을 위험을 초래합니다.
2. 순차적 업데이트: SGDD 와 같은 방법은 깁스 스타일 업데이트에 의존하는데, 이는 순차적이며 고차원 공간에서 혼합 (mixing) 이 느립니다.
3. 전문화: APS 와 같은 방법은 특정 손상 과정 (예: 마스크된 확산) 과 학습 패러다임에 밀접하게 결합되어 있어 일반성이 제한됩니다.

목표는 완전히 이산적으로 유지되며, 모든 토큰 차원을 병렬로 업데이트하고, 이산 확산 사전 모델의 특정 학습 패러다임에 무관한 사후 샘플러를 개발하는 것입니다.

방법론: ∆LPS

저자들은 ∆LPS(Discrete Langevin-Inspired Posterior Sampler, 이산 랑제빈 영감 사후 샘플러) 를 제안합니다. 이는 이산 상태 공간을 벗어나지 않으면서 그래디언트 정보를 활용하여 이산 전이를 안내하는 학습 없는 알고리즘입니다.

핵심 공식화

문제는 인코더 - 디코더 모델의 이산 잠재 공간 $z_0 \in \{1, \dots, K\}^L$ 에서 공식화됩니다. 목표는 사전 $p_\theta$ 가 사전 학습된 이산 확산 사전일 때, 사후 $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$ 에서 샘플링하는 것입니다. 샘플러는 측정치와 사전에 일관된 깨끗한 잠재 $z_0$ 를 먼저 추론한 후, 이를 다음 시간 단계로 재잡음화함으로써 사후 전이 $p(z_s|z_t, y)$ 를 근사합니다.

이산 랑제빈 제안

핵심 혁신은 랑제빈 역학을 이산 공간에 적응시키는 것입니다. 연속 설정에서 랑제빈 업데이트는 로그 사후 잠재 $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$ 의 그래디언트 방향으로 샘플 $z_0$ 를 이동시킵니다. $z_0$ 가 이산적이므로 저자들은 그래디언트 벡터를 직접 추가하지 않습니다. 대신, 연속 랑제빈 역학의 가우시안 제안 분포를 사용하여 유효한 이산 후보 $z'_0$ 에 대한 전이 확률을 정의합니다:

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

지수를 전개하고 유효 토큰의 유한 집합에 대해 정규화하면 제안이 토큰 좌표 간에 분해됩니다. 이를 통해 샘플러는 각 토큰 위치 $\ell$ 에 대해 독립적으로 범주형 분포를 계산할 수 있습니다:
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

이 메커니즘은 깁스 샘플러의 순차적 병목 현상을 피하면서 모든 $L$ 개의 토큰 차원에 걸쳐 완전한 병렬 업데이트를 가능하게 합니다.

일반화 및 전조건부

• 사전 무관성: 이 방법은 범주형 상태에 대한 예측 확률만 필요로 하므로, 마스크된 확산 (MDLM) 과 균일 상태 확산 (Duo) 사전과 호환됩니다.
• 블라인드 역문제: 미분 가능하지 않거나 알려지지 않은 순방향 연산자의 경우, 가능도 그래디언트는 데이터와 측정에 대한 인코더를 사용하는 학습된 대조적 대리 모델로 대체되어 부드러운 안내 신호를 제공합니다.
• Adam 전조건부: 고차원 잠재 공간에서 poorly scaled 된 그래디언트를 처리하기 위해, 원시 그래디언트는 Adam 전조건부 방향으로 대체됩니다. 이는 업데이트를 적응적으로 재조정하여 노이즈가 많은 큰 그래디언트는 감쇠시키고 일관되게 유용한 그래디언트는 보존함으로써 혼합과 수렴을 개선합니다.

주요 기여

1. 완전한 이산 사후 샘플링: ∆LPS 는 모든 반복을 이산 상태 공간에 유지하는 학습 없는 샘플러로, 소프트 토큰 최적화와 투사 오류를 방지합니다.
2. 효율적인 병렬 업데이트: 제안이 좌표 간에 분해됨으로써, 이 방법은 모든 토큰 차원을 병렬로 업데이트하여 순차적 깁스 스타일 샘플러보다 효율성을 크게 향상시킵니다.
3. 사전 무관 설계: 샘플러는 특정 손상 과정 (마스크 대 균일) 에 묶이지 않으며, 예측 확률을 제공하는 모든 사전 학습된 이산 확산 사전과 작동합니다.

실험 결과

저자들은 MNIST, CIFAR-10, FFHQ 에 걸쳐 선형, 비선형, 블라인드 설정을 아우르는 다양한 역문제 세트를 대상으로 ∆LPS 를 평가했습니다.

• 이미지 복원 (FFHQ): 인페인팅, 가우시안 디블러링, 초해상도, HDR, 모션 디블러링을 포함한 작업에서 ∆LPS 는 연속 기반 모델 (예: DPS, PSLD) 과 동등하거나 더 나은 성능을 달성하며, 지각적 품질 (LPIPS) 측면에서 기존 이산 기반 모델 (SGDD, G2D2, APS) 을 능가합니다.
• 이산 작업 (MNIST/CIFAR): 이 방법은 다양한 난이도 (Easy, Medium, Hard) 에서 견고함을 입증합니다. 토큰 정확도와 PSNR 에서 SGDD 및 G2D2 를 일관되게 일치시키거나 능가하며, "전체" SGDD 변형보다 훨씬 낮은 계산 예산 (적은 함수 평가) 으로도 달성됩니다.
• 블라인드 역문제: 순방향 연산자가 알려지지 않은 보행 궤적로부터 평면도를 추론하는 공간 매핑 작업에서 ∆LPS 는 F1 점수와 교집합 비율 (IoU) 에서 연속 기반 모델 (CoGuide, DPS 변형) 을 능가합니다.

중요성과 주장

본 논문은 완전히 이산적이고 그래디언트 정보를 활용한 사후 샘플러가 이산 표현에 대한 역문제 해결을 위한 확장 가능하고 일반적인 경로를 제공한다고 주장합니다.

• 확장성: 병렬 업데이트 메커니즘은 고차원 공간에서 이산 샘플러의 혼합 속도 한계를 해결합니다.
• 일반성: 이 방법은 추론 알고리즘을 사전의 특정 학습 패러다임에서 분리하여 텍스트, 코드, 이미지, 오디오 등 성장하는 이산 확산 모델 생태계를 활용할 수 있게 합니다.
• 실용성: 연속 완화를 피함으로써 ∆LPS 는 최종 샘플이 유효한 이산 토큰임을 보장하며, 연속 근사가 일관된 이산 신호로 디코딩되지 않는 "투사 격차"를 방지합니다.

저자들은 이산 생성 모델이 확장됨에 따라 ∆LPS 와 같은 원칙적인 추론 방법이 생체 분자 설계, 상징적 음악, 코드 완성 및 데이터가 본질적으로 이산적인 기타 도메인에서 점점 더 중요해질 것이라고 결론지었습니다.