기술 요약: 사전 학습된 확산 모델을 활용한 이미지 인페인팅을 위한 감가상각 안내 (AID)

문제 정의

이미지 인페인팅은 가시적인 픽셀과 일관성을 유지하면서 이미지의 누락된 영역을 재구성하는 것을 목표로 하는 조건부 생성 작업입니다. 사전 학습된 확산 모델이 이 능력을 크게 발전시켰지만, 기존 접근 방식은 효율성과 적응성 사이에서 이분법에 직면해 있습니다:

1. 작업별 학습: LaMa, MAT, BrushNet과 같은 방법들은 전용 모델이나 대규모 적응 모듈을 학습합니다. 이는 추론 비용을 감가상각하지만, 상당한 학습 시간이 필요하고 모델 복잡도를 크게 증가시킵니다.
2. 인스턴스별 최적화: RePaint, DPS, MCG와 같은 방법들은 사전 학습된 확산 백본을 고정된 채로 두되, 테스트 시 각 마스크 처리된 이미지에 대해 별도로 적응시킵니다. 이는 사전 학습된 모델을 보존하지만 배포 중 반복적인 최적화 비용을 발생시켜, 모든 이미지를 새로운 추론 문제로 취급합니다.

대규모 작업별 모델을 학습하거나 테스트 시 인스턴스별 최적화를 수행하지 않고도 추론 비용을 감가상각할 수 있는 "중도" 해결책에 대한 공백이 존재합니다.

방법론: 확산을 통한 감가상각 인페인팅 (AID)

저자들은 사전 학습된 확산 백본을 고정된 채로 두면서 작고 재사용 가능한 안내 모듈을 오프라인에서 학습하는 AID라는 프레임워크를 제안합니다. 이 모듈은 이후 테스트 시 추가 최적화 없이 모든 마스크 처리된 이미지에 걸쳐 재사용됩니다.

1. 제어 공식화

이 문제는 결정론적 제어 문제로 공식화됩니다.

• 동역학: 역확산 과정은 고정된 사전 학습된 스코어 함수 $\hat{S}$에 의해 주도되는 확률 흐름 ODE 에 의해 지배됩니다.
• 안내: 관측 가능한 마스크와 가시적인 픽셀을 나타내는 $\xi = (M, y)$인 작업 의존적 제어 항 $u(t, x; \xi)$가 드리프트에 추가됩니다. 이 제어는 숨겨진 깨끗한 이미지 $x^\dagger$에 의존하지 않습니다.
• 목적: 가시 픽셀과 누락 영역에 대한 충실도를 균형 있게 맞추는 말단 손실 $\Psi$를 최소화하고 안내장의 크기에 대한 페널티를 부과하는 실행 비용을 최소화하는 것이 목표입니다. 이는 최적 제어 $u^*$가 결정론적인 해밀턴 - 자코비 - 벨만 (HJB) 방정식으로 이어집니다.

2. 이론적 교량: 결정론적에서 확률론적으로

고차원 결정론적 HJB 방정식을 직접 푸는 것은 계산적으로 불가능합니다. 저자들은 연속 시간 강화 학습 (CT-RL) 을 통해 문제를 학습 가능하게 만들기 위해 가우시안 정책을 사용하는 보조 무작위화 제어 문제를 도입합니다.

• 보조 문제: 결정론적 제어 대신 시스템은 평균 $\mu$가 후보 안내 법칙인 가우시안 정책 $\pi^{(\lambda)} \sim \mathcal{N}(\mu, \sigma^2)$을 사용합니다.
• 최적화자 보존 교량: 이 논문은 두 가지 중요한 정리를 증명합니다:
  1. 충분성: 최적 결정론적 안내 $u^*$는 정책 평균을 $u^*$로 설정함으로써 최적 가우시안 정책으로 승격될 수 있습니다.
  2. 필요성: 보조 문제에 대한 임의의 최적 가우시안 정책은 그 평균이 최적 결정론적 안내 $u^*$를 정확히 복원해야 합니다.
• 의미: 이 교량은 배포 시 사용되는 목표 결정론적 안내를 변경하지 않으면서 (기울기 기반 업데이트를 가능하게 하는) 확률론적 정책을 학습할 수 있음을 보장합니다. 가우시안 정책의 분산은 고정되어 있으며 엔트로피 정규화제가 아닌 학습 장치로만 작용합니다.

3. 감가상각 액터 - 크리틱 알고리즘

