기술 요약: DiffRGD

문제 정의

확산 모델(Diffusion models)은 이미지 생성의 근간이 되었으나, 비용이 많이 드는 재학습이나 미세 조정 없이 특정 작업(예: 이미지 복원, 조건부 생성)을 위해 이를 제어하는 것은 여전히 과제로 남아 있습니다. 기존의 추론 시 가이던스 방법들(예: DPS, FreeDoM)은 미분 가능한 목적 함수로부터 얻은 그래디언트를 잠재 공간(latent space)에 직접 주입하여 샘플링 과정을 유도합니다. 그러나 이러한 방법들은 확산 과정의 기저에 깔린 가우시안 구조를 존중하지 않고 주변 공간(ambient space)에서 작동합니다.

이러한 분포 인식의 결여는 잠재 변수가 원래의 가우시안 매니폴드에서 벗어나게 되는 분포 드리프트(distributional drift) 현상을 초래합니다. 샘플링 단계가 누적됨에 따라, 이 "오프-매니폴드(off-manifold)" 현상은 최종 결과물인 깨끗한 샘플의 품질을 저하시킵니다. MPGD(국소 선형성 가정)나 DSG(고정 반경 초구체로 투영)와 같은 이전의 제약 시도들은 완벽한 인코더를 가정하거나, 가우시안 잠재 분포의 특정 특성을 완전히 준수하지 못하여 때때로 균등 분포로 퇴화하는 문제를 보였습니다.

방법론: DiffRGD

저자들은 가우시안 구조를 명시적으로 보존하는 플러그 앤 플레이(plug-and-play) 추론 시 가이던스 프레임워크인 DiffRGD를 제안합니다. 핵심 통찰은 각 DDIM 샘플링 타임스텝 $t$에서 잠재 분포가 예측된 평균 $\mu_t$를 중심으로 하는 등방성 가우시안(isotropic Gaussian)이라는 점입니다.

1. 극 분해(Polar Decomposition) 및 매니폴드 구축

이 방법은 등방성 가우시안 분포의 극 분해 성질을 활용합니다. 임의의 타임스텝 $t$에서 잠재 샘플 $x_t$는 통계적으로 독립적인 반지름 성분 $r$과 방향 성분 $u$로 분해될 수 있습니다:
$$x_t = \mu_t + \sigma_t r u$$
여기서 $r$은 스케일링된 카이 분포(chi-distribution)를 따르며, $u$는 단위 구 위에서 균등하게 분포합니다.

주변 공간이나 고정 반경 구체에서 최적화하는 대신, DiffRGD는 유도된 구형 매니폴드 $\mathcal{S}_{t,r}$를 구축합니다:
$$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$$
여기서 반지름 $\sigma_t r$은 분포의 자연스러운 스케일로부터 샘플링되어, 허용 집합(feasible set)이 확산 과정의 가우시안 특성을 존중하도록 보장합니다.

2. 리만 그래디언트 디센트 (Riemannian Gradient Descent, RGD)

각 타임스텝에서의 가이던스 문제는 다음과 같은 제약 최적화 문제로 정식화됩니다:
$$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$$
여기서 $\mathcal{L}$은 가이던스 손실(예: 복원을 위한 재구성 손실 또는 조건부 생성을 위한 특징 정렬)입니다.

이를 해결하기 위해 DiffRGD는 구형 매니폴드 $\mathcal{S}_{t,r}$ 상에서 **리만 그래디언트 디센트(Riemannian Gradient Descent)**를 수행합니다:

1. 접공간 투영(Tangent Space Projection): 손실의 유클리드 그래디언트를 현재 점에서의 구체의 접공간으로 투영하여, 업데이트 방향이 매니폴드에 접하도록 반지름 성분을 제거합니다.
2. 리트랙션(Retraction): 접공간을 따라 단계를 밟은 후, 리트랙션 연산을 통해 업데이트된 점을 구형 매니폴드 $\mathcal{S}_{t,r}$로 다시 매핑하여 반지름 $\sigma_t r$을 보존합니다.

이 과정은 각 확산 타임스텝당 $K$번의 내부 반복(통상적으로 $K=3$) 동안 반복됩니다. 저자들은 $L$-매끄러움(L-smoothness) 가정 하에 이 최적화가 정지점(stationary point)으로 수렴함을 증명합니다.

주요 기여

