기술적 요약: 정렬 및 역전 (Align & Invert)

문제 정의

역문제 (Inverse problems) 는 노이즈가 있거나 열화된 측정값 $y$로부터 깨끗한 신호 $x_0$를 복원하는 것을 목표로 합니다. 사전 훈련된 확산 (diffusion) 및 흐름 기반 (flow-based) 모델은 이러한 작업에 강력한 사전 지식 (prior) 으로 자리 잡았지만, 심각한 열화 (예: 극단적인 초해상도, 모션 블러) 가 발생하거나 복잡한 장면이 포함된 시나리오에서는 고도로 디테일한 이미지를 생성하는 데 어려움을 겪는 경우가 많습니다. 잠재 확산 모델 (Latent Diffusion Models, LDMs) 의 경우, 인코더 - 디코더 아키텍처의 비선형성이 역샘플링 과정에서 아티팩트를 유발할 수 있다는 특정 한계가 존재합니다. 또한, 기존 방법들은 고품질 재구성을 달성하기 위해 많은 수의 이산화 단계를 필요로 하며, 추론 과정에서 지각적으로 현실적인 해법으로 궤적을 유도할 수 있는 강력한 의미론적 제약이 부족합니다.

방법론

본 논문은 역문제 해결을 위한 확산 및 흐름 기반 솔버의 추론 과정에 **표현 정렬 (Representation Alignment, REPA)**을 통합하는 Align & Invert 프레임워크를 제안합니다.

핵심 개념

이 방법은 생성 모델 (확산/흐름) 의 내부 표현과 사전 훈련된 자기지도 학습 비전 인코더 (구체적으로 DINOv2) 의 특징 사이의 정렬을 강제합니다. REPA 는 이전에 생성 품질을 향상시키기 위한 훈련 정규화제로 사용되었으나, 본 연구는 정답 신호 (ground-truth) 가 없는 역문제 추론에 이를 적용합니다.

기술적 구현

1. 프록시 구성 (Proxy Construction): 정답 $x_0$를 알 수 없으므로, 본 방법은 누락된 목표 특징을 대신할 프록시 표현 $c_{proxy}$를 구성합니다.

   • 초기에 $c_{proxy}$는 사용 가능한 측정값 $y$의 DINOv2 특징에서 파생됩니다.
   • 역확산 과정이 진행됨에 따라 프록시는 모델의 현재 노이즈 제거 추정치 $f_{DINOv2}(D(E[z_0 | z_t]))$의 DINOv2 특징으로 점진적으로 업데이트됩니다.
   • 이 전략은 일반적인 이미지 열화에 대한 DINOv2 특징의 견고성에 의존합니다.

2. 기울어진 분포 (Tilted Distribution): 샘플링 과정은 정렬 항을 보상으로 포함하는 기울어진 분포에 의해 안내됩니다:
   $$\tilde{p}(x_t | y) \propto p(y | x_t) p_t(x_t) \exp(\lambda \cdot \text{REPA}(\bar{x}_0, h_t))$$
   여기서 $p(y | x_t)$는 측정 일관성을 강제하고, $p_t(x_t)$는 무조건부 사전 분포이며, REPA 항은 확산 모델의 중간 은닉 상태 $h_t$와 프록시 DINOv2 특징 사이의 정렬을 장려합니다.

3. 알고리즘 통합: 이 접근법은 REPA 목적 함수에 기반한 기울기 단계를 역샘플링 루프에 추가함으로써 기존 최첨단 솔버 (예: Latent DPS, ReSample) 에 통합됩니다. 이 단계는 잠재 확산 상태를 목표의 의미론적 특징과 정렬시킵니다.

4. 적응성: 원래 REPA 로 사전 훈련되지 않은 모델의 경우, 본 프레임워크는 확산 모델의 내부 표현을 DINOv2 특징 공간으로 매핑하는 경량의 사후 (post-hoc) 프로젝션 헤드 ($g_\phi$) 를 학습합니다. 이 매핑은 확산 모델을 고정된 상태로 유지하면서 최적화되므로, 적응 과정이 계산적으로 비용이 적게 듭니다.

이론적 통찰

본 논문은 REPA 정규화 항을 특징짓는 두 가지 주요 이론적 결과를 제공합니다:

