On the Redundancy of Timestep Embeddings in Diffusion Models

기술 요약: 확산 모델에서 타임스텝 임베딩의 중복성에 대하여

문제 제기
U-Net과 확산 트랜스포머(DiT) 아키텍처를 모두 포함하는 확산 모델은 근본적으로 다양한 노이즈 스케일에 걸쳐 디노이징 과정을 조절하기 위해 명시적인 타임스텝 임베딩에 의존한다. 이러한 시간적 신호는 일반적으로 포지셔널 인코딩이나 적응형 레이어 정규화(AdaIN)를 통해 네트워크에 현재 노이즈 수준을 알리기 위해 주입된다. 이러한 메커니즘이 표준적인 관행이기는 하지만, 본 논문은 이것이 절대적으로 필요한지에 대한 기존의 통념에 이의를 제기한다. 조사된 핵심 문제는 명시적인 시간적 신호가 디노이징 과정을 학습하는 데 필수적인지, 아니면 모델이 오염된 입력으로부터 노이즈 스케일을 직접 암시적으로 추론할 수 있어 임베딩이 중복되는 것인지 여부이다.

방법론
저자들은 이론적 분석과 엄격한 경험적 절제 연구(ablation studies)를 결합한 이중 접근 방식을 채택하였다:

1. 이론적 프레임워크:

   • 본 논문은 확산 학습 목적 함수의 전역 최솟값(global minimizer)을 분석한다. 저자들은 특정 조건 하에서, 시간 불변(time-agnostic) 손실 함수의 전역 최솟값이 표준적인 시간 조건부(time-conditioned) 손실의 전역 최솟값과 동일함을 입증한다.
   • 정리 1은 오염된 입력 $x_t$로부터 암시적 노이즈 스케일($\bar{\alpha}_t$)을 복구할 수 있는 측정 가능한 매핑이 존재한다면, 네트워크가 최적의 노이즈 예측을 수행하기 위해 명시적인 타임스텝 조건화가 필요하지 않음을 확립한다.
   • 정리 2는 데이터 성분의 경험적 2차 모멘트가 상수 $C \neq 1$ 주변으로 집중될 때 이러한 매핑이 존재함을 증명하는 고차원 분석을 제공한다. 고차원 공간에서 오염된 샘플의 노름(norm)은 노이즈 스케일의 결정론적 함수 주변으로 집중되며, 이를 통해 복구가 가능하다. 저자들은 이 조건이 데이터 공간(예: CIFAR-10)에서는 자연스럽게 충족되지만, 정규화된 잠재 공간(예: VAE로 인코딩된 CelebA)에서는 $C \approx 1$이 되어 약화될 수 있음을 언급한다.
   • 저자들은 또한 "지역적 식별 가능성(Local Identifiability)"을 논의하며, 컨볼루션 아키텍처(U-Net)가 겹치지 않는 어텐션 그리드(DiT)에 비해 정규화된 공간에서의 국소적 분산 변동을 더 잘 활용할 수 있다고 제안한다. 이는 U-Net의 밀집된 슬라이딩 윈도우 수용 영역 덕분이다.

2. 경험적 절제 연구:

   • 저자들은 U-Net과 DiT 아키텍처 모두에서 명시적인 타임스텝 조건화를 체계적으로 제거하였다.
   • 실험은 CelebA 데이터셋(128×128 해상도, VAE 잠재 공간에서 학습)과 CIFAR-10 데이터셋(32×32 해상도, 데이터 공간에서 학습)에 대해 수행되었다.
   • 모델은 Fréchet Inception Distance (FID), 정밀도(Precision), 재현율(Recall) 지표와 함께 정성적인 시각적 검사를 통해 평가되었다.

주요 기여

1. 엄격한 절제 연구: 본 연구는 컨볼루션(U-Net) 및 트랜스포머 기반(DiT) 백본 모두에 대해 시간 불변 변형 모델과 표준 조건부 모델을 포괄적으로 비교한 첫 번째 연구를 제공한다.
2. 중복성에 대한 이론적 증명: 고차원 집중 가설 하에서 노이즈 스케일이 오염된 입력으로부터 점근적으로 식별 가능하다는 것을 증명함으로써, 타임스텝 임베딩이 기능적으로 중복될 수 있다는 이론적 근거를 제시한다.
3. 아키텍처적 통찰: 분석 결과, 타임스텝 임베딩의 필요성은 데이터 분포(잠재 공간 vs 데이터 공간)와 아키텍처의 귀납적 편향(밀집된 컨볼루션 vs 희소한 어텐션) 사이의 상호작용에 따라 달라질 수 있음을 밝혀냈다.

결과

• CIFAR-10 (데이터 공간): 타임스텝 임베딩을 제거했을 때 유의미한 성능 향상이 나타났다. DiT와 U-Net 모두 조건부 모델에 비해 실질적으로 낮은 FID 점수(예: U-Net는 52.67에서 41.18로 개선, DiT는 51.32에서 34.74로 개선)와 더 높은 정밀도 및 재현율을 달려냈다.
• CelebA (잠재 공간): 결과는 엇갈렸으나 대체로 견고했다. U-Net는 FID(69.82에서 63.63)와 재현율에서 주목할 만한 개선을 보였다. DiT는 FID가 약간 증가(70.04에서 70.48)했으나 대등한 성능을 유지했는데, 이는 잠재 공간 정규화($C \approx 1$)가 이론적 복구 메커니즘을 약화시키더라도 모델이 여전히 견고함을 시사한다.
• 효율성: 모든 경우에서 타임스텝 임베딩을 제거함으로써 시간적 조건화 메커니즘과 관련된 계산 오버헤드를 없애 추론 시간을 단축시켰다.
• 정성적 결과: 시간 불변 모델은 높은 구조적 충실도를 유지했다. 일부 경우, 이들은 얼굴 특징, 질감, 포즈에서 더 큰 다양성을 보였으며, 이는 개선된 다양성 지표와 상관관계가 있었다.

의의 및 주장
본 논문은 생성 모델링의 오랜 시간 조건화 패러다임에 도전한다고 주장한다. 저자들은 아키텍처가 입력 구조로부터 노이즈 스케일을 암시적으로 추론할 수 있으므로, 명시적인 타임스텝 임베딩이 효과적인 디노이징을 위해 반드시 필요한 것은 아니라고 상정한다. 이러한 발견은 샘플 품질을 희생하지 않으면서도, 때로는 오히려 향상시키면서, 계산 오버헤드를 줄일 수 있는 더 효율적이고 구조에 집중된 생성 아키텍처를 향한 경로를 제시한다.

저자들은 본 연구의 범위에 대해 겸허한 태도를 유지하며, 본 연구가 저~중해상도 데이터셋에 초점을 맞추고 있음을 명시한다. 이들은 이러한 결과의 고해상도 합성 및 더 복잡한 데이터 분포로의 확장 가능성이 아직 검증되어야 하며, 데이터 복잡성이 증가함에 따라 타임스텝 임베딩이 다시 필요해질 수 있는지 판단하기 위한 후속 연구가 필요하다고 명시적으로 밝혔다.