기술적 요약: 확산 인코더

문제 제기

본 논문은 오토인코더 프레임워크 내에서 확산 모델을 인코더로 활용하는 과제를 다룹니다. 전통적인 변분 오토인코더 (VAE) 에서 인코더는 일반적으로 대각 공분산을 가진 단순한 가우시안 분포로 구성됩니다. 이러한 선택은 재매개변수화 기법을 통한 공동 훈련을 가능하게 하지만, 다음과 같은 중대한 한계를 부과합니다:

1. 제한된 표현력: 최적의 인코더 (실제 사후분포) 는 종종 비가우시안적이므로, 선택된 분포족이 최적 해를 포함하지 못할 수 있습니다.
2. 상각 격차 (Amortization Gap): 모든 입력에 공유되는 인코더 매개변수로 인해 모델이 특정 데이터 포인트에 대한 잠재 표현을 최적화하지 못하게 되어, 사후분포의 부Optimal한 근사로 이어집니다.
3. 공적응 (Co-adaptation): 디코더가 더 나은 사후분포 근사를 향해 인코더를 유도하기보다는 경직된 인코더 가족의 한계에 적응하게 되어, 결국 재구성 충실도가 저해됩니다.

확산 모델은 높은 표현력과 반복적 생성 과정을 통해 복잡한 상관관계를 모델링할 수 있는 능력을 제공합니다. 그러나 이들은 일반적으로 고정된 목표 분포를 위해 설계됩니다. 오토인코더에서는 인코더와 디코더가 최적의 표현을 협상함에 따라 잠재 분포가 훈련 중에 진화해야 합니다. 확산 모델에 재매개변수화 기법을 직접 적용하는 것은 잠재 변수와 네트워크 매개변수 간의 복잡한 함수적 의존성으로 인해 계산 비용이 많이 들고 불안정합니다. 또한, 인코더와 디코더를 독립적으로 업데이트하면 서로 반대 방향으로 이탈할 수 있는 동기화 문제가 발생합니다.

방법론

저자들은 기대값 - 최대화 (EM) 알고리즘에서 영감을 받은 교차 최적화 방식을 통해 확률적 인코더로 확산 모델을 활용하는 프레임워크를 제안합니다.

1. 확률적 인코더 공식화

직접적인 매핑 대신 인코더는 확률적 과정으로 모델링됩니다. 저자들은 확률 미분방정식 (SDE) 에 의해 지배되는 잠재 변수의 체인 $\{z_0, \dots, z_n\}$을 구성합니다. 목표는 최적 사후분포를 근사하는 분포 $q_\theta(z_T|x)$를 학습하는 것입니다.

2. 교차 훈련 방식 (알고리즘 1)

동기화 문제와 이동 표적 문제를 해결하기 위해 저자들은 인코더와 디코더 업데이트를 분리합니다:

• E-Step (사후분포로의 평형화):

  • 디코더 $p_\psi(x|z)$는 고정됩니다.
  • 최적 사후분포는 $q^*(z|x) \propto p(z) \exp(\frac{1}{\gamma} \log p_\psi(x|z))$로 정의됩니다.
  • 이 사후분포는 특정 SDE (랜지빈 동역학 과정) 의 평형 분포입니다.
  • 전체 생성 체인을 통해 역전파하는 대신, 현재 인코더의 출력에서 시작하여 짧은 랜지빈 체인을 실행함으로써 샘플 $z^*$를 생성합니다. 이 체인은 로그 가능도 $\nabla_z \log p_\psi(x|z)$의 기울기와 샘플을 사전분포 쪽으로 정규화하는 아핀 드리프트 항에 의해 안내되며 디코더 지형을 탐색합니다.
  • 이 단계는 인코더의 "목표"를 효과적으로 업데이트하여 잠재 샘플이 현재 디코더 상태와 일관되도록 보장합니다.

• M-Step (디코더 업데이트):

  • 인코더는 고정됩니다.
  • 디코더 $p_\psi(x|z)$는 E-Step 에서 생성된 샘플 $z^*$를 사용하여 재구성 손실을 최소화하도록 업데이트됩니다.

3. 확산 인코더 훈련

인코더 자체는 E-Step 에서 생성된 샘플 $z^*$의 분포와 일치하도록 훈련되는 확산 모델입니다.