1. 발산 최소화 (Divergence Minimization): REPA 목적 함수를 최대화하는 것은 특징 불일치 및 프록시 근사 오차를 고려한 항을 제외하고, 재구성된 이미지와 정답의 DINOv2 특징 분포 간의 최대 평균 불일치 (Maximum Mean Discrepancy, MMD) 를 최소화하는 것과 동등함이 입증되었습니다.
2. 잠재 상태 수축 (Latent State Contraction): 확산 인코더와 프로젝션 헤드의 야코비안 (Jacobian) 에 관한 정칙성 가정 하에서, 단일 REPA 업데이트는 확산 상태 공간에서 수축을 유도합니다. 구체적으로, 현재 상태가 해의 국소 근방 내에 있다면, 이는 현재 잠재 상태 $z_t$를 깨끗한 이미지의 잠재 표현 $z^*$에 더 가깝게 유도합니다.

실험 결과

저자들은 ImageNet 및 FFHQ 데이터셋을 사용하여 초해상도 (4×, 8×, 16×), 박스 인페인팅, 가우시안 디블러링, 모션 디블러링의 네 가지 역문제에 대해 이 방법을 평가했습니다.

• 지각적 품질 (Perceptual Quality): REPA 를 통합한 것은 모든 작업과 솔버 (Latent DPS, ReSample, Latent DAPS) 에서 지각적 지표 (LPIPS 및 FID) 를 일관되게 향상시켰습니다. 예를 들어, ImageNet 에서 4× 초해상도 작업 시 Latent DPS + REPA 는 LPIPS 를 0.238 에서 0.217 로, FID 를 123.82 에서 86.88 로 감소시켰습니다.
• 픽셀 지표 (Pixel Metrics): 표준 픽셀 단위 지표 (PSNR, SSIM) 는 베이스라인과 유사하게 유지되었으며, 이는 REPA 가 픽셀 단위의 정확도보다 의미론적 현실성을 우선시하는 잘 알려진 지각 - 왜곡 트레이드오프를 반영합니다.
• 효율성 (Efficiency): 중요한 발견은 REPA 가 솔버들이 훨씬 적은 이산화 단계로 동일한 수준의 지각적 품질에 도달할 수 있게 한다는 것입니다. 초해상도 및 가우시안 디블러링에서 이 방법은 베이스라인보다 약 4 배 적은 단계로 동등한 LPIPS 점수를 달성하여, 단계당 추가 계산이 있음에도 불구하고 전체 실행 시간을 단축시켰습니다.
• 견고성 (Robustness): 이 방법은 심각한 열화 (예: 16× 업스케일링, 고강도 블러) 하에서도 효과적이었으며, 학습된 프로젝션 헤드를 통해 독립적으로 사전 훈련된 모델에 적용되었을 때도 유효했습니다.

중요성 및 주장

본 논문은 표현 정렬이 추론 시 생성 역문제 솔버를 개선하는 효과적인 도구라고 주장합니다.

• 의미론적 안내 (Semantic Guidance): 내부 모델 표현을 견고한 비전 특징 (DINOv2) 과 정렬함으로써, 이 방법은 정답 감독 없이도 재구성을 지각적으로 의미 있는 해법으로 안내하는 강력한 의미론적 사전 지식을 제공합니다.
• 이론적 정당성 (Theoretical Justification): 이 작업은 정렬이 왜 작동하는지, 즉 특징 공간 발산을 최소화하고 잠재 상태를 깨끗한 이미지 해쪽으로 수축시키는 메커니즘으로 작용함을 규명합니다.
• 일반적 적용성 (General Applicability): 이 프레임워크는 일반적이며 다양한 확산/흐름 기반 솔버와 호환되며, 경량의 사후 매핑을 통해 원래 정렬로 훈련되지 않은 모델에도 적응 가능합니다.
• 효율성 (Efficiency): 이는 확산 샘플링의 계산 병목 현상을 해결하여 수렴을 가속화함으로써 더 적은 샘플링 단계로 고품질 재구성을 가능하게 합니다.

저자들은 이 접근법이 확산 모델의 생성 능력과 자기지도 학습 인코더의 의미론적 이해 사이의 간극을 메우며, 까다로운 역문제 설정에서 재구성 품질을 향상시키는 원칙적인 방법을 제공한다고 결론지었습니다.