Timestep Rescheduling in Diffusion Inversion

기술 요약: 디퓨전 인버전에 있어서의 타임스텝 리스케줄링 (Timestep Rescheduling in Diffusion Inversion)

문제 정의
디퓨전 인버전(Diffusion inversion)은 실제 이미지 $z_0$를 디퓨전 모델의 가우시안 잠재 공간($z_T$)으로 매핑하는 과정으로, 이미지 재구성 및 편집과 같은 작업을 가능하게 한다. Denoising Diffusion Implicit Models (DDIM)가 결정론적이고 효율적인 정식화를 제공함에도 불구하고, 인버전 과정은 본질적으로 근사적이다. 표준적인 방식은 디노이징 단계를 분석적으로 역전시키는 것인데, 이 과정에서 발생하는 편차(deviation)가 시간이 흐름에 따라 누적되어 눈에 띄는 재구성 오류와 편집 충실도(fidelity) 저하를 초래한다.

기존 방법들은 인버전 방정식을 정교화하기 위해 고정점 문제(fixed-point problems)를 해결하려고 시도하지만, 이러한 접근 방식들은 타임스텝 스케줄링(디퓨전 타임스텝의 선택 및 간격)이 인버전 충실도에 미치는 영향을 간과하고 있다. 대부분의 파이프라인은 균등하게 간격이 배치된 타임스텝을 사용하는데, 저자들은 이것이 최적이 아니라고 주장한다. 논문은 타임스텝의 분포가 전역적 인버전 오류(global inversion error)의 누적에 중요한 영향을 미치지만, 디퓨전 샘플링이나 학습에서의 역할에 비해 이 관계에 대한 연구가 미흡하다고 지적한다.

방법론: TRDI (Timestep Rescheduling in Diffusion Inversion)
저자들은 추가적인 파라미터나 추론 단계에서의 계산 오버헤드 없이 전체 인버전 오류를 최소화하도록 설계된 비균등 타임스텝 스케줄러인 TRDI를 제안한다. 이 방법론은 인버전 오류에 대한 이론적 분석에 근거하며, 2단계 "코스-투-파인(coarse-to-fine)" 전략을 통해 구현된다.

1. 이론적 통찰 (Theoretical Insight):
   저자들은 타임스텝 점프 $\Delta t$에 의해 유도되는 인버전 오류가 **스케일링된 고정점 문제(scaled fixed-point problem)**로 재정의될 수 있음을 도출했다. 오류 크기 $\delta$는 스케일링 계수 $c_\alpha(t, \Delta t)$와 연속적인 모델 예측 간의 차이($\Delta \epsilon_\theta$)의 곱에 비례한다.

• 핵심 관찰: 스케일링 계수는 타임스텝 인덱스에 대해 **포물선 형태의 추세(parabolic trend)**를 보인다. 오류는 작은 타임스텝(초기 디퓨전 단계)과 큰 타임스텝(후기 단계) 양쪽 모두에 집중되며, 중간 지점에서 감소한다. 또한, 더 큰 스텝 사이즈($\Delta t$)는 본질적으로 더 큰 오류를 유발한다.

2. 전역적 리스케일링 (Global Rescaling - Coarse Adjustment):
   포물선형 오류 분포를 해결하기 위해, 저자들은 하이퍼파라미터 $\gamma$에 의해 제어되는 거듭제곱 법칙 변환(power-law transformation)을 사용하여 기존의 균등 타임스텝 시퀀스 $t_k$를 수정하는 전역 리스케일링 메커니즘을 도입한다.
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

• $\gamma > 1$일 때, 타임스텝은 초기 단계 쪽으로 늘어나며(오류가 높은 초기 영역에 더 많은 해상도를 할당),
• $\gamma < 1$일 때, 타임스텝은 마지막 단계 쪽으로 밀집된다.
  이 단계는 스텝의 할당을 전역적으로 조정함으로써 정밀한 최적화를 위한 준비를 마친다.

3. 동적 계획법을 통한 국소적 리스케줄링 (Local Rescheduling via Dynamic Programming - Fine Adjustment):
   누적 오류를 더욱 최소화하기 위해, 저자들은 윈도우 기반의 동적 계획법(DP) 접근 방식을 사용한다.

