[논문 리뷰] Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models

Zijian Zhang, Zhou Zhao외 1명|arXiv (Cornell University)|2022. 12. 26.

Generative Adversarial Networks and Image Synthesis인용 수 15

한 줄 요약

이 논문은 전처리된 확산 확률 모델(DPM)을 미세조정 없이도 자동에코잉 프레임워크 내에서 비지도 표현 학습에 활용할 수 있도록 하는 PDAE를 제안한다. 이는 전방 과정에서 발생하는 정보 손실으로 인해 발생하는 사후 평균 간격을 메우기 위해 학습 가능한 평균 이동 예측기(learnable mean shift predictor)를 도입함으로써 이루어진다. PDAE는 Diff-AE보다 더 나은 재구성 품질과 훨씬 빠른 훈련 효율성을 달성하여, 훈련 시간을 절반 이하로 줄였음에도 불구하고 이미지 재구성 및 무조건적 생성 작업에서 더 뛰어난 성능을 보였다.

ABSTRACT

Diffusion Probabilistic Models (DPMs) have shown a powerful capacity of generating high-quality image samples. Recently, diffusion autoencoders (Diff-AE) have been proposed to explore DPMs for representation learning via autoencoding. Their key idea is to jointly train an encoder for discovering meaningful representations from images and a conditional DPM as the decoder for reconstructing images. Considering that training DPMs from scratch will take a long time and there have existed numerous pre-trained DPMs, we propose extbf{P}re-trained extbf{D}PM extbf{A}uto extbf{E}ncoding ( extbf{PDAE}), a general method to adapt existing pre-trained DPMs to the decoders for image reconstruction, with better training efficiency and performance than Diff-AE. Specifically, we find that the reason that pre-trained DPMs fail to reconstruct an image from its latent variables is due to the information loss of forward process, which causes a gap between their predicted posterior mean and the true one. From this perspective, the classifier-guided sampling method can be explained as computing an extra mean shift to fill the gap, reconstructing the lost class information in samples. These imply that the gap corresponds to the lost information of the image, and we can reconstruct the image by filling the gap. Drawing inspiration from this, we employ a trainable model to predict a mean shift according to encoded representation and train it to fill as much gap as possible, in this way, the encoder is forced to learn as much information as possible from images to help the filling. By reusing a part of network of pre-trained DPMs and redesigning the weighting scheme of diffusion loss, PDAE can learn meaningful representations from images efficiently. Extensive experiments demonstrate the effectiveness, efficiency and flexibility of PDAE.

연구 동기 및 목표

• 기존에 미리 훈련된 DPM을 활용하여 확산 오토에코더(Diff-AE)를 다시 훈련하는 데 드는 비효율성과 열악한 성능 문제를 해결하기 위해.
• 전방 과정에서의 정보 손실로 인해 미리 훈련된 DPM이 잠재 변수로부터 이미지를 재구성하지 못하는 실패 원인을 해결하기 위해.
• 사전 훈련된 DPM의 구성 요소를 재사용하고, 학습 가능한 평균 이동 예측기를 통해 간격 메우기 기법을 통해 표현 학습을 향상시키는 일반화 가능하고 효율적인 방법을 개발하기 위해.
• 최소한의 미세조정으로도 고품질의 이미지 재구성과 무조건적 생성을 가능하게 하기 위해, 사전 훈련된 DPM을 디코더로 활용할 수 있도록 하기 위해.

제안 방법

• PDAE는 역방향 확산 과정에서 예측된 사후 평균과 진짜 사후 평균 사이의 간격을 메우기 위해 학습 가능한 기울기 추정기(gradient estimator)를 도입한다.
• 모델은 잠재 표현에 담긴 정보량을 최대화하도록 인코더를 훈련시켜, 평균 이동 예측기가 원본 이미지를 효과적으로 재구성할 수 있도록 한다.
• 사후 평균 간격에 더 풍부한 정보가 포함된 시점에 더 무게를 두는 재설계된 손실 가중치 기법을 도입하여 표현 학습의 효율성을 향상시킨다.
• PDAE는 사전 훈련된 DPM의 U-Net 아키텍처 일부를 재사용하여 수렴 속도를 높이고 훈련 비용을 절감한다.
• 분류기 유도 샘플링과 사후 평균 간격 간의 연결 고리를 활용하여, 분류기의 기울기를 정보 손실을 복구하는 메커니즘으로 해석한다.
• 무조건적 샘플링 품질을 향상시키기 위해, 평균 이동 예측기를 샘플링 과정의 마지막 30퍼센트에만 적용하여 불확실성을 줄이고 도메인 외 샘플을 방지한다.

실험 결과

연구 질문

• RQ1 미리 훈련된 DPM을 자동에코잉 프레임워크 내에서 비지도 표현 학습에 효과적으로 재활용할 수 있는가?
• RQ2 왜 미리 훈련된 DPM은 잠재 변수로부터 이미지를 재구성하지 못하는가? 그리고 사후 평균 간격은 무엇에 기인하는가?
• RQ3 사후 평균 간격을 이용해 더 정보가 풍부한 표현을 학습하는 데 사용할 수 있는가?
• RQ4 DPM을 다시 훈련하지 않고도 손실된 정보를 효율적으로 재구성할 수 있도록 평균 이동 예측기를 어떻게 설계할 수 있는가?
• RQ5 사전 훈련된 DPM 구성 요소를 재사용하고 손실 가중치를 최적화하면 훈련 효율성과 성능이 향상되는가?

주요 결과

• PDAE는 Diff-AE보다 더 나은 이미지 재구성 품질을 달성하면서도 훈련 시간을 절반 이하로 줄여 뛰어난 훈련 효율성을 입증했다.
• FFHQ128, Horse128, Bedroom128, CelebA64 등 여러 데이터셋에서 FID 점수로 측정한 결과, 무조건적 샘플 품질이 크게 향상되었다.
• 샘플링 과정의 마지막 30퍼센트에만 평균 이동 예측기를 적용함으로써 불확실성이 감소하고 도메인 외 생성을 방지하여 샘플 품질이 향상되었다.
• 사후 평균 간격은 시간 단계에 따라 정보 농도가 다르게 나타나며, 이는 효율적인 표현 학습을 위한 재설계된 손실 가중치 기법의 타당성을 뒷받침한다.
• PDAE는 어떤 사전 훈련된 DPM에도 일반적인 부스터로 적용되어 DPM 재훈련 없이도 재구성 및 생성 품질을 향상시킬 수 있다.
• 사전 훈련된 DPM이 동결된 상태에서도 효과가 있음을 입증하여, 실세계 적용에 있어서의 호환성과 실용성을 입증했다.