이론적 교량에 기반하여 저자들은 완전히 데이터 기반인 액터 - 크리틱 알고리즘을 개발했습니다:

• 매개변수화: 액터 (안내 평균 $\hat{\mu}_\phi$) 와 크리터 (가치 함수 $\hat{V}_\theta$) 는 신경망으로 매개변수화됩니다.
• 학습: 알고리즘은 보조 HJB 방정식에서 유도된 마팅갈 직교성 조건을 사용합니다. 이는 작업 분포에서 샘플링된 궤적을 기반으로 한 확률적 근사를 사용하여 액터와 크리터를 업데이트합니다.
• 배포: 학습이 완료되면 안내 모듈은 결정론적으로 배포됩니다. 새로운 마스크 처리된 이미지의 경우, 시스템은 역시간 솔버 내에서 학습된 평균 안내 $\hat{\mu}_\phi$를 사용하고 정책 노이즈는 폐기합니다.

주요 기여

1. 방법론: AID 는 오프라인에서 경량 안내 모듈을 학습함으로써 추론을 감가상각하고, 인스턴스별 최적화 없이 작업 전반에 걸쳐 재사용할 수 있는 제어 안내 프레임워크를 도입합니다.
2. 이론: 이 논문은 결정론적 안내 문제와 보조 가우시안 공식화 사이의 엄격한 "정책 동등성 교량"을 확립합니다. 보조 문제가 최적화자 보존임을 증명하여, 결정론적 제어장을 학습하기 위해 확률론적 액터 - 크리틱 방법을 사용하는 것을 정당화합니다.
3. 효율성: 이 접근 방식은 사전 학습된 백본 대비 학습 가능한 매개변수를 1% 미만으로 추가하면서도 품질 - 속도 트레이드오프를 크게 개선합니다.

경험적 결과

저자들은 AFHQv2, FFHQ(픽셀 공간 EDM 파이프라인 사용), ImageNet(잠재 공간 EDM2 파이프라인 사용) 에서 AID 를 평가했습니다.

• 고정 백본 베이스라인 대비 성능: AID 는 여러 지표 (PSNR, SSIM, LPIPS) 와 마스크 유형 (자유 형태, 중앙, 스트립) 에서 강력한 고정 백본 베이스라인 (UnGuided, Replacement, MCG, DPS) 보다 일관되게 우수한 성능을 발휘했습니다.
• RePaint 대비 성능: AID 는 품질 면에서 강력한 반복 최적화 베이스라인인 RePaint 를 능가하면서 샘플링 예산의 약 10% 만 사용합니다 (예: RePaint 의 350 단계 대비 35 단계).
• LatentPaint 대비 성능: ImageNet 에서 AID 는 45% 미만의 학습 가능한 매개변수를 가진 안내 모듈을 사용하면서도 모든 지표에서 재사용 가능한 잠재 공간 적응 방법인 LatentPaint 를 능가했습니다.
• 지연 시간: AID 는 성능을 유지하면서 추론 시간을 더욱 단축하는 낮은 지연 시간 운영 지점 (12 단계) 을 지원하며, 이는 RePaint 가 아닌 베이스라인보다 우수한 성능을 보이고 RePaint 와도 경쟁력 있는 수준을 유지합니다.
• 일반화: 동일한 학습된 안내 모듈은 재학습 없이 보이지 않는 마스크 유형과 다른 샘플링 예산으로도 효과적으로 이전됩니다.

의미 및 주장

이 논문은 AID 가 다음과 같은 최초의 감가상각 제어 안내 프레임워크라고 주장합니다:

1. 사전 학습된 확산 백본을 고정 상태로 유지합니다.
2. 작고 재사용 가능한 안내 모듈만 학습합니다.
3. 배포 시 사용되는 결정론적 안내장으로의 정확한 최적화자 보존 교량을 확률론적 액터 - 크리틱 학습에서 제공합니다.

저자들은 이 작업을 제어 및 RL 기반 프레임워크를 사용하여 확산 모델을 연구하고 개선하기 위한 이론적으로 근거 있는 방법을 제공하는 단계로 위치시킵니다. 그들은 연속 시간 RL 이 엄격한 분석과 실용적인 알고리즘 설계 모두의 기반이 되며, 휴리스틱 인스턴스별 최적화나 무거운 작업별 재학습에 대한 원칙적인 대안을 제공한다고 강조합니다.