• 목적 함수: 인코더는 경로 측도 간의 KL 발산을 최소화하는 표준 데노이징 목적 함수 (가중치 부여된 데노이징 스코어 매칭 손실 최소화) 를 사용하여 훈련됩니다.
• 아키텍처: 인코더 아키텍처는 시간 의존적 방식으로 입력 $x$를 쿼리하는 어텐션 메커니즘을 사용하며, 데노이징 과정이 진화함에 따라 $x$로부터의 정보를 잠재 변수 $z_t$로 점진적으로 주입합니다.
• 정규화: 정규화 압력은 인코더에서 E-Step 의 랜지빈 드리프트 항으로 이동됩니다. 이는 인코더가 디코더와의 정렬을 처리하는 랜지빈 동역학에 맡기면서, 안정적이고 표준적인 목적 함수로 훈련되는 "수동적" 역할에 머무를 수 있게 합니다.

주요 기여

1. 인코더로서의 확산: 본 논문은 확산 모델이 가우시안 인코더보다 더 표현력 있는 대안으로 성공적으로 인코더로 기능할 수 있음을 입증합니다.
2. 교차 훈련 알고리즘: 주요 기여는 알고리즘 1 로, 다음을 교차하는 안정적인 훈련 루틴입니다:
   • 짧은 랜지빈 체인을 통한 사후분포 샘플 생성 (E-Step).
   • 디코더 업데이트 (M-Step).
   • 표준 데노이징 목적 함수를 사용하여 이러한 샘플과 일치하도록 확산 인코더 업데이트.
3. 동기화 메커니즘: 랜지빈 체인을 사용하여 현재 디코더 지형에 잠재 변수를 평형화함으로써, 독립적 업데이트 중 인코더와 디코더가 정렬에서 벗어나려는 경향을 억제합니다.
4. 정보 이론적 연결: 프레임워크는 "신경 엔트로피 (Neural Entropy)" 개념을 활용하여 속도 - 왜곡 트레이드오프에서 속도 (상호 정보) 를 상한으로 제한하여 확산 인코더의 용량에 대한 이론적 근거를 제공합니다.

결과

저자들은 제안한 확산 오토인코더 (dAE) 를 MNIST, CIFAR-10, TinyImageNet, CelebA-HQ 에서 평가했습니다.

• 안정성: 훈련 과정은 안정적이고 매끄러운 것으로 관찰되었으며, 손실 프로파일은 EM 유사 방식과 일치했습니다. 저자들은 다양한 하이퍼파라미터 영역에서도 유의미한 훈련 불안정성이나 사후분포 붕괴가 보고되지 않았다고 밝혔습니다.
• 재구성: 모델은 일관된 재구성을 생성합니다 (그림 5-8 에 시각화됨).
• 속도 - 왜곡 성능: 속도 - 왜곡 (RD) 곡선에 플롯했을 때, dAE 는 표준 VAE 의 프런티어에 근접하지만 이를 능가하지는 않습니다. VAE 는 더 단순한 가우시안 인코더에도 불구하고 주어진 속도에서 더 낮은 재구성 오차를 달성합니다.
• 분석: 저자들은 dAE 의 약간 낮은 성능이 신경 엔트로피가 속도에 대한 상한일 뿐이며, 샘플로부터 사후분포를 모델링하는 잠재적 오류에서 비롯될 수 있다고 가설을 세웁니다. 또한, 토이 모델에서 확산 인코더는 실제 사후분포의 대칭성을 포착했지만 여전히 거의 가우시안 근사를 생성했다고 지적하여, 추가 조사가 필요한 가우시안 인코더의 잠재적 보편성을 시사합니다.

중요성과 주장

본 논문은 자신의 기여를 가우시안 인코더를 대체하는 최첨단 솔루션이 아닌, 확산 모델의 새로운 응용에 대한 개념 증명 (proof-of-concept) 으로 겸손하게 제시합니다.

• 동기: 이 작업은 오토인코더의 맥락에서 표현 능력과 훈련의 단순성을 위해 확산 모델의 사용을 동기화하는 것을 목표로 합니다.
• 강건성: 제안된 방식은 랜지빈 동역학을 통한 인코더 지연의 명시적 수용으로 인해 VAE 에서 흔히 발생하는 사후분포 붕괴에 대한 복원력을 보입니다.
• 향후 작업: 저자들은 랜지빈 체인의 순차적 오버헤드를 포함한 한계를 인정하며, 적응형 단계 크기 또는 임계 감쇠 동역학이 효율성을 개선할 수 있음을 제안합니다. 그들은 이 프레임워크를 확산 기반 인코더에 대한 추가 연구와 최대 엔트로피 강화 학습에서의 응용을 위한 비옥한 토양으로 간주합니다.