• 비용 맵 (Cost Map): $k$번째 인버전 단계에서 $t$번째 단계에 도달하기 위한 최소 누적 오류를 추적하기 위해 비용 맵 $E[k, t]$를 초기화한다.
• 재귀 (Recursion): 현재 단계 주변의 폭 $2d+1$인 슬라이딩 윈도우 내에서, 알고리즘은 이전의 누적 오류와 현재 단계의 오류 계수의 합을 최소화하는 인접 타임스텝 간의 국소적으로 최적화된 전이를 재귀적으로 탐색한다.
• 백트래킹 (Backtracking): 최종 단계에서 첫 번째 단계로 백트래킹하여 최적의 타임스텝 시퀀스를 추출한다.

전체 리스케줄링 과정은 알려진 노이즈 스케줄을 기반으로 오프라인(또는 전처리 단계)에서 수행되므로, 실제 이미지 인버전이나 편집 과정에 가해지는 런타임 오버헤드는 무시할 수 있는 수준이다.

주요 기여

• 이론적 정식화: 본 논문은 타임스텝 점프와 인버전 오류 사이의 이론적 연결 고리를 도출하여, 오류를 노이즈 스케줄에 의존하는 스케일링된 고정점 문제로 프레임화하였다.
• 적응형 스케줄러: 오류가 낮은 영역에는 큰 스텝을, 오류가 높은 영역에는 미세한 스텝을 적응적으로 할당하여 효율성과 정확성의 균형을 맞추는 코스-투-파인 리스케줄링 프레임워크(TRDI)를 제안한다.
• 플러그 앤 플레이(Plug-and-Play) 강화: 이 방법은 추가적인 파라미터, 학습, 또는 상당한 계산 오버헤스를 요구하지 않는 "기성(off-the-shelf)" 강화 기능으로 설계되어 기존의 인버전 파이프라인과 호환된다.

실험 결과
저자들은 다양한 디퓨전 모델(Stable Diffusion v1.5, SDXL, SDXL Turbo)과 베이스라인 인버전 방법(DDIM, ReNoise, NPI, GNRI)을 사용하여 이미지 재구성(MSCOCO 데이터셋) 및 이미지 편집(PIE-Bench)에 대해 TRDI를 검증하였다.

• 재구성 (Reconstruction): SD v1.5에서 TRDI를 통합했을 때 모든 베이스라인에서 일관되게 지표가 개선되었다. 예를 들어, DDIM 베이스라인의 경우 PSNR은 20.07에서 20.21로 증가했고, LPIPS는 193.97에서 187.85로 감소하여(지각적 품질 향상) 나타났다. 정성적 결과에서는 베이스라인이 정확하게 재구성하지 못한 기차 창문이나 과일 모양과 같은 구조적 세부 사항의 교정이 확인되었다.
• 편집 (Editing): PIE-Bench에서 TRDI는 높은 CLIP 유사도를 유지하면서 구조 보존 및 배경 보존 성능을 크게 향상시켰다. 특히, 배경 퇴화가 일반적으로 심각하게 발생하는 SDXL Turbo(4 steps)와 같은 적은 단계의 모델에서, TRDI는 배경 보존 성능(예: DDIM의 Structure Distance를 17.43% 감소)을 실질적으로 개선했다.
• 어블레이션 연구 (Ablation Studies): 연구를 통해 스케일링 인자 $\gamma \approx 1.05$와 윈도우 폭 $d=8$이 구조적 충실도와 의미적 편집 가능성 사이의 최적의 절충안을 제공함을 확인하였다.

의의 및 주장
본 논문은 타임스텝 스케줄링이 디퓨전 인버전 충실도에 있어 결정적이지만 이전에 간과되었던 요소임을 주장한다. 타임스텝 크기와 인덱스, 그리고 오류 누적 사이의 포물선 관계를 밝힘으로써, 저자들은 전략적이고 비균등한 계산 단계의 할당이 전역적 인버전 오류를 크게 줄일 수 있음을 입증하였다.

저자들은 TRDI를 더 강력한 인버전 정식화(예: 고정점 반복 또는 확률적 인버전)의 대체물이 아니라, **보완적인 경량 스케줄러 레벨의 개조(retrofit)**로 위치시킨다. 이는 기존의 결정론적 인버전 파이프라인이 계산 비용을 높이거나 모델을 재학습시키지 않고도 더 우수한 재구성 및 편집 결과를 얻을 수 있게 한다. 이 연구는 단계를 어떻게 밟느냐만큼이나 언제 밟느냐를 최적화하는 것이 중요하다는 점을 강조한다.