1. 분포 인식 기하학(Distribution-Aware Geometry): 저자들은 확산 가이던스를 위한 분포 유도 기하학을 밝혀내는 극 분해 성질을 식별했습니다. 이들은 잠재 분포의 기하학을 존중하는 리만 프레임워크를 제안하여, 샘플 품질을 저하시키는 오프-매니폴드 드리프트를 완화합니다.
2. 이론적 분석: 본 논문은 분포 보존에 대한 분석을 제공하며, 리만 최적화 알고리즘이 정지점으로 수렴함을 증명합니다. 실제 환경에서 제안된 기하학적 최적화는 ADMM 기반 방식보다 빠르게 수렴합니다.
3. 플러그 앤 플레이 성능: DiffRGD는 재학습이나 미세 조정이 필요 없는 플러그 앤 플레이 방식으로 작동하며, 다양한 작업에서 최첨단(state-of-the-art) 결과를 달성함이 입증되었습니다.

실험 결과

저자들은 HuggingFace Diffusers 프레임워크를 사용하여 단일 NVIDIA GeForce RTX 4090 GPU에서 DiffRGD를 DPS, FreeDoM, MPGD, DSG, ADMMDiff와 비교 평가했습니다.

이미지 복원 (Image Restoration)

FFHQ 256×256 및 ImageNet 256×256 데이터셋에 대해 네 가지 작업(인페인팅, 4배 초해상도, 가우시안 디블러링, 모션 디블러링)을 평가했습니다.

• 정량적 지표: DiffRGD는 대부분의 작업에서 모든 베이스라인 대비 PSNR, SSIM, LPIPS, FID에서 통계적으로 유의미한 개선을 달성했습니다. 예를 들어, 인페인팅 작업에서 DiffRGD는 두 번째로 우수한 방법인 ADMMDiff의 32.38 대비 34.04의 PSNR을 기록했습니다.
• 강건성: DPS 및 ADMMDiff와 같은 방법들이 눈에 띄는 성능 저하를 보인 것과 달리, DiffRGD는 적은 샘플링 단계(100 steps)에서도 강력한 성능을 유지했습니다.
• 정성적 지표: 시각적 결과물은 DPS 및 MPGD에 비해 아티팩트가 적고 배경 일관성이 뛰어났으며, ADMMDiff와 대등하거나 더 우수한 품질을 보이면서도 계산 시간은 적게 소요되었습니다.

조건부 생성 (Conditional Generation)

CelebA-HQ 256×256 데이터셋에 대해 세그멘테이션 맵, 스케치, FaceID 가이던스를 평가했습니다.

• DiffRGD는 정렬 지표(mIoU, Sketch-$\ell_2$, FaceID-$\ell_2$)와 생성 품질(FID, KID) 측면에서 기존 추론 시 방법들(FreeDoM, DSG, ADMMDiff)을 일관되게 능가했습니다.
• 정성적 결과는 DiffRGD가 DSG(종종 불안정한 제어를 보임)에 비해 조건에 대한 정렬이 더 뛰어나고 자연스러우며 시각적으로 매력적인 결과를 생성함을 보여주었습니다.

가이던스 강도 (Guidance Strength)

가이던스 강도($\eta_t$)에 대한 실험 결과, DiffRGD는 모든 설정에서 성능이 저하되는 DSG와 비교하여 더 넓은 범위의 강도에서 정렬(mIoU)과 생성 품질(FID) 사이의 더 나은 트레이드오프를 달성했습니다.

의의 및 주장

본 논문은 DiffRGD가 추론 시 가이던스의 핵심적인 간극, 즉 단계별 가우시안 잠재 구조를 보존하지 못하는 문제를 해결한다고 주장합니다. 분포 유도 구형 매니폴드 상에서 제약 최적화를 수행함으로써, DiffRGD는 샘플 품질을 저하시키는 "오프-매니폴드" 현상을 방지합니다.

저자들은 **기하학적 일관성을 갖춘 가이던스(geometry-consistent guidance)**가 주변 공간에서 직접 그래디언트를 적용하거나 제한적이고 분포를 인식하지 못하는 매니폴드를 사용하는 것보다 유효하고 우월한 대안임을 단언합니다. 이 방법은 값비싼 재학습 없이도 사전 학습된 확산 모델 고유의 고충실도 샘플 품질을 유지하면서 작업 지표를 개선하는 간단한 플러그 앤 플레이 솔루션을 제공합니다.