다양하지만 얕지 않게

Abstract

 Denoising Diffusion Probabilistic Models (DDPMs)는 적대적 훈련 없이도 고품질의 이미지 생성을 달성했지만, 샘플을 생성하기 위해 많은 단계에서 마르코프 연쇄를 시뮬레이션해야 한다. 샘플링 속도를 높이기 위해 우리는 DDPMs와 동일한 훈련 절차를 가진 더 효율적인 반복적 암시적 확률 모델 클래스인 Denoising Diffusion Implicit Models (DDIMs)을 제안한다. DDPMs에서는 생성 과정이 특정 마르코프 확산 과정의 역방향으로 정의된다. 우리는 동일한 훈련 목표를 달성하는 비마르코프 확산 과정 클래스를 통해 DDPMs를 일반화한다. 이러한 비마르코프 과정은 결정론적인 생성 과정에 해당할 수 있으며, 이는 훨씬 더 빠르게 고품질 샘플을 생성하는 암시적 모델을 만든다. 우리는 실험적으로 DDIMs가 DDPMs에 비해 10배에서 50배 빠르게 고품질 샘플을 생성하고, 계산과 샘플 품질을 맞바꾸고, 잠재 공간에서 직접 의미 있는 이미지 보간을 수행하며, 매우 낮은 오류로 관찰값을 재구성할 수 있음을 보여준다.

1. Introduction

2. Background

 q(x_0)를 따르는 x_0에 대해 q(x_0)를 근사하는 p(x_0)를 학습하여 쉽게 샘플링하는 것에 관심이 있다. 아래 식은 DDPM에서 확인할 수 있음.

 파라미터 θ는 variational lower bound를 maximizing하면서 q(x_0)를 학습한다.

 VAE와 달리 DDPM은 fixend inference procedure q(x_{1:T}|x_0)를 사용하고 latent variables의 크기가 원본 이미지와 같아 상대적으로 차원이 높다. DDPM 논문에서 사용한 Markov chain with Gaussain transitions는 다음과 같음.

 위 과정을 forward process라 하고 p(x_{0:T})가 generative process로 intractable한 q(x_{t-1}|x_0)를 근사하여 samling 하는 process임. 직관적으로 forawrd process는 점진적으로 noise를 첨가하고 generative process는 점진적으로 noise를 제거한다.

 forward pprocess의 놀라운 성질은 x_0에서 x_t를 한번에 구할 수 있다.

 만약 α_T가 작다면 q(X_T|x_0)는 standard Gaussian for all x_0에 가까워져 p(x_T)를 N(0, I)로 잡는 것이 자연스럽다. 마만약 모든 p(x_{t-1}|x_t)가 학습되는 mean functions와 고정된 variances를 가지면 (2)에서 얻은 objectives는 다음의 식으로 간단해진다.

여기서 ε_θ는 ε _t를 예측하는데 쓰인 함수이고 γ는 α에 의존하는 positive 계수들이다. DDPM의 초기 논문에서는 이 γ를 1로 두었고 이 objective는 score matching을 기반으로 하는 noise conditional score networks의 objective와 동일하다. 이렇게 학습된 모델이 x_0는 x_T가 먼저 sampling 되고 x_{t-1}이 하나씩 sampling되면서 최종적으로 만들어진다.
 DDPM에서 forward process의 lengh를 결정하는 hyperparameter인 T는 매우 중요하다. Variational 관전에서 large T는 reverse process가 Gaussian에 가깝도록 만들어준다. 따라서 DDPM에서는 T를 큰 값인 1000을 사용하지만 이 때문에 sampling process가 다른 모델들에 비해서 느리다.

3. Variational Inference for non-Markovian forward process

 Generative model에서 reverse process를 inference process로 근사하기 때문에 이 과정에서 iterations의 숫자를 줄일 수 있는 방법을 생각해보았다. 우리의 중요한 관찰을 DDPM의 obejctive인 L_γ는 모든 latent variable의 joint인 q(x_{1:T}|x_0)에 직접적으로 의존하는 것이 아니라 오직 marginals인 q(x_t|x_0)에 의존한다는 점이다. 
 그런데 이 marginal 분포를 만드는 다양한 inference distribution이 있으므로 inference process(=sampling process)에 다른 non-Markovian 분포를 찾아볼 것이고 이를 통해 새로운 generative process를 유도해볼 것이다. 이 새로운 non-Markovian inference process는 똑같은 surrogate objective function을 만족하고 Gaussian case가 아닌 경우에도 non-Markovian 관점을 적용해볼 것이다.

3.1 Non-Markovian Forward Process

 Inference distribution인 Q를 생각해보면 

다음의 식이 성립한다. q(x_{t-1}|x_t, x_0)의 평균과 분산은 이 분포의 marginal인 q(x_{t-1}|x_0)가

위를 만족시키기 위해 만들어짐. 따라서 이 (7)에 해당하는 분포는 위에서 말한 요구사항이었던 marginals를 만족한다. 이 분포로부터 새로 만든 forward process는 베이즈 정리에 의해

이 되고 이 역시 가우시안이다.

이 분포와 달리 (8)의 분포는 더 이상 Markovian이 아니고 파라미터 σ는 forward process의 stochastic한 정도를 다룬다. σ가 0이면 전혀 노이즈가 추가되지 않는 경우인데 이 경우에 x_0와 x_t를 안다면 x_{t-1} 역시 fixed한 값으로 정해진다.

3.2 Generative Process and Unified Variationl Inference Objective

  다음은 trainable generative process인 p_θ(x_{0:T})인데 p_ θ^t(x_{t-1}|x_t)는 q(x_{t-1]|x_t,x_0)에 대한 지식으로 만들어진다. 직관적으로 noisy observation인 x_t가 주어졌을 때 먼저 대응되는 x_0를 예측하고 이를 사용해 x_t-1을 우리가 정의한 분포인 q(x_{t-1]|x_t,x_0)로 sampling하는 방식이다.
 먼저 q(x_0)를 따르는 x_0에 대해 x_t를 (4)를 통해 얻으면 다음과 같다.

역으로 x_t로부터 x_0를 예측해야하기 때문에 x_0로부터 더해진 noise인 ε을 예측해야한다.(ε_θ(x_t)가 그 역할) 따라서 아래 식인 denoised observation을 얻을 수 있다.

 그럼 우린 generative process를 정의할 수 있고 p(x_T) = N(0, I)를 따른다고 하면

가 된다.(DDPM과 달리 x_1에서 x_0를 생성하는 과정에도 Gaussian noise를 첨가함.)
우리는 아래의 variational inference objective를 minimize하는 θ를 구한다.

 J의 정의를 보면 마치 σ가 달라짐에 따라 여러 model을 훈련해야 하는 것처럼 보인다. 하지만 특정 가중치 γ에 대해 J_σ와 L_γ가 동일하다고 알려져있다.
 

 Variational objective L_γ는 모델의 파라미터 θ가 다른 t에 걸쳐 공유되지 않는다면 θ의 최적해가 weights인 γ에 의존하지 않는다는 점에서 특별하다.(합계에서 t에 의존하는 각 항을 별도로 최적화하여 전역 최적해를 달성하기 때문) 이러한 L의 성질은 두 가지 의미를 가진다. 첫째는 DDPM에서 variational lower bound로 L_1을 사용하는 것을 정당화하고 둘째는 J가 어떤 L과 동일하므로 J의 최적해는 L_1의 최적해와 동일하다는 것이다.
 따라서  모델의 파라미터 θ가 t에 걸쳐 공유되지 않는다면 L_1을 J의 surrogate objective로 사용할 수 있다.
(θ가 t에 걸쳐 공유된다는 것은 모든 시점 t에 대해 동일한 θ를 사용한다는 것을 의미)

4. Sampling from generalized generative process

 목적 함수로 L_1을 사용함으로써, 우리는 Markovian inference process뿐만 아니라, 우리가 설명한 σ로 매개변수화된 많은 non-Markovian forward process에 대한 생성 과정도 학습하고 있다. 따라서 사전 훈련된 DDPM을 새로운 objectives의 해결책으로 보아 σ를 변경하여 우리의 필요에 맞는 샘플을 생성하는데 더 나은 generative process를 찾는 것에 집중할 수 있다.

4.1 Denoising Diffusion Implicit Models

 Eq.(10)을 통해 x_{t-1}을 x_t로부터 sampling하면 다음과 같다.

 σ를 달리하면 generative process가 달라지지만 동일한 모델의 θ를 사용하므로 re-training은 필요하지 않다. 만약

라면 forward process는 Markovian이 되고 DDPM의 generative process가 된다.
 우리는 σ가 모든 t에 대해 0인 special한 상황에 주목하였는데 이 상황의 경우 forward process가 given x_{t-1} and x_0에 대해 deterministic해진다. 이렇게 얻어지는 모델은 implicit 모델(fixed procedure과 latent varaibles를 통해 samples가 만들어짐)인데 우리의 모델은 implicit model이면서 DDPM objective로 훈련되므로 denoising diffusion implicit model(DDIM)으로 명하였다. (그런데 이제 forward process는 더 이상 diffusion이 아님)

4.2 Accelerated Generation Process

 Generative process는 reverse process의 근사이다. 이 때 forward process가 T step을 거쳐야하기 때문에 generative process역시 T steps를 거쳐야했다. 그러나 denoising objective L_1은 q(x_t|x_0)가 고정되어 있는 한 특정한 forward procedure에 의존하지 않으므로 우리는T보다 짧은 길이의 forward procedure만 고려할 수 있었다. -> 이를 통해 generative process에서 다른 모델을 훈련할 필요없이 가속화할 수 있음.
  Forward process의 모든 잠제변수를 고려하는 것이 아니라 [1, ... , T]의 sub-sequence이면서 길이가 S인 τ에서 {x_ τ1, ..., x_ τS}를 생각하면 

q(x_ τi|x_rj, x_0)의 marginal이 q(x_ri|x_0) = N((α_ τi)**0.5*x_0, (1- α_ τi)I)가 되도록 잡아 sampling 과정의 반복되는 특성을 사용해 계산 효율성을 늘릴 수 있따.
 DDIM 역시 trainin에서는 section 3에서 언급한 대로 L_1을 목적 함수로 훈련할 수 있기 때문에 training은 변화가 필요하지 않다. 다만 생성시

에 약간의 변경을 추가하여 새로운, 더 빠른 생성 과정을 얻는다. 

5. Experiments

 이 section에서는 적은 iterations로 image generation을 수행할 때 DDIM이 DDPM보다 우수함을 보여주며 DDPM generation process에 비해 10배에서 100배의 속도 향상을 보인다고 한다. 게다가 DDPM과 달리 초기 잠재 변수인 x_T가 고정되면 DDIM은 generation trajectoryy에 관계없이 high-level image features를 유지하므로 latent space에서 직접적으로 보간을 수행할 수 있다. DDIM은 또한 latent code에서 재구성할 수 있기에 sample을 인코딩하는 데 사용할 수 있지만 DDPM은 확률적 샘플링 과정으로 이를 수행할 수 없다.
 
 각 데이터셋에 대해 동일한 훈련된 모델인 T=1000, 목적함수로는 L_1을 사용한다. Section 3에서 주장하는 바와 같이, 훈련 절차와 관련하여 변경 사항이 필요하지 않고 샘플을 생성하는 방식이 달라진다. 샘플이 얼마나 빨리 얻어지는지 제어하는 요소인 τ와 deterministic DDIM과 stochastic DDPM 사이의 비율을 결정하는 σ만 달라진다.
 
 우리는 [1, ...T]의 다양한 sub-sequences인 τ와 다양한 분산 하이퍼파라미터인 σ를 고려한다. 비교를 단순화 하기 위해 다음 형식의 σ를 고려한다.

 위 식의 η가 직접 제어할 수 있는 하이퍼파라미터로 η가 1일 때 원래의 DDPM을 의미하며 η가 0일 때 DDIM을 나타낸다.

5.1 Sample Quality and Efficiency 

 Table 1을 보면 샘플을 생성하는 데 사용된 timesteps 수(S)와 η를 달리하였다. 당연히 샘플의 품질은 S가 증가할 수록 높아졌으며 이는 샘플  품질과 계산 비용 간의 tradeoff를 나타낸다. S가 작을 때 DDIM이  최고의 샘플 품질을 달성하고 DDPM에 가까워질 수록 샘플 품질이 일반적으로 더 낮았다. 하지만 S가 클 경우 DDPM이 좋은 성능을 봉ㅆ다.
 
 Figure 3에서는 동일한 샘플링 단계 수와 다양한 η값을 가진 CIFAR10 및 CelebA 샘플을 보여주는데 DDPM의 경우 샘플링 경로가 10단계에 일 때 샘플 품질이 급격히 저하된다. η 값의 경우 생성된 이미지가 짧은 경로에서 더 많은 noisy perturbations를 갖는 것으로 판단된다. (이 때문에 FID 점수가 낮은 것으로 판단됨.)
 
 Figure 4에서는 샘플을 생성하는 데 필요한 시간이 샘플 경로의 길이에 따라 선형적으로 증가함을 보여줌. DDIM이 훨씬 적은 단계로 샘플을 생성할 수 있으므로 DDIM이 샘플을 더 효율적으로 생성할 수 있음을 시사한다. 

5.2 Sample Consistency in DDIMs

 DDIM의 경우 생성 과정은 결정론적이기 때문에 x_0는 오직 초기 상태인 x_T에 의존한다. Figure 5에서 보면 동일한 초기 x_T로 시작하면서 다른 generative trajectories(different τ)에서 생성된 이미지를 관찰했는데 흥미롭게도 동일한 초기 x_T로 생성된 이미지들은 생성 경로와 관계없이 대부분 고수준 특징들이 유사한다. 많은 경우에 단지 20단계로 생성된 샘플들이 1000단계로 생성된 샘플들과 고수준 특징 면에서 유사하며, 세부 사항에서만 약간의 차이를 보였다. 이는 x_T 자체가 이미지의 충분한 정보를 담고 있는 잠재 인코딩임을 나타내고 세부 사항은 생성 경로가 긴 부분에서 담당하는 것으로 보인다.

5.3 Interpolation in Deterministic Generative Process

 DDIM sample의 고수준 특징이 x_T에 의해 인코딩되므로 다른 implicit probabilistic models에서 관찰된 것과 유사한 semantic interpolation effect를 내는지 확인해보았다. DDPM에서는 동일한 x_T가 확률적 생성 과정으로 다양한 x_0를 생성할 수 있기 때문에 DDPM의 보간과는 다르다. 
 Figure 6에서 x_T에서의 간단한 보간이 두 샘플 사이에서 의미론적으로 유의미한 보간을 이끌어 낼 수 있음을 보여준다. 이를 통해 DDIM에서는 DDPM에서는 못하는 잠재 변수를 통해 생성된 이미지를 고수준에서 직접적으로 제어할 수 있다.

5.4 Reconstruction from Latent Space 

 블로그에선 생략하였지만 DDIM이 특정 ODE에 대한 오일러 적분 방식이기 때문에 x_0에서 x_T로 인코딩하고 x_T를 재구성할 수 있는 확인해보았다. CIFAR-10 모델로 CIFAR-10 test set에서 인코딩 및 디코딩을 S steps 수행해보았다. Error는 per-dimension mean squared error를 사용하였다. Table 2를 보면 우리의 결과는 DDIM이 더 큰 S 값에서 더 낮은 재구성 오류를 가지는데 이는 neural ODEs 및 normalizing flows와 유사한 특성을 가지고 있음을 보여준다. 동일한 것을 DDPM에 대해서는 수행할 수 없는데 이는 DDPM의 확률적 특성 때문이다.

Modeling Tabular Data using Conditonal GAN

Abstract

 표형 데이터의 확률 분포를 모델링하고 현실적인 합성 데이터를 생성하는 것은 어려운 작업이다. 표형 데이터는 이산형 및 연속형 열이 혼합되어 있으며, 연속형 열은 여러 모드를 가질 수 있고, 이산형 열은 종종 불균형하여 모델링을 어렵게 한다. 기존 통계 모델과 심층 신경망 모델은 이러한 데이터를 적절히 처리하지 못한다. 이를 해결하기 위해 우리는 조건부 생성기를 사용하는 CTGAN을 설계했다. 7개의 시뮬레이션 데이터셋과 8개의 실제 데이터셋을 포함한 벤치마크를 통해 CTGAN의 성능을 평가했으며, 대부분의 실제 데이터셋에서 베이지안 방법보다 우수한 성능을 보였다. 다른 딥러닝 방법들은 이러한 성과를 내지 못했다.

1. Introduction

2. Related Work

3. Challenges with GANs in Tabular Data Generation Task

이 섹션에서는 표형 데이터에 GAN을 사용하는 데 있어 고유한 도전 과제를 설명한다:

• 혼합 데이터 유형: 이산형 및 연속형 열을 동시에 생성해야 한다.
• 비정규 분포: 표형 데이터의 연속 값은 종종 비정규 분포를 가지며, 이는 정규화를 복잡하게 만든다.
• 다중 모드 분포: 연속형 열은 여러 모드를 가질 수 있다.
• 희소 원-핫 인코딩 벡터: 이산형 열이 희소 벡터로 표현될 때 문제가 발생한다.
• 불균형 범주형 열: 범주형 열의 높은 불균형은 모드 붕괴와 소수 범주의 불충분한 표현을 유발한다.

4. CTGAN Model

CTGAN 모델은 확인된 문제를 해결하기 위해 여러 기술을 도입한다:

• 모드 특정 정규화: 각 연속 값을 모드와 해당 모드 내의 정규화된 값의 조합으로 나타내어 비정규 및 다중 모드 분포를 처리하는 새로운 정규화 방법이다.
• 조건부 생성기 및 샘플링 학습: 생성기는 이산형 열 값을 조건으로 하여 모든 범주가 훈련 중에 잘 대표되도록 한다. 샘플링 학습 기법은 불균형한 범주형 열을 처리하는 데 도움을 준다.
• 네트워크 구조: 모델은 특정 활성화 함수와 정규화 기술을 사용하여 고품질 합성 행을 생성하는 완전 연결 네트워크를 사용한다.

+ TVAE Model

VAE를 변형하여 tabular data에 사용할 수 있게 함. 훈련할 때에는 ELBO loss를 사용함.

5. Benchmarking Synthetic Data Generation Algorithms

이 섹션에서는 CTGAN을 평가하기 위한 벤치마킹 프레임워크와 데이터셋에 대해 설명한다:

• 베이스라인 및 데이터셋: 벤치마크에는 베이지안 네트워크(CLBN, PrivBN)와 여러 심층 학습 접근 방식(MedGAN, VeeGAN, TableGAN)이 포함된다. 데이터셋에는 UCI 머신러닝 저장소와 Kaggle에서 가져온 7개의 시뮬레이션 데이터셋과 8개의 실제 데이터셋이 포함된다.
• 평가 메트릭 및 프레임워크: 평가에는 두 가지 메트릭이 사용된다: 시뮬레이션 데이터의 우도 적합성과 실제 데이터의 머신러닝 효능. 우도 적합성은 합성 데이터가 훈련 데이터와 동일한 분포를 따르는지 확인하고, 머신러닝 효능은 합성 데이터를 사용하여 훈련된 모델이 실제 테스트 데이터에서 얼마나 잘 수행하는지를 평가한다.

6. Conclusion

 이번에 리뷰할 논문은 diffusion model의 기초를 세운 DDPM을 다룬 논문인 'Denoising Diffusion Probabilistic Models'이다. 이 논문 이전에도 diffusion 아이디어는 있었지만 그 과정을 수식적으로 명확히 다루고 품질인 좋은 샘플들을 생성하였으며 학습 안전성을 높혔기 때문에 유명한 논문이다.(이 논문을 기점으로 diffusion model에 관한 관심이 늘어났다.) 읽어보면서 VAE와의 연관성을 생각해보면 p(x)를 다루기 어려워 latent variable을 통해 분포를 예상한다음 새로운 샘플을 생성해내는 것, reparametrization trick이 사용되는 점 등등이 있겠다. 논문은 아래 사이트에 기재되어있다.

https://arxiv.org/abs/2006.11239

Abstract

- 우리는 nonequilibrium thermodynamics에서 영감 얻은 diffusion probabilistic models(a class of latent variable models)를 사용하여 높은 품질의 이미지를 생성해냈다.

-> 간단히 찾아보니 nonequilibrium thermodynamics는 비평형 열역학으로 시스템이 평형 상태에 있지 않을 때의 동적 거동을 설명하는 이론이라고 한다. Diffusion의 diffusion process와 reverse diffusion process가 엔트로피의 변화와 유사한 원리를 바탕으로 작동한다고 한다.

-  우리의 가장 좋은 결과는 diffusion probabilistic models와 denoising socre matching with Langevin dynamics 사이의 새로운 연결에 따라 만들어진 weighted variational bound를 training함으로써 얻어냈다.

- 또한 우리의 모델은 progressive lossy decompression scheme을 허용하며, 이는 autoregressive decoding의 일반화로 해석될 수 있다.

-> Progressive lossy decompression scheme은 여러 단계에 걸쳐서 압축된 데이터를 복원하는 구조를 말하고(손실은 압축시 일부 정보가 사라질 수 있음을 의미, decompression은 압축된 데이터를 다시 원래 상태로 복원하는 과정) autoregressive decoding은 데이터를 순차적으로 예측하며 생성하는 방식이다. Autoregressive decoding의 핵심 아이디어를 포함하면서도, 더 넓은 범위의 문제를 해결할 수 있어 일반화라는 표현을 사용.

- unconditional CIFAR 10 dataset에서 우리는 9.46의 Inception 점수와 3.17의 sota FID 점수를 얻었다.

-> Inception 점수: 생성된 이미지가 얼마나 다양하고, 또 얼마나 실제 데이터와 유사한지 평가하는 지표, 높을수록 생성된 이미지가 다양하고 높은 품질을 가짐을 의미

-> FID 점수: 생성된 이미지와 실제 이미지 분포 차이를 측정하는 지표, 낮을수록 생성된 이미지가 실제 이미지와 매우 유사하며, 높은 품질을 가짐을 의미

-> unconditional은 클래스 레이블 정보를 사용하지 않고 이미지를 생성하는 것을 의미한다.

- 256x256 LSUN에서는 ProgressiveGan과 유사한 샘플 품질을 얻었다. 

1. Introduction

- 딥러닝을 사용하는 generative models는 다양한 data 종류에 대해서 좋은 퀄리티의 samples를 만들어냈다. 

- 이 논문은 diffusion probabilistic models의 progress를 제안한다. Diffusion probabilistic model은 매개변수화된 Markov chain으로 유한 시간내에 data에 대응되는 samples를 생성하기 위해 variational inference를 사용한다.

- 위 그림에서 q로 표현된 markov chain이 diffusion process로 noise를 추가하는 forward process임. p로 표현된 markov chain은 sampling process로 noise를 제거하는 denoising과정으로 reverse process임. DDPM은 p로 표현된 sampling process를 학습함.(q는 가정하기 때문에 학습하지 않음.)

- 만약 이렇게 추가하는 noise가 작은 양의 Gaussian noise라면 sampling chain의 transitions은 conditional Gaussian으로 표현하기에 충분함. -> 따라서, 간단히 신경망으로 평균과 분산만 예측하면 됨.

- Diffusion model는 정의하기 쉽고 훈련 효율도 높지만 지금까지 고품질 샘플을 생성할 수 있다는 증거가 없었다. 우리는 diffusio models도 좋은 품질의 samples를 만들 수 있음을 보였고 때로는 다른 유형의 생성 모델에서 발표된 결과보다 더 나은 성과를 보이기도 함.

- 추가로 이러한 매개변수화가 denoising score matching over multiple noise levels during traings와 annealed Langevin dynamics during sampling이 동일함을 보여줌.

-> denoising score matching: during traing: 데이터의 확률 밀도 함수를 추정하기 위한 방법으로, 데이터 포인트 주변의 확률 밀도의 기울기를 학습함. Denoising score matching에서는 노이즈가 추가된 데이터에 대해 스코어(확률 밀도 함수의 로그에 대한 기울기)를 추정하고 이를 통해 원래의 데이터 분포를 복원함.

-> annealed Langevin dynamics during sampling: 여러 단계에 걸쳐 노이즈를 점진적으로 줄여가며 데이터 샘플을 생성하는 방법. 각 단계는 앞 단계에서 학습된 스코어(확률 밀도 함수의 로그에 대한 기울기)를 이용하여 노이즈를 제거하고, 점차적으로 더 깨끗한 샘플을 만듦.

-> 두 방식 모두 데이터의 노이즈를 제거하고 원래의 데이터를 복원하는 것인데 본질적으로 같은 원리를 따을 논문에서 보인 것

- 우리의 모델은 다른 likelihood-based models에 비해 경쟁력 있는 log likelihoods를 가지지 못했는데 그 이유는 lossless codelength의 대부분이 사람이 인식할 수 없는  image의 세부적인 부분을 설명하는데 사용되고 있었기 때문이다. 이 현상을 손실 압축의 관점에서 분석해 개선하였다. 확산 모델의 sampling 과정은 progressive decoding의 한 유형으로 autoregressive decoding과 유사하지만, 더 일반화된 방식으로 작동하였다.

-> lossless codelength: 데이터를 압축할 때 손실 없이 필요한 비트 수

-> 손실 압축 관점은 데이터를 완벽히 복원하지 않아도 되지만, 중요한 정보는 남겨두는 방식이기 때문에 사람이 인식할 수 없는 세부사항을 생락하거나 덜 정확하게 표현하는 것. => 효율적이고 효과적인 모델

-> 그 결과 만들어진 것이 progressive decoding(이미지를 한 번에 새엉하는 대신 여러 단계에 걸쳐 점진적으로 만들어나가는 방식)이고 이는 데이터를 한 번에 하나씩 순차적으로 생성하는 방식인 autoregressive decoding보다 일반화된, 유연한 방식임.

2. Background

- 일단 가장 중요한 거

학습하려는 대상이 p가 아닌 q라는 점. VAE랑 또 다르다. 또한 왜 diffusion이 noising인지 단어의 의미로 해석해보면 데이터의 정보가 점점 퍼져서(확산되어서) 노이즈로 가득 차게 되는 것을 표현하기 위함이라 한다. 

- Diffusion models은 latent variable model로 p(x_0)는 다음과 같이 형성된다.

- 그런데 여기서 알아야 할 것이 이 x_1, ... x_t인 latents가 모두 x_0와 같은 차원이라는 것이다. -> VAE와 다른 점.

- x_0는 q(x_0)를 따른다.

- x_T는 Gaussian 분포를 따른다.

- reverse process: p(X_T)로 시작하는 학습되는 Gaussian trainsitions으로 이루어진 Markov chain

- forward process: 데이터에 가우시안 노이즈를 추가하는 과정으로 x_0에서 schedule된 베타들로 노이즈들이 추가된다.

 사실 베타들은 reparameterization으로 학습해도 되고 그냥 상수로 둬도 된다. 베타 값들은 크기가 작은데 그 이유가 베타가 작은 경우 p와 q가 유사한 함수 형태를 갖기 때문에 reverse process의 표현력이 증가하기 때문이다.  

 또한 forward process의 놀라운 성질은 x_t를 t에 관계 없이 x_0로부터 closed form으로 샘플링할 수 있다는 점이다. 

위 두 표현으로 q(x_t|x_0)를 다음과 같이 쓸 수 있다.

-> 위 과정이 왜 놀랍냐면 sampling process를 매우 단축시킬 수 있기 때문이다.(reparameterization trick을 쓰면 하나의 noise를 샘플링하면 끝임)

- training은 negative log likelihood의 variational bound를 optimizing하는 것이 목표이다.

- 효율적인 학습을 위해 L을 다시 써서 variance를 낮추면 아래와 같은 식이 된다.(왜 variane가 낮고 아래 식이 유도되는 지는 appendix B에 첨부하겠다.)

위 식에서 처음 보는 term이 가운데 term인데 이 term은 p(x_{t-1}|x_t)와 q(x_{t-1}|x_t, x_0) -> forward process posteriors를 비교하는 KL divergence term이다. 갑자기 forward process의 분포에 given x_0가 생겼는데 이는 위의 L의 variance를 낮추는 과정에서 생긴 것이다.(이렇게 초기 추정량에 조건부 기댓값을 추가하여 분산을 낮추는 방법을 Rao-Blackwellization이라 한다.) 이 분포는 아래와 같이 정의되고

모든 term들이 gaussian을 다루거나 gaussian 사이의 KL-divergence를 구하기 때문에 closed form을 갖는다.

3. Diffusion models and denoising autoencoders

- Diffusion models은 마치 제한된 형태의 잠재 변수 모델처럼 보일 수 있지만, 구현에 있어 많은 자유도를 허용함.

-> forward process의 분산인 베타

-> reverse process의 모델 architecture, gaussian distribution parameters

을 선택해야함.

3.1 Forward process and L_T

- 베타 역시 reparameterization으로 학습될 수 있지만 여기선 constants로 고정함. 따라서 q에서 학습될 파라미터가 없기 때문에 식 (5)에서 L_T는 무시함.(p(x_T)도 학습할 거 없으니까)

3.2 Reverse process and L_{1:T-1}

- 먼저 p(x_{t-1}|x_t)(1<t<=T)에서 평균과 분산을 어떻게 고를지 생각해보자.

분산은 학습되지 않지만 시간에 따라 다른 상수로 정하였는데 

시그마의 후보론

 가 있다. 첫번째 선택은 x_0가 정규분포 N(0,I)를 따를 때 최적이고 두 번째 선택은 x_0가 한 점으로 정해질 때 최적이다. 

-> 이미지로 생각하면 x_0의 각 픽셀이 독립적으로 평균 0, 분산 1을 갖는 정규 분포를 따를 때 최적(다양한 이미지가 균일하게 분포되어 있을 때), x_0의 각 픽셀이 고정된 값을 갖는 것을 의미한다.(모든 이미지가 거의 동일한 구조와 패턴을 가지는 경우) 

-> coordinatewise unit variance를 갖는 data의 reverse process entropy의 상한과 하한에 대응된다.

분산을 위와 같이 잡아 reverse process의 분포가 다음과 같이 변하면

L_{t-1}을 아래와 같이 쓸 수 있다.

이 식은 당연히 최소화해야 하기 때문에 가장 간단한 접근은 forward process posterior의 평균으로 뮤(x_t,t)를 잡는 것이다.

위 식에서 reparameterization을 통해 

에서 x_0를 얻고 

를 (8)에 대입하면 

를 얻게 된다.

이 식에서 알 수 있는건 forward process posterior의 평균이

으로 표현할 수 있어서 µ(x_t,t)를 아래와 같이 잡는 것이다.

e(x_t)는 아래 x_t의 e를 예측하기 위해 

모델이 학습하여 예측하려는 노이즈 추정량이다.

그럼 reverse process에 x_{t-1}을 sampling하는 것은

reparameterization trick으로

을 계산하는 것과 같다.

예측한 평균인 (11)을 (10)에 넣으면

가 된다.

-> 이 식이 denoising score matching과 비슷하고, Langevin-like reverse process의 variational bound와 같아 둘이 동듬함을 보였다고 한다.

- 전체 알고리즘은 다음과 같다.

-> Training 과정에서 t를 sampling하여 접근하는 방식은 계산의 효율성을 높이고 다양한 노이즈 수준을 균형 있게 학습할 수 있게 한다.

+ Diffusion 모델은 다양한 수준의 noise를 학습하기 위해 시간 term을 sampling하고 특정 시간 단계 t에서도 다양한 노이즈에 대해 학습하기 위해 e를 sampling함. 모델은 주어진 x_t와 시간 t를 입력 받아 해당 시점의 노이즈를 e를 예측하도록 학습됨. 이렇게 충분한 훈련이 되면 x_t와 t를 입력받아 노이즈를 예측하고 예측된 노이즈를 사용하여 현재 이미지에서 노이즈를 제거해 x_{t-1}을 계산한다. 이 과정을 무수히 반복해 원본 이미지와 같은 수준의 이미지를 생성해감.

- 요약하자면, rever process의 mean function approximator인 µθ를 예측하도록 학습할 수 있고, 또는 매개변수를 수정하여 ε을 예측하도록 학습할 수 있음.(x_0를 예측하게 할 수도 있으나 초기 실험에서 이것이 더 낮은 품질을 생성함.) ε-prediction parameterization이 Langevin dynamics와 denoising score matching과 유사한 것을 최적화 하는 것이 목표로 하는 단순한 diffusion model의 variational bound 둘과 유사함을 보였다. 그럼에도 불구하고 이것은 단지 pθ​(xt−1​∣xt​)의 또 다른 매개변수화일 뿐이므로, 4장에서 ϵ을 예측하는 것과 μt를 예측하는 것을 비교하는 실험을 통해 그 효과를 검증함.

-> ablation study: 모델의 특정 부분을 제거하거나 변경하여 성능을 비교하는 방법

3.3 Data scaling, reverse process decoder, and L_0

- 우리는 이미지 데이터를 픽셀 단위로 scaling해서 [0, 255]를 [-1,1]로 linearly scaling해줌.(시작이 표준 정규분포이니 신경망은 모든 데이터 포인트를 동일한 범위에서 처리하게 됨.)

- discrete log likelihoods를 얻기 위해 마지막 항을 독립된 이산 디코더로 설정함. 이는

이 가우시안 분포로부터 유도된 것으로

왜 이런 짓을 하냐면 애초에 이미지 데이터는 이산 데이터이기 때문. 연속적인 분포를 사용하여 이산 데이터의 로그 우도를 계산하려면 이렇게 이산 값에 따른 정해진 부분을 적분하여 이산값에 대한 확률을 계산해야 함.

-> 이렇게 하면 추가적인 노이즈를 넣거나 복잡한 수학 연산을 할 필요 없이 variational bound와 lossless codelength of discrete data가 같다고 알려져있다고 한.

-> 마지막 x_0를 구할 때는 노이즈를 넣지 않음.

3.4 Simplified training objective

- 3.1과 3.2, 3.3으로 구성한 variational bound는 세타에 대해 미분 가능해 training을 쓰일 수 있다. 그런데 여기서 기존 variational bound를 대신할 수 있는 간소화된 손실 함수를 사용할 것을 주장하는데 식은 다음과 같다.

t=1일 때 위 식은 L_0에 대응되는데 (13)에서 정의한 discrete decoder definition의 적분이 gaussian probability density funtion과 bin width를 곱한 것으로 근사된다.(시그마 제곱과 edge effects 무시)

t>1일 때 위 식은 L_1,...,L_t-1이 대응 되는데 (12) 가중치 없는 버전으로 NCSN denoising score matching과 유사하다 한다.

이 식은 표준 변분 경계와 비교했을 때 작은 t에 해당하는 손실 항목의 가중치를 낮추는 역할임. 따라서 네트워크가 더 큰 t값에서 더 어려운 노이즈 제거 작업에 집중할 수 있게 하고 이는 실제 샘플 품질을 높이는 것으로 이어졌다고 함.

4. Experiments

- 모든 실험에 T를 1000으로 맞춤

- 순방향 과정의 분산을 β_1 = 10^-4, β_T = 0.02로 선형 증가하는 상수로 설정함.

-> β 스케줄링 방법에는 여러가지가 있음

1. 선형 스케줄링 (Linear Scheduling)

: β를 일정하게 증가시키는 방법

2. 제곱 스케줄링 (Quadratic Scheduling)

: β 값을 제곱 함수 형태로 증가시키는 방법, 초반과 후반의 차이를 더 극대화

3.  코사인 스케줄링(cosine scheduling)

: α_t를 코사인 함수로 정의하고 β_t 값을 유도하는 방식

-> β들이 [-1,1]로 스케일된 데이터에 비해 작게 설정되어 reverse process와 forward process가 대략 같은 함수 형태를 유지하게 하면서 신호 대 잡음 비율을 가능한한 작게 유지하려 함.(L_T=D(q(x_T|x_0) || N(0, I))≈10^-5 bits per dimension, L_T가 실제 신호 대 잡음 비율을 의미하는 것은 아니지만 L_T가 작게 하면 신호 대 잡음 비율이 줄어듬.)

- reverse process를 나타내기 위해 그룹 정규화(배치 정규화와 유사하지만, 미니배치의 크기에 의존하지 않는 정규화기법, 채널을 여러 그룹으로 나누어 각 그룹 내에서 정규화를 수행한다.)를 전체적으로 사용하는 PixelCNN++과 유사한 U-Net backbone 구조를 가짐. 파라미터는 시간에 걸쳐 공유되며, Transformer 모델에서 사용하는 사인파 위치 인코딩을 적용하여 각 시간 단계에 맞는 위치 정보를 제공함. 또한 16*16 크기의 feature map에서 self-attention mechanism을 사용하여 중요한 부분에 집중할 수 있게 함.

4.1 Sample quality

- 아래 Table 1은 CIFAR10에 대해 Inception scores, FID scores, negative log likelihoods(lossless codelengths)을 기록해놓았다. FID score이 3.17로 우리의 uncoditional model은 알려진 대부분의 model보다 좋은 sample quality룰 보였다. 테스트 세트를 기준으로 계산한 FID score는 5.24로 알려진 다른 모델들의 test set FID score보다 좋다.

- 위 표를 보면 simplified objective로 구할 때보다 true variational bound로 구할 때 더 좋은 codelengths(작을수록 이득임.)를 보였는데 후자가 더 좋은 sample quality(큰 IS, 작은 FID)를 보였다.

다음은 실제로 만든 이미지들

4.2 Reverse process parameteriation and training objective ablation

- 아래 Table 2는 reverse process parameterization과 training objectives의 sample quality effects를 보여줌.

- 우리는 µ를 예측하는 방식은 true variational bound를 구할 때만 잘 작동하는 것을 확인함.

- 또한 reverse process의 variances를 학습하는 것도 실험해보았지만 고정 분산과 비교했을 때 불안정한 학습과 더 낮은 샘플 품질을 초래함.

- ε를 예측하는 방식은 고정 분산으로 variational bound를 사용하여 학습했을 때는 µ를 예측하는 것과 거의 동일하게 수행되었지만, simplified objecive로 학습한 경우 훨씬 더 잘 수행됨.

4.3 Progressive decoding

- Table 1에 CIFAR 10 model의 codelengths도 보여주는데 train과 test 사이의 gap이 최대 0.03 bits per dimension이므로 이는 다른 model에서 보고된 차이보다 적은 차이로 우리의 diffusion model이 overfitting되지 않았음을 나타냄. 

- 하지만 여전히 우리의 lossless codelengths는 energy based models와 score matching에서 보고된 큰 추정치보단 좋지만, 다른 유형의 likelihood-based generative models와 비교할 때는 경쟁력이 없음.

-> 우리의 샘플들이 높은 품질을 유지하고 있음에도 불구하고 우리는 diffusion models이 lossy compression에서 높은 품질의 샘플을 생성할 수 있는 inductive bias를 가지고 있다고 판단함.

-> L_1+...+L_T를 rate, L_0를 distortion으로 보면 rate는 1.78 bits/dim이고 distortion은 1.97bits/dim임.(둘이 더해서 3.75)

-> 그런데 이 1.97 비트/차원은 원래 픽셀 값이 [0, 255]로 변환했을 때 RMSE가 0.95에 해당하는 매우 작은 값임.

즉, 절반이 넘는 lossless codelength가 지각할 수 없는 distortions를 설명하는 것임.

Progressive lossy compression

- 우리는 progressive lossy code를 도입하여 이 rate-distortion behavior를 더 탐구할 수 있는데 이 부분은 어려워서 생략...

 아무튼 결론은 bits의 대부분이 impreceptible distortions를 설명하고 있다는 거임.

Progressive generation

- 이번엔 아래 식을 통해 각 reverse process에서 x_0를 예측해봄.

이 때 아래 그림처럼 대규모 이미지 특징이 먼저 나타나고 세부 사항은 나중에 나타남.

아래 그림은 다양한 t에 대해 x_t(4개 중 오른쪽 아)가 고정된 상태에서 x_0(가장 오른쪽 아래 사진)를 예측한 결과를 보여줌. t가 작을 때는 거의 모든 것이 보존되고 t가 클 때는 대규모 특징만 보존됨.

-> conceptual compression의 hints가 된다 하는데 conceptual compression은 중요한 대규모 특징을 먼저 보존하고, 세부 사항을 나중에 보존하는 방식임.

Connection to autoregressive decoding

- variational bound를 아래와 같이 다르게 쓸 수 있는데

만약 diffusion process length T가 data의 dimensionality라 하고 forward process q(x_t|x_0)를 t번째 좌표를 만들어내는 것 같이 조건들을 설정해주면 여기서 p_θ를 훈련하는 것이 마치 autoregressive model이 됨.

-> 이 때문에 Gaussian diffusion model을 autoregressive model의 일반화라 볼 수 있음.

4.4 interpolation

 - 우리는 stochastic encoder q를 사용하여 x_0, x_0'을 interpolate할 수 있음. q를 사용하여 x_t와 x_t'으로 변환한 후 선형 보간으로 새로운 벡터를 만든 다음 reverse process를 통해 이미지 공간으로 decoding 하는 것임.

즉, 아래와 같이 decoding하는 것.

이는 위 그림의 왼쪽에 모사된 것처럼, 소스 이미지의 손상된 버전을 선형 보간 후 reverse process를 통해 인위적인 흔적을 제거하는 효과를 가짐. 

 오른쪽 그림은 노이즈를 고정한 채 람다를 달리하면서 보간 및 복원을 보여주는 데 높은 품질의 복원을 생성하며, 포즈, 피부, 톤, 헤어스타일, 표정 및 배경과 같은 속성이 부드럽게 변하는 그럴듯한 보간을 생성하지만 안경은 그렇지 않음. 

5. Related Work

- 확산 모델은 흐름(flow) [9, 46, 10, 32, 5, 16, 23]과 VAE [33, 47, 37]와 비슷할 수 있지만, 확산 모델은 q에 매개변수가 없고 최상위 잠재 변수 x_T가 데이터 x_0x와 거의 제로의 상호 정보를 가지도록 설계되었다. 우리의 ϵ-예측 역방향 과정 매개변수화는 확산 모델과 여러 노이즈 수준에서의 denoising score matching과 annealed Langevin dynamics를 통한 샘플링 사이의 연결을 확립한다 [55, 56]. 그러나 확산 모델은 간단한 로그 가능도 평가를 허용하고, 학습 절차는 변분 추론을 사용하여 Langevin dynamics 샘플러를 명시적으로 학습한다 이 연결은 또한 특정 가중 형태의 denoising score matching이 Langevin-like 샘플러를 학습하기 위한 변분 추론과 동일하다는 역방향 의미도 가진다. 마코프 연쇄의 전이 연산자를 학습하는 다른 방법들로는 infusion training [2], variational walkback [15], generative stochastic networks [1], 그리고 기타 [50, 54, 36, 42, 35, 65]가 있다.

- Score matching과 에너지 기반 모델링의 알려진 연결에 의해, 우리의 작업은 에너지 기반 모델에 대한 최근 연구 [67–69, 12, 70, 13, 11, 41, 17, 8]에 영향을 미칠 수 있다. 우리의 비율-왜곡 곡선은 하나의 변분 경계 평가에서 시간에 걸쳐 계산되며, 이는 annealed importance sampling의 한 번의 실행에서 왜곡 페널티에 대한 비율-왜곡 곡선을 계산하는 방법을 연상시킨다 [24]. 우리의 점진적 디코딩 논증은 convolutional DRAW 및 관련 모델 [18, 40]에서 볼 수 있으며, 오토회귀 모델 [38, 64]에 대한 더 일반적인 설계 또는 샘플링 전략으로 이어질 수도 있다.

6. Conclusion

 우리는 diffusion models를 사용하여 높은 품질의 이미지 샘플을 제시했으며 diffusion models과 variational inference를 통한 Markov chains training, denoising score matching과 Langevin dynamics(이를 확장한 에너지 기반 모델), autoregressive models, progressive lossy compression 간의 연결을 발견했따. Diffusion models이 이미지 데이터에 대해 뛰어난 inductive bias를 가지고 있는 것 같으므로 이 모델이 다른 데이터 형태에서도 유용한지, 다른 유형의 생성 모델 및 머신 러닝 시스템의 구성 요소로서 어떻게 활용될 수 있는지 조사해보고자 한다.

Broader Impact

- 우리의 확산 모델 연구는 기존의 다른 유형의 딥 생성 모델에 대한 연구와 유사한 범위를 갖는다. 예를 들어, GAN, 흐름, 오토회귀 모델 등의 샘플 품질을 개선하려는 노력이 그렇다. 이 논문은 확산 모델을 이러한 기법들의 범주 내에서 일반적으로 유용한 도구로 만드는 데 기여하며, 이는 생성 모델들이 세상에 미친(그리고 앞으로 미칠) 영향을 확대하는 데 기여할 수 있다.

- 불행히도, 생성 모델의 악의적인 사용은 많이 알려져 있다. 샘플 생성 기술은 정치적 목적으로 고위 인물들의 가짜 이미지와 비디오를 생성하는 데 사용될 수 있다. 소프트웨어 도구가 등장하기 전에도 수작업으로 가짜 이미지를 만드는 일이 있었지만, 생성 모델과 같은 도구는 이 과정을 더 쉽게 만든다. 다행히도, 현재 CNN이 생성한 이미지에는 미묘한 결함이 있어 이를 감지할 수 있지만, 생성 모델의 발전은 이를 더 어렵게 만들 수 있다. 또한, 생성 모델은 훈련된 데이터셋의 편향을 반영힌다. 많은 대규모 데이터셋이 인터넷에서 자동화된 시스템에 의해 수집되기 때문에, 특히 이미지가 라벨링되지 않은 경우 이러한 편향을 제거하기가 어렵다. 이러한 데이터셋을 기반으로 훈련된 생성 모델의 샘플이 인터넷에 확산되면, 이러한 편향은 더욱 강화될 것이다.

- 반면에, 확산 모델은 데이터 압축에 유용할 수 있다. 데이터가 더 높은 해상도를 가지게 되고 글로벌 인터넷 트래픽이 증가함에 따라, 이는 인터넷 접근성을 넓은 대중에게 보장하기 위해 중요할 수 있다. 우리의 연구는 이미지 분류부터 강화 학습에 이르는 다양한 다운스트림 작업을 위한 라벨이 없는 원시 데이터에 대한 표현 학습에 기여할 수 있으며, 확산 모델은 예술, 사진, 음악에서 창의적인 용도로도 사용될 수 있다.

 원래 autoencoder는 1980년대 만들어진 신경망 모델 중 하나로 데이터를 효율적으로 표현하기 위해서 고안되었다.  Encoder 구조를 통해 데이터를 압축하고 decoder 구조로 다시 복원하는 과정을 통해 데이터의 중요한 특징을 추출하는 것이 목표였는데 이 구조를 유지한 채 각각의 목적에 맞게 설계되어 noise를 제거하는 DAE, 새로운 데이터를 생성하는 VAE 등으로 발전하였다.

 오늘 다뤄볼 모델은 VAE이다. 참고한 논문은 'Auto-Encoding Variational Bayes' by. Diederik P Kingma, Max Welling이고 아래 사이트에 기재되어있다.

[1312.6114] Auto-Encoding Variational Bayes (arxiv.org)

수식은 아래 사이트를 참고하였습니다.

[논문리뷰] Auto-Encoding Variational Bayes(VAE) 모든 수식 알아보기 (tistory.com)

Abstract

• 우리는 어떻게 intractable posterior distributions를 갖는 continuous latent variables가 존재하고 large datasets에서 효율적인 inference와 directed probabilistic models를 학습할 수 있을까?
  • by. stochastic variational inference & learning algo.
• 우리는 이 논문을 통해 크게 2가지를 기여함.
  • reparameterization of the variational lower bound로 간단한 SGD로 최적화 가능한 a lower bound estimation를 만든 것
  • approximate inference model을 proposed lower bound estimator를 통해 intractable posterior에 fitting 시켜 posterior inference가 효율적으로 수행되게 하는 것

1. Introduction

• Directed probabilistic models에서 연속적인 잠재 변수와 파라미터가 intractable(앞으로 '난해한, 복잡한'으로도 자주 쓰겠다.) posterior distributions(p(z|x))를 가지면 어떻게 효율적으로 근사추론하고 학습할 수 있을까?
  • variational Bayesian(VB) approach
    • 복잡한 사후 분포를 더 간단한 분포로 근사하여 계산의 용이성을 도모하는 접근법
    • 기본 아이디어는 근사 분포 선택과 ELBO(Evidence Lower Bound, evidence는 x의 log-likelihood) 최적화로 이루어짐.
    • variational은 원함수가 어려울 때 다른 함수로 대체하는 방법을 의미

    • 

    • 

  • VB의 방법 중 하나인 mean-field접근법이 있는데
    • 근사 분포를 각 잠재 변수가 독립적인 형태로 근사해 다변수 분포를 단변수 분포의 곱으로 나타낸다.
    •  

      

    • mean-field 접근법은 물리학에서 유래되어 복잡한 시스템에서 모든 상호작용을 계산하기 어려워 평균적은 장으로 근사하는 방식임.
    • 하지만 이런 가정을 하더라도 general case에서 approximate posterior q(z)는 적분하기 어렵다.(논문에서 말하는 analytical solutions은 수학공식을 통한 적분을 말함.)
    • 따라서. variational lower bound를 reparameterization을 통해 simple differentiable unbiased estimator of the lower bound를 만들고 이 SGVB(Stochastic Gradient Variational Bayes) estimator가 efficient approximate posterior inference에 사용됨을 보일 것이다.
• i.i.d. dataset(각 data points 독립적, from 동일한 분포)과 data point마다 continuous latent variable이 있는 경우 Auto-Encoding VB(AEVB) algo. 를 제안한다.
• 이 알고리즘은 SGVB estimator로 inference와 learning을 효율적으로 가능하게 해 recognition model을 최적화할 수 있다.
• 그 결과, simple ancestral sampling(모델의 구조에 따라 순차적으로 sampling 하는 것)으로도 매우 효율적인 approximate posterior inference가 가능하다. -> 빠른 learning으로 이어짐., 비싼 interative infernce(MCMC) 같은 거 안 써도 됨.
• 이런 learned approximate posterior inference model은 다양한 task에도 쓰일 수 있고 특히 neural network와 결합하면 variational auto-encoder가 된다.

2. Method

• 이번 장에선 lower bound estimator(a stochastic objective function)을 유도해 볼 것이다.
• 사전에 하는 가정은 i.i.d dataset with latent variables per datapoint, maximum likelihood(ML) or maximum a poseteriori(MAP)가 목표이고 latent variables에 대해 variational inference를 하는 것이다.
  • variational inference: 복잡한 확률 분포를 계산하기 쉬운 분포로 근사하여 접근하는 방법
• 이 시나리오를 글로벌 파라미터에 대해서도 variational inference를 수행하는 경우로 확장하는 것은 간단하다. 해당 알고리즘은 appendix F에 올려두지만 이 경우에 대한 실험은 앞으로 연구가 필요하다.(VAE는 잠재 변수에 대해서만 variational inference를 수행함.)
• 우리의 아이디어는 실시간으로 연속적인 생성과 전달되어 분석 및 처리가 요구되는 streaming data에서도 사용할 수 있지만 간단히 하기 위해 fixed dataset으로 가정하겠다.

• 

  
  
  

• 

  
  
  

• 

  
  
  
• 2.1 Problem scenario
  • 우리는 데이터 x가 보이지 않는 continuous random variable z로부터 random process를 통해 만들어졌다고 가정하자.
  • 이 프로세스는
    • z는 true parameter 세타*로 이루어져 있는 prior distribution p_{세타*}(z)로부터 생성되고
    • x가 conditional distribution p_{세타*}(x|z)로부터 만들어진다.
  • 이렇게 two steps로 이루어진다.
  • p_{세타*}(z), p_{세타*}(x|z)는 각각 parametric families of distributions에서 온다고 가정하는데 이는 각각의 분포가 특정한 parameter를 통해 결정됨을 의미한다.(가우시안이면 평균과 분산)
  • p_{세타*}(z), p_{세타*}(x|z)는 세타와 z에 대해 미분 가능도 가정
  • 그리고 이러한 과정은 대부분 hidden process이기 때문에 true pparameter 세타*, latent variable z 역시 알지 못한다.
  • 매우 중요한 것이, marginal과 posterior에는 simplifying assumptions를 적용하지 않은 것!

  • 

  • 우리는 다음의 case에서도 작동하는 일반적인 알고리즘에 관심 있다.
    • Intractability: marginal likelihood인 p(x)= ∫p(z)p(x|z)dz가 intractable한 상황. 즉, p(x)를 계산할 수 없고 미분할 수 없다. + true posterior density p(z|x) = p(x|z)p(z)/p(x)가 intractable하여 EM algorithm을 사용할 수 없는 상황 + mean-field VB algo.에서 필요한 integrals 역시 intractable한 상황. 
      • likelihood functions p(x|z)가 적당히 복잡하기만 하더라도 자주 생기는 상황이다.(neural network with nonlinear hidden layer이면 당연)
      • EM algo를 사용하려면 p(x)를 알아야 한다.

      • 

    • A large dataset: 데이터 셋이 많아 모든 데이터를 한꺼번에 처리하는 batch optimzation을 못하기 대문에 minibatches가 single datapoints로 update를 하려 한다. 이때, sampling methods는 각 데이터 포인트 혹인 mini batches마다 반복적으로 sampling을 수행해야 하므로 너무 많은 비용이 든다.
  • 우리는 이에 대한 해결책을 제시하기 위해 세 가지 관련된 문제로 쪼갰다.
    • Efficient approximate ML or MAP estimation for the parameter theta
      • hidden process를 모방하고 real data와 비슷한 데이터를 만들 수 있다.
    • Efficient approximate posterior inference of the latent variable z given an observed value x for a choice of parameters theta
      • coding or data representation tasks에 적합하다.
    • Efficient approximate marginal inference of the varible z
      • 여러 방면에 사용된다.(prior of x이니까) 
  • 위 문제들을 해결하기 위해 recognition model q_{파이}(z|x)를 도입한다. 이는 intractable true posterior p(z|x)를 근사한 것이다. 
  • 또한 알아야 할 것이 mean-field variational inference와 다르게 factorial할 필요가 없고(factorial하다는 것은 잠재변수끼리 독립적일 필요가 없다는 것) recognition model의 파라미터 파이는 단순한 수학적 기댓값 공식으로 계산될 필요가 없다.
    • 대신 recognition model parameters 파이와 generative model parameters 세타를 동시에 학습하는 learning method를 도입할 것이다.
  • Coding theory 관점에서 unobserved variables z는 latent representation or code로 볼 수 있어 recognition model q(z|x)를 확률적 encoder로 볼 수 있다. 주어진 datapoint x에 대해 z에 대한 분포를 만들기 때문.(단, x가 생성될 수 있는 z들로 한정) 같은 관점으로 p(x|z)는 확률적 decoder가 된다. 주어진 code z로 대응되는 x에 대한 분포를 만들기 때문.
• 2.2 The variational bound 
  • marginal likelihood는 individual data points의 marginal likelihoods의 합으로 구성된다.

  • 

  • 그리고 이 각 데이터 포인트의 marginal likelihood는 다음과 같이 쓸 수 있다.

  • 

     
  • 
    
    

    

  • 우측항의 왼쪽 KL divergence는 true poseterior와 그것의 근사에 대한 차이이고 이 값은 non-negative이기 대문에 두 번째 항은 L(세타, 파이; x)는 (variational) lower bound on the marginal likelihood of datapoint i와 동일하다.
  • 이 식은 다시 아래와 같이 쓸 수 있는데
  •  

  

  • 

  • 우리는 lower bound를 variational parameters인 파이와 generative parameters 세타로 미분하고 최적화하려 하는데 파이에 대한 미분은 좀 문제가 있다.
  • Monte Carlo gradient estimator를 기댓값에 적용한 식은 다음과 같은데 

  • 

  • 이 값이 매우 높은 분산을 보여 실용적이지 않다는 것이다.(아래 2.4 section에서 reparameterization trick과 score function을 비교해 보겠다.)
• 2.3 The SGVB estimator and AEVB algorithm
  • 이번 section에선 실용적인 lower bound estimator랑 그것의 미분된 것을 구해볼 것인데 여기서 사용되는 기법이 posterior을 근사한다고 가정한 q(z|x) 뿐만 아니라 q(z)에도 적용될 수 있다.(변분 추론 기법이 적용 범위가 넓음을 의미) 
  • section 2.4에서 소개될 조건 하에서 선택된 approximate posterior q(z|x)에 대해 q(z|x)를 따르는 random variable z틸다를 differentiable transformation g(e, x)와 보조 noise variable e를 통해 reparameterize할 수 있다.

  • 

  • section 2.4에 적절한 분포 p(e)와 함수 g(e, x)를 잡는 법을 소개한다. 우리는 이제 f(z)의 Monte Carlo estimates of expectations을 다음과 같이 계산할 수 있다.

  • 

  • 이를 eq. (2)에 적용하여 일반적인 Stochastic Gradient Variational Bayes (SGVB) estimator를 얻는다.

  • 

  • 

    
    
    
  • eq.(3)의 KL-divergence D(q(z|x)||p(z))는 주로 적분가능하다.(appendix B에서 볼 거임.) 따라서 오른쪽에 있는 expected reconstruction error E_q[log(p(x|z))]의 sampling만 근사하면 된다. 식의 의미를 살펴보면 KL-divergence term은 approximate posterior를 prior p(z)에 가깝게 하는 파이에 대해 regularizing 하는 term이다. eq. (3)에 reparameterization trick을 사용하면

  • 

  • 
    
    
  • 이렇게 두 번째 버전 SGV estimator가 나오고 이 식은 일반적으로 generic estimator보다 작은 분산을 갖는다.
    • 왜 그런지에 대한 얘기는 없는데 단순히 z에 의존하는 term이 두 개여서 그런가 싶다...
  •  data set X가 총 N datapoints가 있다면 minibatches로 전체 marginal likelihood를 다음과 같이 구할 수 있다.

  • 

  • 그럼 Auto-Encoding VB (AEVB) algo. 다음과 같이 구현된다.

  • 

  • auto-encoders관의 관계는 eq.(7)을 보면 명확하다. 
    • 첫 번째 term은 KL divergence of the approximate posterior from the prior로 regularizer의 역할을 수행한다.
    • 두 번째 term은 expected negative recounstruction term으로 함수 g()는 x와 noise vector e를 approximate posterior에서 뽑은 sample인 z로 mapping해주는 역할이고 그 z가 두번째 term으로 들어간다.
      • log(p(x|z))는 under generative model, given z에서 datapoint x의 확률을 의미하는데 당연히 크면 클수록 좋으니 negative가 붙어 negative reconstruction error이다.
• 2.4 The reparameterization trick
  • 우리는 문제를 해결하기 위해 q(z|x)에서 샘플을 생성하는 부분을 
  • score function trick에서

  • 

  • reparameterization trick으로 

  • 

  • 넘어왔다.
  • -------------------------------------------------------------------------------------------------------------------------------------------------
  • + 이 부분은 논문에 없지만 score function trick과 reparameterization trick의 분산을 비교해 보겠다.
  •  

    

  • 

    
    
    

  • 

    
    
    
  • 그럼 왜 높은 분산을 가지면 안 좋은 걸까?
    • 높은 분산은 학습에 불리하기 때문. 학습 속도가 느려지고 최적화 과정에서 적절한 방향으로 수렴하지 못할 가능성이 높아짐. 즉 학습이 불안정해지는 것. + overfitting, underfitting 가능성 모두 증가
  • -------------------------------------------------------------------------------------------------------------------------------------------------
  • 재매개변수화 트릭의 본질은 간단하다. 연속 확률 변수 z가 q(z|x)라는 분포를 다른다고 가정하면 random variable z를 z=g(e, x) 같이 결정론적 변수로 표현할 수 있다.(e는 independent marginal p(e)를 갖는 보조 변수, ㅎ()는 파라미터 파이에 의해 매개변수화된 벡터 함수)
  • 이 트릭이 왜 우리에게 유용한가면 이를 통해 q(z|x)의 Monte carlo estimate of the expecation를 미분가능하게 해 주기 때문이다.
    • 증명은 z = g(z|x)인 결정론적 mapping이 주어졌을 때
    • 확률 분포를 변환하면(이 부분이 잘 이해가 안 갔는데 수리 통계를 공부해야 할 거 같다..)
    •  

      

    • 가 성립하고  이어서

    • 

    • 가 성립한다. 2.3에서 이 트릭으로 variational lower bound에 대해 미분가능한 estimator를 얻은 것.
  • 몇 가지 예시를 들어보면 만약 z가 단변량 가우시안 분포를 따른다고 하면

  • 

    
    
    
  • 유효한 reparameterization은 z=μ+σϵ이다.
  • 따라서 기댓값이 다음과 같이 변하고
  •  
    

  

• • 몬테카를로 추정을 하면
  •  

    

    
    
    
  • 가 된다.
  • 이렇게 특정 분포 q(z|x)에 대해 미분 가능한 변화 g(e, x)아 보조 변수 e를 선택할 수 있는 건 다음 3가지 접근 방식 중 하나라도 가능할 때이다.
    • Tractable inverse CDF: 만약 q(z|x)의 역 누적 분포 함수를 계산할 수 있는 경우 적용할 수 있다. 이 경우 g(e, x)에서 e는 U[0, I]를 따른다.
      • examples: Exponential, Cauchy, Logistic, Rayleigh, Pareto, ...
    • Gaussian example과 유사하게 "location-sacle" family of distributions에 해당하는 경우 standard distribution을 따르는 e를 auxiliary variable로 사용할 수 있다.
      
      • examples: Gaussian, Laplcae, Elliptical, Student's t, Logistic, Uniform, ...
    • 랜덤 변수를 보조 변수의 다양한 변환으로 표현할 수 있는 경우
      • examples: Log-Normal(정규 분포된 변수의 지수 변환), Gamma(지수 분포된 변수들의 합), Dirichlet(감마 변수들의 가중합), Chi-Squared, Beta, ...
  • 이러한 접근법들이 모두 불가능한 경우 inverse CDF에 대한 좋은 근사 방법이 있긴 한데 PDF의 시간 복잡도와 필적한 계산이 든다.

3. Example: Variational Auto-Encoder

• 이번 section에선 probabilistic encoder로 neural network를 사용한 example인 variational auto-encoder에 대해 알아보겠다. parameters 파이와 세타는 AEVB algo.에 의해 jointly optimize된다.
• latent variables의 prior p(z)를 centered isotropic multivariate Gaussian으로 가정한다.(평균이 0이고 공분산 행렬이 단위행렬인 가우시안 분포)

  • 

    
    
    
  • 이 경우 사전 분포는 파라미터를 가지지 않는다. 평균과 공분산 행렬이 고정되어 있으므로 학습할 파라미터가 없다.
•  p(x|z)의 경우 x가 실수값 데이터의 경우 multivariate Gaussian, 이진 데이터의 경우 Bernoulli 분포이다. 이 분포의 파라미터는 MLP(단일 은닉층을 가진 완전 연결 신경망)을 통해 z로부터 계산된다. -> appendix C
• 실제 사후 분포 p(z|x)는 계산하기 어렵다. 반면 이를 근사한 q(z|x)의 형태에는 많은 자유도가 있지만 우리는 실제(계산 불가능하지만) 사후 분포가 대각 공분산을 가진 가우시안 분포를 따른다고 가정한다. 이 경우 variational approximate posterior 역시 대각 공분산 구조를 가진 다변량 가우시안 분포로 설정할 수 있다.

  

• 2.4에서 설명했듯이 z를 아래와 같이 sampling한다

• 

  
  
  
• 이 모델의 경우 p(z) -> prior과 q(z|x)를 Gaussian으로 가정하기 때문에 eq.(7)에 estimator를 사용할 수 있다.(KL-divergence는 근사 없이 계산 가능함!  -> appendix B)
• 따라서 data point x와 모델에 대한 최종 estimator는

• 

  
  
  
• 와 같다.
• 뒤에 붙는 decoding term log(p(x|z))는 modeling하는 data의 type에 따라 Bernoulli or Gaussian MLP이다.

4. Related Work

• 연속적인 잠재 변수를 사용한 model에 적용 가능한 온라인 학습 방법으로 wake-sleep 알고리즘만이 존재한다.
  • 우리의 방법과 마찬가지로 실제 사후 분포를 근사하는 recoginition model을 사용한다.
  • wake-sleep algo.의 단점은 두 objective functions를 동시에 근사해야 하는데 두 목표함수가 결합되어 이는 marginal likelihood의 최적화와 다르다는 것이다.
  • 장점은 이 알고리즘이 이산 잠재 변수를 포함하는 모델에도 적용될 수 있다는 점이다.
  • 데이터 포인트당 계산 복잡도는 AEVB와 동일하다.
• 최근에 stochastic variational inferencec가 주목받고 있다. 최근에 2.1에서 논의된 단순 gradient esstimator의 높은 분산을 줄이기 위해 제어 변수를 도입하였고 이를 exponential family approximations of the posterior에 적용하였다. 좀 더 일반적인 방법으로 이 paper와 유사한 reparameterization이 exponential-family approximating distributions의 natural parameters를 학습하기 위한 stochastic variational inference algo.에 사용되었다.
• AEVB algo.는 directed probabilistic models(trained with a variational objective)와 auto-encoders 사이의 관계를 보여준다. 선형 auto-encoders와 generative linear-Gaussian models 사이의 관계는 오래전부터 알려져 있었다. 
  • PCA와 사전 분포 p(z) = N(0, I)와 조건부 분포 p(x|z) = N(x:Wz, eI)(특히 e가 매우 작은 경우)를 가지는 특수한 경우의 linear-Gaussian models의 ML의 해와 일치함을 보였다.
• autoencoders에 관한 최근 관련된 연구는 비정규화 오토인코더의 훈련 기준이 입력 X와 잠재 표현 Z 사이의 mutual information의 하한을 최대화하는 것과 일치함을 보여주었다. mutual information을 최대화하는 것은 conditional entropy를 최대화 하는 것과 같으며 이는 autoencoding 모델에서 데이터의 expected loglikelihood(negative reconstrunction error)의 하한이다. 그러나 이런 reconstruction criterion으론 유용한 표현을 학습하기에 충분하지 않다는 것이 알려져 있다. 유용한 표현을 학습하기 위해 정규화 기법이 제안되었으며, 그 예로 denoising , constractive and sparse autoencoder 변형이 있다.
• SGVB objective는 variational bound에 의해 정해지는 regularization term이 포함되어 있으며 이는 유용한 표현을 학습하는데 필요한 일반적인 regularization hyperparameter가 필요하지 않다. Predictive sparse decomposition(PSD)와 같은 encoder-decoder 구조와도 관련이 있으며, 여기서부터 영감을 얻었다. 또한 최근에 도입된 Generative Stochastic Networks와도 관련 있다. 이 네트워크에서는 노이즈가 있는 auto-encoders가 마르코프 체인의 transition operator를 학습하여 data distribution에서 샘플링한다. Deep Boltzmann Machines의 효율적인 학습을 위해 recognition model을 사용한 사례가 있다. 이러한 방법들은 unnormalized model(Boltzman machines같은 undirected models)이나 sparse coding models에 중점을 둔 반면, 우리의 제안된 알고리즘은 일반적인 directed probabilistic models를 학습하는데 중점을 둔다.
• 최근 제안된 DARN method 역시 auto-encoding 구조를 사용하여 directed probabilistic model을 학습하지만 이 방법은 binary latent variables에 제한된다. 더욱 최근에는 auto-encoders와 directed probabilistic models 및 우리가 paper에서 얘기한 reparameterization trick을 사용하는 stochastic variational infernce의 연결에 대해 제시했다. 그들의 연구는 우리와 독립적으로 개발되었으며 AEVB에 대한 추가적인 관점을 제시한다.

5. Experiments

• 우리는 MNIST와 Frey Face 데이터셋의 이미지로 훈련시켰고 variational lower bound과 marginal likelihood의 측면에서 학습 알고리즘을 비교했다.
• section 3에서 설명된 generative model(encoder)와 variatinal approximation(decoder)를 사용했으며 인코더와 디코더는 같은 수의 hidden units를 가진다. Frey Face 데이터는 연속적이기 때문에 Gaussian outputs를 갖는 decoder를 사용했으며 encoder와 동일하지만 디코더의 출력에 sigmoid activation function을 사용하여 평균을 (0, 1)의 범위로 제안했다. 여기서 hidden units는 인코더와 디코더 neural networks의 hidden layer를 나타낸다.
  • Frey Face 데이터의 얼굴 이미지의 각 픽셀 값은 0에서 255 사이의 값을 자유롭게 가지지만 MNIST 데이터의 이미지의 각 픽셀 값은 0 or 255 같이 특정 값만 갖는다.
• 파라미터는 algorithm 1. 에서 이야기한 differentiating the lower bound estimator에 의해 계산된 gradients를 쓰는 stochatic gradient ascent로 업데이트되었다. 또한 사전 분포인 p(세타) = N(0, I)에 작은 wieght decay term이 추가되었다. 이 목표를 최적화하는 것은 likelihood gradient가 gradient of the lower bound로 근사되는 approximate MAP estimation을 수행하는 것과 동등하다.
• 우리는 AEVB 알고리즘과 wake-sleep algorithm을 비교했다. wake-sleep 알고리즘과 variational auto-encoder 모두 동일한 encoder( = recognition model)을 사용하였다. 모든 파라미터, variational, generative paramters 모두 N(0, 0.01)에서 random sampling으로 초기화되었다고 MAP를 기준으로 사용하여 jointly stocastically optimize되었다. Stepsizes는 Adagrad로 조정되었으며, Adagrad global stepsize parameters는 초기 몇 번의 반복에서 훈련 세트 성능을 기준으로 {0.01, 0.02, 0.1} 중에서 선택되었다. 미니배치 크기는 M=100으로, 데이터 포인트당 샘플 L=1개를 사용하였다.
• Likelihood lower bound
  • 우리는 MNIST의 경우 500개의 hidden units을, Frey Face 데이터셋의 경우 200개의 hidden units를 가지는 generative models(decoders)와 대응되는 encoders(a.k.a recognition models)로 훈련하였다. Frey Face dataset의 경우 상당히 작은 dataset이기 때문에 overfiting을 막기 위해 hidden units의 개수를 줄였다. 선택된 hidden units의 수는 auto-encoders에 대한 기존 문헌을 기반으로 하였으며, 다양한 알고리즘의 상대적인 성능은 이러한 선택에 크게 민감하지 않았다.
  • 다음 그림은 lower bounds를 비교한 결과를 알려준다.

  • 

  • 흥미롭게도, 과도한 잠재 변수가 overfitting을 초래하지 않았으며, 이는 variational bound의 regularizing nautre로 설명된다.
• Marginal likelihood
  • 매우 저 차원의 잠재 공간의 경우 learned generative models의 marginal likelihood를 MCMC estimator를 통해 추정할 수 있다.(자세한 내용은 appendix D에) Encoder와 decoder로 neural networks를 사용했으며 이번에는 100 hidden units과 3개의 잠재변수를 사용하였다.(높은 차원의 잠재 공간에서는 추정치가 신뢰할 수 없게 되었다.) 
  • 다시 한번 MNIST dataset을 사용하였다. AEVB와 Wake-Sleep method를 hybrid Monet Carlo(HMC) sampler를 사용한 Monte Carlo EM(MCEM)과 비교하였다.(자세한 내용은 appendix E에) 우리는 작은 훈련 세트와 큰 훈련 세트 크기에 대해 세 알고리즘의 수렴 속도를 비교했다. 
  • 결과는 아래 그림에 나와있다.

  • 

    
    
    
• Visualization of high-dimensional data
  • 만약 저 차원 잠재 공간(ex. 2D)를 선택하면 learned encoders(recognition model)을 통해 high-dimensional data를 low-dimenstional manifold로 project할 수 있다. MNIST와 Frey Face 데이터 셋의 2D latent manifolds 시각화는 appendix A에 있다.

6. Conclusion

• 우리는 연속적인 잠재 변수를 사용한 효율적인 근사 추론을 하는 variational lower bound의 새로운 estimator인 Stochastic Gradient VB(SGBV)를 소개했다.
• 제안된 estimator는 standard stochastic gradient method를 통해 간단히 미분 및 최적화가 가능하다.
• i.i.d. dataset과 continuous latent variables per datapoint의 경우, 우리는 SGVB estimator를 사용하여 근사 추론 모델을 학습하는, 효율적인 추론 및 학습 알고리즘인 Auto-Encoding VB (AEVB)를 소개한다.
• 이론적인 장점은 실험 결과에 반영되었다.

7. Future Work

• SGVB estimator와 AEVB algorithm은 continuous latent variables를 사용하는 거의 모든 inference and learning problem에 적용될 수 있으므로, 미래 연구 방향이 많이 있다.
  • AEVB와 jointly 훈련된 encoders와 decoders에 더 깊은 신경망(ex. convolutional networks)를 사용하여 hierachical generative architectures를 학습하는 것
  • time-series models(= dynamic Bayesian networks)
  • SGVB를 global parameters에 적용
  • 복잡한 noise distributions를 학습하는데 유용한 잠재 변수를 가진 supervised model

Appendix

A. Visualisations

 SGVB로 학습된 모델의 잠재 공간과 이에 대응하는 observed space의 시각화는 아래 두 그림이다.

B. Solutions of -D(q(z)||p(z)), Gaussian case

 Variational lower bound(the objective to be maximized)는 KL term을 포함하는데 이는 종종 analytically 적분 된다.

 VAE에서는 prior p(z)를 N(0, I)로 가정하고 posterior approximation q(z|x) 역시 가우시안이기 때문에 위 항이 생각보다 깔끔하게 적분 되는데 식을 전개해 보겠다.

C. MLP's as probabilistic encoders and decoders

 VAE에서 신경망은 probabilistic encoders와 decoders로 사용된다. Encoders와 decoders의 선택은 데이터와 모델의 유형에 따라 다양한데 논문에서는 비교적 간단한 신경만인 multi-layered perceptrons(MLPs)를 사용하였다. 인코더로는 가우시안 출력을 가지는 MLP를 사용하였고, 디코더로는 데이터의 유형에 따라 가우시안 또는 베르누이 출력을 가지는 MLP를 사용하였다.

- C.1 Bernoulli MLP as decoder

 이 경우 p(x|z)를 z에서 계산된 확률을 갖는 다변량 베르누이 분포로 가정한다. 이는 단일 은닉층을 가진 완전 연결 신경망으로 이루어져 있다.

- C.2 Gaussian MLP as encdoer or decoder

 이 경우 인코더와 디코더는 대각 공분산 구조를 가진 다변량 가우시안 분포로 가정한다.

D. Marginal likelihood estimator

 우리는 샘플링된 공간(잠재변수 z의 공간)의 차원이 낮은 경우(5차원 이하)와 충분한 샘플이 있는 경우 marginal likelihood에 대한 좋은 estimates를 만드는 marginal likelihood estimator를 유도하였다. p(x,z) = p(x)p(x|z)를 생성 모델이라 하고 given x에 대해 marginal likelihood p(x)를 estimate해보자.

 추정 과정은 세 단계로 구성된다.

1. Hybrid Monte Carlo같은 gradient-based MCMC를 사용하여 posterior에서 z 값을 L개 샘플링한다.

2. z samples에 대해 density estimator q(z)를 fit한다.

3. 다시, posterior에서 L개의 값을 샘플링하고 이러한 samples와 잘 학습된 q(z)를 다음 estimator에 대입한다.

E. Monte Carlo EM

 Monte Carlo EM algorithm은 인코더를 사용하지 않고 잠재 변수z의 posterior p(z|x)에서 직접 샘플링한다. 이 때 posterior의 gradient를 사용하는데 gradient는 다음과 같다.

 Monte Carlo EM은 accpetance rate가 90%가 되도록 자동으로 조정된 stepsize로 10번의 HMC leapfrog 단계를 수행한 후, 획득된 sample을 이용하여 5번의 가중치 업데이트를 수행한다. 모든 알고리즘에서 파라미터는 Adagrad stepsizes(annealing schedule 동반)를 사용하여 update 된다.

 Marginal likelihood는 학습 및 테스트 세트의 첫 100개의 데이터 포인트에서 추정되었으며, 각 데이터 포인트에 대해 4번의 리프로그 단계를 사용하여 Hybrid Monte Carlo로 잠재 변수의 posterior에서 50개의 값을 샘플링했다.

F. Full VB

 논문에서 설명한 바와 같이, 잠재 변수에 대해서만 수행한 것과는 달리, 파라미터 세타와 잠재 변수 z 모두에 대해 variational inference를 수행할 수 있다. 여기서는 세타에 대한 estimator를 유도하겠다.

 알파로 매개변수화된, 위에서 소개한 parameters에 대한 hyperprior p_{알파}(세타)가 있다고 하자. 그럼 marginal likelihood는 다음과 같이 쓸 수 있다.

... 보충 필요

 이번엔 Diffusion Model 중 하나인 DDPM에 대해 알아보자.

2. MLDL2 수업 - Diffusion Models

**무단 배포 금지** 본 게시물의 저작권은 김태섭 교수님께 있습니다.

1. VAE to Diffusion Models

사실 교수님께서 diffusion을 알기 위해선 VAE가 필수라고 하셔서 VAE를 짚고 넘어가는 거긴함.

Diffusion 출발은 저번 포스트의 마지막인 Hierachical Variational Auto-Encoder(HVAE)이다.(이말이 꼭 HVAE를 기반으로 diffusion이 만들었다는 말은 아니고 이해하는 관점에서 HVAE와 비교하면서 보는 것이 타당하다.)

 HVAE는 VAE를 generalize한 모델이다. 기존에 VAE에선 latent variable이 하나였지만 이 variables를 늘리고 이 들을 hierachical하게 설계하여 일반화한 것이다. Hierachical하게 설계했다는 것은 latent variable사이 순서, 계층 구조가 있음을 의미한다. 따라서 latent variable역시 더 높은 level = 더 abstract한 latent variable로부터 생성된다. 

 HVAE에선 abstract할 수록, higher-level로 갈수록 그 크기가 줄어든다. = compact해진다.(encoder에서 출력하는 다음 층의 평균과 분산이 계속 줄어드는 것)

그리고 joint-distribution에 하나의 성질이 있는데 markov property이다. 즉, 다음 요소를 결정하는 것은 바로 이전의 요소만이라는 것.(만약 이 성질이 없다면 저 수많은 latent variable로부터 계산할 수 없을 것이다..)

 마찬가지로 HVAE에서의 encoder는 x로부터 low-level, high-level의 z를 만들어내고

이 역시 markov property를 가지며 이 q(z_{1:T}|x)는 p(z_{1:T}|x)를 추정하는 분포이다.

 HVAE에서 evidence를 써보면 다음과 같이 전개되는데

 ELBO에 위 두 그림에서 얻은 식을 넣으면 최종 식이 나오게 된다.

 2. Denoising Diffusion Probabilistic Models(DDPM)

  DDPM은 diffusion model의 초기 model로 VAE에 추가적인 3가지 restrictions를 넣는다.

(수업 자료에는 VDM으로 명시하였는데 나는 틀렸다고 생각한다. VDM은 DDPM과 VAE를 결합한 것이기 때문에 잠재 공간에 대한 이야기가 다루어져야 하는데 전혀 다루지 않기 때문이다. 생각해보면 이런 제약을 넣는 것 자체가 더 기초적인 모델임을 의미한다. 즉 VDM은 아래에서 설명하는 첫번째 제약이 없어지기 때문에 잠재 공간에 대한 자유성이 부여되는 한 단계 더 발전한 모델인 것이다. 더 공부해보고 혹시 내 생각이 틀렸다면 수정하겠다.)

(VAE와 DDPM의 공통점은 간단한 분포로 근사해서 생각하는 variational inference, VI를 한다는 것인데 다루는 공간이 각각 latent space, data space자체를 다루는 것이 차이이다.)

 첫번째가 latent variable의 size와 input data의 size가 갖다는 것인데 위의 HVAE에서 higher-level일수록 compact해진다는 특성이 있지만 VDM에서는 모든 latent variable과 input의 size가 같다.

 두번째는 encoder는 학습되지 않고 pre-defined된 linear Gaussian Model를 사용한다.

 여기서 알파는 hyperparameter이다.

 세번째로 위 gaussian parameters인 알파가 위치에 따라(t로 표현해 시간에 따라가 맞는 표현일 수 있다.) 다르고 마지막 latent variabel x_{T}의 분포는 standard Gaussian N(0,I)이라는 것이다.

 그럼 forward diffusion process는

 이 식을 통해 진행되며 이미지에 noise가 첨가되기 때문에 

 encoding 과정을 noising이라고도 부른다.

 + 왜 정규분포가 저렇게 생겼는가?

 일단 정규분포를 사용하는 이유는 정규분포를 다르는 noise를 넣기 위함이고 그럼 위 식에 의해

 이렇게 다음 단계의 latent variable이 생성된다.(VAE에서 reparameterization trick과 동일한데 사실 꼭 이 방법을 사용하지 않아도 된다. 미분 가능하면서 sampling과 같은 효과를 내는 방식이 또 있다는 뜻) 이 때 앞 계수들의 의해 x_{t}의 분산이 x_{t-1}과 같아지게 된다.

 반대로 reverse denoising process는 생성하는 과정으로

 이 과정을 아래의 식으로 하기 때문에

 평균과 분산을 출력할 network 학습이 필요하다.

 정리해보면

이러한 과정이고

 원 data로 부터 잠재변수의 분포를 추론해가는 과정으로 볼 수 있다.

 이번엔 DDPM에서 ELBO을 살펴보자.

 Jensen's Inequality를 사용하였고 식을 더 전개하면

 마지막 식은 markov property로 확률을 대체하고 의미없는 변수들을 다 적분해버리면 나오게 된다. 그럼 이 식은

 이런 세가지 term이 되는데

 먼저 reconsturction term은 VAE와 동일하고

 그 다음 prior matching term은 

 두번째 제약 조건에 의해 encoder는 아래 분포로 고정되고

 세번째 제약 조건에 의해 마지막 잠재변수의 분포는 

 이기 때문에 학습되는 파라미터가 없다. 즉, constant이다.

 마지막 consistency term은 

 디코더의 분포와 인코더의 분포의 차이의 기대값을 모든 항에 대해 계산하는 term으로 의미적으론 reverse process를 forward process에 맞추는 term이다. 그러나 이 term의 Monte Carlo estimate(확률론적 방법을 사용하여 기댓값이나 적분을 계산하는 방식, 주로 random sampling을 사용한다.)는 두 개의 변수를 매 time steps마다 다루기 때문에 높은 variance를 야기한다. 따라서 이를 해결하기 위해 한 step마다 하나의 변수에 대한 기댓값을 구하도록 바꿔어야 한다.

 이는 생각보다 간단하게 해결되는데 markov property로 식을 변형하면 해결된다.

 encoder의 변분 분포를 이렇게 바꿔주면

log내에 곱으로 되어 있는 항들이 cancellation이 일어나게 된다. 

그럼 최종식은

이렇게 되어

 consistency term이 denoising matching term으로 바뀐 것을 알 수 있다.

 Denoising matching term을 좀 더 자세히 살펴보면 왼쪽 식은 encdoer에서 만든 denoising ground-truth signal인데

이렇게 표현된다. 이 수식이 참 만만치 않은데 x_t를 reparameterization으로 구했다 가정하고 접근해보겠다.

 이 수식에서 이상한 부분이라면

 이 부분인데 쉬운 말로 '정규분포는 정규분포를 낳는다.'라는 말에서 유래한 것이다. 자세한 수식은 잡지식에 올려두겠다.

(왜 정규분포는 정규분포를 낳을까?)

2024.05.31 - [다양한 분야의 잡지식] - 말 그대로 다양한 분야의 잡지식

 그럼 이 과정을 확장하여

에 도달하고 추가로 쓰였던 성질의 함수를 전개하면

다시 돌아와서 denoising matching term은

decoder와 ecoder 사이의 차이를 최소화하는 term이다. 그런데 두 정규분포 사이의 KL divergence는 비교적 쉽게 계산되는데

이므로 세타를 최적화 하는 것은 결국 생성 과정의 뮤세타가 인코더의 뮤큐와 같게하는 것이다. 따라서 뮤세타를 뮤큐와 같은 형태로 설계하는 것이 자연스럽고 T시점부터 출발하기 때문에 모르는 x_0를 각 층마다 예측하게 되는 것이다. 이렇게 정한 뮤세타와 뮤큐를 다시 대입하면 

이러한 결론에 도달한다.

DDPM에서는 다른 추가적인 FC layers와 U-Net architecture를 둬서 t시점의 x_0를 예측한다.

 위 그림에 표현되지는 않았지만 skip connections(고해상도의 이미지를 저해상도에 전달하는 것)가 있고 이 과정은 매 time마다 진행된다.

 그런데 여기서 DDPM을 더 개선한 방식이 있는데

 x_0를 마치 x_t와 노이즈의 결합으로 해석하는 방식이다. 그럼 예측한 뮤세타 역시

이런 꼴로 바뀌게 되고 목적함수는

 noise에 대한 식으로 바뀌게 된다.

이렇게 노이즈를 예측하는 방식이 주로 쓰인다고 하는데 그 이유는 노이즈 예측 방식이 더 안정적이고, 학습 과정에서 더 좋은 성능을 보이는 것으로 나타났기 때문이다.

 추가로 DDPM을 응용하는 것으로 주어진 conditions하에서 generate하는 방식이 있는데 encoder는 fix하고

 이미지를 generate하는 방식에서 조건 이미지를 추가로 넣어 U-Net 구조를 학습시킨다.

 그럼 이렇게 훈련된 모델은 각각 새로운 이미지에 대해 해상도를 복원할 수 있고 색체화가 가능하다.

 마지막으로 여러 generative models를 비교해보면

 Diffusion Models는 아무래도 time의 길이가 길어 fast sampling은 힘들지만 high quality와 diversity한 sample을 만든다는 강점이 있다.

 3. 마무리

 이번 포스트에선 diffusion의 가장 간단한 모델인 DDPM을 알아보았다. 이제 preview 영상으로 가봅시다.

3. Preview 영상

 스터디가기 전에 보기로한 영상에 대해 정리해보겠다.

Diffusion Model 수학이 포함된 tutorial (youtube.com)

 읽어보니까 꽤나 DDPM을 이 post에서 깊게 다루어서 생각보다 흐름을 이해하는데는 어렵지 않았다. Preview 영상이니 자세한 수식보단 큰 흐름과 위에 없는 내용 위주로 요약하겠다.

• Denoising Diffusion Models는 두 과정으로 이루어진다.:
  • Foward diffusion process -> noising
  • Reverse denoising precess -> denoising

• 이렇게 알파 바를 잡았던 것이 특정 시점의 x를 한번에 예측하기 위해서였다.

• intractable은 다루기 어려움을 의미한다.

• ELBO 대신 variational upper boudn로 loss를 정의한다.

• noise만으로도 예측 가

• diffusion은 weighting을 equally하게 적용하지만 advanced weighting도 있다.

• Time representation에 사용되는 것으로 sinusoidal positional embeddings(트랜스포머에서 쓰임)와 random Fourier features가 있다.

• DDPM은 reverse의 시그마 t를 베타 t로 두지만 이 역시 학습가능하다.

• T가 작을수록 high-frequency content = low-level details를 생성한다.

• DDIM: 마르코브 property가 아닌 x_0가 모든 x에 영향을 준다.(Non-Markovian diffusion process)

• Score-based Generative Modeling with Differential Equations

SDE 하에서 score-based model과 diffusion이 사실 같다는 놀라운 사실

 수식이 꽤나 어려워보인다.

• 좀 더 deterministic한 모델

• 특정 조건하에서 생성하는 모델

• 새로운 classifier를 도입하여 구현하는 방식

• 마지막으로 직접 만드신 tree

 Preview 영상이 이해가 하나도 안가서 이번학기에 들은 수업자료를 먼저 되짚어보겠다.

1. MLDL2 수업 - VAE 

0. Intro.

 Data distribution을 알면 p(x)를 기반으로 새로운 samples를 만들거나 새로운 데이터 $\widetilde{x}$가 기존에 관측된 data와 유사한지 판단할 수 있다. 따라서 우리는 model을 통해 data distribution을 approximate하는 것이 목표이고 data distribution을 approximate한다는 것은 observed data x에 대해 likelihood p(x)를 maximize하는 것과 같다.(model에서도 관찰된 데이터가 '관찰'될 확률을 높인다는 것) 

 하지만 우리의 data는 보통 high-dimensional하기 때문에 그 분포는 매우 complex하다. 이를 해결하기 위해 indirect approach인 laten variable methods를 사용한다.

 이 접근방법은 우리의 관측된 x가 관측되지 않는 미지의 변수인 z로부터 생성됬다고 가정하는 것이다. 보통 이 미지의 변수를 latent variable 혹은 hidden variable인데 low-dimensional임이 자연스럽다.(high-dimensional하면 도입할 이유가 없다!) 

 이제 우리는 modeling하려던 p(x) 대신 p(x,z)를 예측하면 되고 이 p(x,z)는 p(z)와 p(x|z) 두개의 곱으로 이루어진다. 즉,

 여기서의 likilihood는 우리가 도입한 변수인 z와 x 사이간의 likelihood이고 p(z)는 우리가 가정한, 추측한 확률이므로 prior probability로 해석한다.

 이렇게 latent variable을 도입한 예시로 gaussian을 활용한 방법이 있다.

 Data distribution이 gaussians의 mixtue N개로 되어있다고 가정하여 data가 그 중 하나로부터 생성되었다고 해석한다.

  이 관점의 특징은 z를 discrete variabe로 본다는 것이다. 따라서 z의 분포는 각 값을 가질 확률로 표현되는 discrete distribution이고 각 data는 자신이 generate된 gaussian의 index인 component index를 갖는다. 

 이 아이디어를 확장한 latent variable model로 non-linear lattent variable model이 있는데 이 모델은 latent variable z가 countinuous variable이고 nomarl prior로 부터 sampling되었다고 가정한다. 즉, 

  그런데 Conditional probability의 mean을 예측하는 f(z)는 신경망 같은 함수인 non-linear function 이기 때문에 위 적분이 불가능하다.(선형이면 가능) 그러니까 지금 생긴 문제를 정리해보면 우리는 p(x)의 likelihood를 maximize해야 하는데 직접 접근하는 것이 어려워 latent variable을 도입했지만 그럼에도 불구하고 closed form이 존재하지 않는 상황인 것이다.

 따라서 이를 해결하기 이해 lower bound를 도입해 proxy objectibve(대리 목적함수)를 사용한다.

1 . ELBO: Evidence Lower Bound

  ELBO에서 evidence는 observed data의 log-likelihood를 의미한다. 즉, ELBO는 위에서 말한 proxu objective인것.

 이제 ELBO를 어떻게 해석할지 알아보기 위해 Jensen's Inequality로 식을 전개해볼 것인데 먼저 Jensen's Inequality는 다음과 같다.

 log function은 concave하기 때문에 평균의 함수값이 함수값의 평균보다 큰 성질을 이용할 것이다. 그럼,

 이렇게 식이 전개되고 

 ELBO는 위와 같이 쓸 수 있다. 여기서 세타는 p를 결정하는 parameters이고 파이는 q를 결정하는 파라미터이다.

2. Variational Inference(VI) 

 그 다음은 ELBO를 최대화 하는 과정을 배울 것인데 처음 배울 때 매우 신기한 아이디어인 VI에 대해 알아보자. ELBO를 이렇게 전개할 수도 있다.

 x와 z의 joint distribution에 z에 대한 posterior를 도입한 것인데(이게 참 웃긴게 우리는 z가 x를 만들었다고 가정했지만 여기서의 확률은 given x, z의 분포에서 얻은 확률이다.)  ELBO가 원래 우리가 구하려던 evidence에 z의 분포와 p(z|x)의 KL-divergene를 뺀 것이라는 것이다. 즉, tight한 바운드는 두 분포 z를 예측한 분포(q), z의 posterior이 같았을 때이다.

 그럼 우리는 ELBO를 어떻게 maximize할까? 우리는 E-M Method(Expectation-Maximization Method, velog에 올라와있는데 추후 가져올 예정..)를 통해 세타와 파이를 번갈아가면서 최적화 해 evidence의 max에 가까워질 수 있다. 즉, 간접적으로 intractable(적분 불가능한) evidence를 optimize하는 것이다.

 위에서 말한 expectation과 maximization 과정은 각각 tight한 bound를 만드는(파이를 update), bound를 개선하는(세타를 update)하는 것인데 그림으로 나타내면 다음과 같다.

 하지만 여기까지 글을 읽을 보면서 이상한 점이 느껴져야 하는데 z의 posterior이 과연 뭐냐는 것이고 계산할 수 있냐는 것이냐이다.

  식을 보면 posterior를 알기 위해선 p(x)를 알아야한다.(p(x)가 어려워서 z도입했더니 p(x)를 계산해야 하는 웃픈 상황) p(x)는 매우 복잡한 분포여서 이 식 역시 계산할 수 없으므로 E-M method를 사용하지 못한다.

 그래서 사용하는 방법이 variational approximation인데 이 방법은 간단한 parametric distribution으로 근사하는 방법으로 z의 posterior를 가우시안 분포로 근사한다.(q를 구하는 parameters인  파이는 variational parameters라 하고 복잡한 함수를 근사하여 추론하는 것을 variational inference라 한다. 그리고 variational parameters는 당연히 x에 딸라 달라지므로 x에 대한 함수로 볼 수 있다.)

 물론 이 발상이 항상 타당한 것은 아니다. 아래 그림에서는 꽤나 괜찮은 근사를 하지만

 이 그림에선 매우 poor하다.

 다시 그동안 했던 것들을 되돌아보면

 Evidence를 maximize하려던 목표는 ELBO을 maximize하는 것으로 바뀌었고 이는 q(z)와 p(z|x)의 차이를 최소화하는 과정이 필요했으나 불가능하여 variational approximation을 사용한다. 마치

를 하는 것과 같아진 것.

3. Variational Auto-Encoder(VAE)

 이번엔 VI를 응용한 신경망인 VAE에 대해 알아보자. VAE는 x를 latent variable z로 encoding하고 이를 다시 decoding하여 x와 유사한 x틸다를 생성하는 것이다. z에 대한 posterior이 아닌 prior를 도입해보자.

 여기서 왼쪽 term은 reconstruction term이라 하는데

 근사된 z(q(z)의 확률 분포를 갖는 z)에서 x가 잘 re-generate되는지 확인하는 term이다. 값은 당연히 크면 클수록 잘 근사한 것이다.

 오른쪽 term은 regularization term인데 

 근사한 z의 분포가 실제 z의 분포(prior)과 유사해는지 확인하는 term이다. 값이 크면 차이가 큼을 의미하므로 -가 붙는다.

 일반적으로 VAE에서 z의 prior는 N(0,1)로 가정한다.

  Regularization term을 보면 (생략되었지만 두 분포 모두 given x에서 얻어졌다.) 따라서 variational parameters 역시 VI와 마찬가지로 x에 대한 함수이다. VI에선 언급하지 않았지만 이 부분이 문제가 되는데 그럼 모든 x에 대해서 파이를 계산해야 하는 문제가 생겨 scalable하지 않다.(scalable하다는 것은 많은 양의 x에서도 합리적인 계산이 가능해 확정가능함을 의미한다.) 

 따라서 VAE에서 도입한 것이 encoder architecture인데 이 구조로 amortized inference가 가능하다.

 뭔 소리냐? x를 입력하면 뮤와 시그마를 출력하는 신경망을 만드는 것. 이 신경망 하나를 학습시키면 각 파라미터를 모두 학습시키는 것보다 훨씬 간단해진다.(amortized는 '공통으로, 한꺼번에'를 의미한다.) 

 이렇게 만든 q(z|x)에서 z를 sampling하고 decoder를 통해 x 틸다를 만들어 낸다.

 그럼 실제로 reconstruction term을 어떻게 계산하는지 알아보자.

 VAE 역시 신경망이다 보니 differentiable이 필수이다. 하지만 sampling process는 일반적으로 non-differentiable이다. 따라서 우린 reparameterization trick을 통해 이 과정을 미분가능하게 할 것인데 z를

로 sampling하고 gradient를 계산하는 것이다. z의 prior이 N(0,1)으로 unbiased하고 minibatch size가 클 때 이렇게 one sample로 gradient를 계산하여도 동일한 효과를 내는 것은 수학적으로 밝혀져있다.(이부분은... 납득!) 

 이후 decoder 과정에서 x 틸다를 N(x|x*, I)에서 sampling하는데(시그마를 1로 가정) 이 때

가 성립한다.(x_{rec} = x 틸다)

 왜냐하면 가우시안 분포가

이렇게 생겼는데 log를 씌우면

 이고 앞 항은 constant여서

이다. 즉, 의미를 분석하면 reconstruction term은 x와 x 틸다 사이의 유클리드 거리의 제곱이 최소화되도록 학습하는 것이다.

 이번엔 regularization term을 살펴보자.

 p(z)는 일반적으로 N(0,1)로 가정하고 q(z)는 z의 posterior 인데 이 역시 Gaussian distribution이므로 위 식은 closed form이 존재해 생각보다 쉽게 계산된다.

 즉, sampling을 포함한 전체 process를 보면

가 된다. 이제 각 term이 어느 층과 연결되는지까지 표현하면

 이렇게 되고 역전파를 나타내면

 와

이다. 결론적으로 우리는 latent variable z를 통해 새로운 것을 'generation' 할 수 있게 되었다.

+

이 주석을 봤을 때 의문이 들어야 하는 것이 'x 틸다를 sampling한다.'는 표현이다. 그럼,

이 backporpagation에서 x 틸다는 sampling process를 거치기 때문에 미분이 가능한지 의구심을 들게 만든다.

결론부터 말하면 둘다 맞다.

 우선 미분이 가능한 이유는 deocder는 주어진 z에 대해 하나의 결정적인 출력을 만든다. 즉, 주어진 z에 대해 항상 동일한 x 틸다를 생성하는 것이다.(x 틸다 = g(z)) 그럼, sampling 했다고 표현할 수 있는지 생각이 들 수 있는데 x가 잠재 변수 z로부터 생성될 때 가우시안 분포로부터 생성된다고 가정했기 때문에 하나의 sample을 추출한 것으로 해석할 수 있다. 당연히 정규 분포에서 가장 가능성이 높은 평균을 sample하는 것이 자연스럽기 때문에 이 x틸다는 z로부터 출력한 p(x|z_ 분포의 평균이고 이를 통해 reconstruction term을 계산한다. 

 정리하면 z를 통해 x가 바로 결정되고 x는 p(x|z)에서 하나 sampling 된 것으로 볼 수 있다.

4. Hierachical Variational Auto-Encoder(HVAE)

 VAE에서 latent variable이 여러 개로 확장되고 각각이 depencency가 존재하는 모델이 HVAE이다. 이는 diffusion의 기반이 되니 다음 포스트에서 또 언급하겠다.

 5. Recurrent VAE

 VAE를 확장한 또다른 구조로 recurrent VAE가 있다. Hidden variables 사이에 recurrence가 존재하는 것이 특징이다.

6. Problems with VAE

1. inefficient to evaluate / measure the likelihood of observed data

: 애초에 위 과정이 모두 이루어진 것이  p(x|z)p(z)dz가 적분이 되지 않아 이런 먼길을 돌아온 것이다. 적분을 학술이는 방법이 있긴 하지만(MCMC(Markov Chain Monte Carlo) 방법은 복잡한 확률 분포에서 샘플링하여 적분 문제를 해결하는 강력한 수치적 방법) 계산 과정이 매우 많아 비효율적이다.

2. not perfect samples from the distribution

: 1번 때문에 간단한 함수로 근사하는 variational inference를 하지만 '간단한' 함수이기 대문에 생성하는데 완벽하지 않다. 즉 normal distribution으로 복잡한 분포는 근사하기에는 충분하지 않다. 실제로 생성해보면 over smoothing한 x 틸다가 아래처럼 생긴다.

3. posterior collapse with strong decoder

: 만약 decoder가 너무 강하면(capacity가 크다면) 잠재변수 z를 무시하고 그냥 x로부터 x틸다를 구성해버린다. Encoder를 사용하지 않지만 마치 학습되는 것처럼 보이기 대문에 모델의 표현력이 개선되지 않는다.

7. 마무리

 이렇게 저번 기말 범위 중 하나였던 VAE에 대한 내용을 마치고 다음 포스트는 HVAE에 제약을 더한 diffusion에 대해 올려보겠다. (지금부터 해야 되긴해..)

 이번엔 기말고사 전에 리뷰했던 U-Net을 구현해보는 시간을 가져볼 것이다. 리뷰했을 때에 언급했지만 U-Net은 biomedical dataset에 특화된 모델이어서 sub task 역시 bio에 관련된 것으로 가져왔다.

1. Task와 Dataset 소개

 오늘의 dataset은 아래에 올라와있다.

https://www.kaggle.com/datasets/andrewmvd/cancer-inst-segmentation-and-classification 

 전체 데이터셋은 크게 3개로 나눠있는데 하나의 part만 가져와도 256*256의 이미지가 2656개가 있어서 충분하다고 판단해(사실 충분한단 것은 gpu memory의 한계를 의미하는 것과 같다.) part 1만 사용하였다.
  Author Notes의 README를 읽어보면 19개의 tissue types의 데이터로 이루어져있으며 각각은 0: Neoplastic cells, 1: Inflammatory, 2: Connective/Soft tissue cells, 3: Dead Cells, 4: Epithelial, 6: Background 이렇게 마스킹 되어있다.
 용어를 하나하나 살펴보면
 

Neoplastic cells : 종양을 형성하는 세포
Inflammatory : 염증성
Connective/Soft tissue cells: 결합조직 및 연조직으로 신체의 구조와 지지 기능을 담당하는 세포들
Dead Cells: 죽은 세포, 더 이상 기능하지 않는 세포
Epithelial: 상피 세포, 신체의 표면과 내부 장기를 덮고 있는 세포
 

인데 오랜만에 생물 공부를 하는 느낌이다.
 그리고 추가적으로 데이터 형식을 살펴보면 특이하게 images와 masks가 images.npy, masks.npy로 되어 있다는 것이다. 
(npy 파일은 NumPy 라이브러리에서 사용하는 파일 형식으로, 다차원 배열 데이터를 효울적으로 저장하고 로드하는데 사용된다.)

import numpy as np

images = np.load(r'C:\Users\james\Desktop\U-net\Part 1\Part 1\Images\images.npy', mmap_mode='r')
images = images.astype('int32')

masks = np.load(r'C:\Users\james\Desktop\U-net\Part 1\Part 1\Masks\masks.npy', mmap_mode='r')
masks = masks.astype('int32')

print(images.shape)
print(masks.shape)

 (코드의 mmap_mode 매개변수는 매모리 맵핑을 사용하여 대용량 배열 데이터를 디스크에서 직접 읽는 방식을 지정하는데 사용한다. 이 값을 기본값으로하게 되면 매모리 맵핑을 사용하지 않아 파일 전체가 메모리에 로드되는데 파일 양이 커서 local의 경우 과부화가 걸린다. 따라서 읽기 전용 모드로 파일을 메모리에 매핑하여 데이터의 일부만 로드하고 필요한 부분만 읽게 하였다. 즉, 메모리 사용량을 줄이고 입출력 성능을 향상시키기 위해 하는 것!)

(2656, 256, 256, 3)
(2656, 256, 256, 6)

 출력해보면 masks의 채널이 왜 6인지는 위에 데이터셋을 보면 알 수 있을 것이다. 일전의 FCN에서 했던 cloth segmentation의 경우 semantic segmentation이기 때문에 2차원 mask에서 각 pixel의 값에 class의 number를 매핑하면 충분하였다. 하지만 이번 task는 instance segmentation을 해야하기 때문에 3차원 mask에 각 채널이 하나의 class를 의미하고 각 class에 해당하는 instance를 숫자로 매핑하게 되었다.
 Semantic segmentation과 instance segmentation이 무엇이지 까먹었다면 이 곳으로

2024.04.26 - [DL/Image Segmentation] - [Paper Review] Fully Convolutional Netowrks for Semantic Segmentation

 그럼 task와 데이터 소개를 마쳤으니 본격적인 구현 단계로 넘어가 보겠다.

2. Dataset  정의 및 Dataset 로딩

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image

class NumpySegDataset(Dataset):
    def __init__(self, images_path, masks_path, transform=None, target_transform=None):
        self.images = np.load(images_path, mmap_mode='r')
        self.masks = np.load(masks_path, mmap_mode='r')
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        image = self.images[idx]
        mask = self.masks[idx]

        if self.transform:
            image = self.transform(image)

        if self.target_transform:
            mask = self.target_transform(mask)

        # image는 정규화하면 소수로 변하니까 float, mask는 정규화 안하니까 int
        image = torch.tensor(image, dtype=torch.float32).permute(2,0,1)
        mask = torch.tensor(mask, dtype=torch.int64).permute(2,0,1)

        return image, mask

 항상 해왔듯이 Dataset class인 NumpySegDataset class를 정의하고 __init__,. __len__함수와 __getitem__ 함수를 정의하였다. 다만 저번 미니 프로젝트와는 다른 점은 __init__ 함수를 만들 때 numpy 배열이기 때문에 np.load와 mmap_mode = 'r'이라는 점이다.

images_path=r'C:\Users\james\Desktop\U-net\Part 1\Part 1\Images\images.npy'
masks_path=r'C:\Users\james\Desktop\U-net\Part 1\Part 1\Masks\masks.npy'

dataset = NumpySegDataset(images_path, masks_path)

 이렇게 dataset을 정의해주었고

total_len = len(dataset)
train_len = int(total_len * 0.8)
val_len = total_len - train_len

train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_len, val_len])
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=4, shuffle=True)

 dataset을 로딩해줄 때 8:2의 비율로 train과 validation을 구분하였다.
 (batch size가 4로 매우 작은 편인데 batch size를 좀만 늘려도 OOM(Out of Memory) 문제가 발생하여서 batch size와 epoch를 모두 줄이게 되었다.)

 3. 모델 architecture 정의

 먼저 참고한 사이트는 다음과 같다.
https://github.com/meetps/pytorch-semseg/blob/master/ptsemseg/models/unet.py

https://www.kaggle.com/code/pedroamavizca/working-with-u-net

 이번에 U-Net class를 만들 때 구현하고자 하는 내용은 다음과 같다.

1. U-Net 구조
2. pixel-wise weight와 cross entropy를 결합한 custom_loss
3. mirroring extrapolation

 그래서 먼저 U-Net 구조를 만들면

이 구조를 만들어야 하기 때문에 반복적으로 있는 conv 3*3, ReLU(위 그림에서 파란색 2개)를 하나의 class로 정의해보겠다.

class UNetConv2(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNetConv2, self).__init__()
        # (입력크기+2*패딩-커널크기)/스트라이드+1
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU()
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        return x

 그 다음 이를 활용하는 전체 U-Net class를 정의하면

class UNet(nn.Module):
    def __init__(self, num_classes=6, in_channel=3):
        # 부모 클래스 초기화되어 자식 클래스에서도 사용 가능
        super(UNet, self).__init__()
        self.conv_1 = UNetConv2(in_channel, 64)
        self.conv_2 = UNetConv2(64, 128)
        self.conv_3 = UNetConv2(128, 256)
        self.conv_4 = UNetConv2(256, 512)

        self.mid_conv = UNetConv2(512, 1024)

        self.conv_5 = UNetConv2(1024, 512)
        self.conv_6 = UNetConv2(512, 256)
        self.conv_7 = UNetConv2(256, 128)
        self.conv_8 = UNetConv2(128, 64)

        self.down = nn.MaxPool2d(kernel_size=2, stride=2)
        self.up_1 = nn.ConvTranspose2d(1024, 512, kernel_size=2, stride=2)
        self.up_2 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.up_3 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.up_4 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)

        self.end = nn.Conv2d(64, num_classes, kernel_size=1, stride=1)

    def forward(self, x):
        padded_x = F.pad(x, (92, 92, 92, 92), mode='reflect')
        conv_1 = self.conv_1(padded_x) # output: 252*252
        if conv_1.size()[2] % 2 != 0:
            conv_1 = F.pad(conv_1, (0, 1, 0, 1))
        pool1 = self.down(conv_1) # output: 126*126

        conv_2 = self.conv_2(pool1) # output: 122*122
        if conv_2.size()[2] % 2 != 0:
            conv_2 = F.pad(conv_2, (0, 1, 0, 1))
        pool2 = self.down(conv_2) # output: 61*61

        conv_3 = self.conv_3(pool2) # output: 57*57
        if conv_3.size()[2] % 2 != 0:
            conv_3 = F.pad(conv_3, (0, 1, 0, 1))
        pool3 = self.down(conv_3) # output: 29*29

        conv_4 = self.conv_4(pool3) # output: 25*25
        if conv_4.size()[2] % 2 != 0:
            conv_4 = F.pad(conv_4, (0, 1, 0, 1))
        pool4 = self.down(conv_4) # output: 13*13

        mid_conv = self.mid_conv(pool4) # output: 9*9

        up_1 = self.up_1(mid_conv)
        scale_idx_1 = (conv_4.shape[2] - up_1.shape[2]) // 2
        cropped_conv_4 = conv_4[:, :, scale_idx_1:-scale_idx_1, scale_idx_1:-scale_idx_1]
        up_1 = torch.cat([up_1, cropped_conv_4], dim=1)
        conv_5 = self.conv_5(up_1)

        up_2 = self.up_2(conv_5)
        scale_idx_2 = (conv_3.shape[2] - up_2.shape[2]) // 2
        cropped_conv_3 = conv_3[:, :, scale_idx_2:-scale_idx_2, scale_idx_2:-scale_idx_2]
        up_2 = torch.cat([up_2, cropped_conv_3], dim=1)
        conv_6 = self.conv_6(up_2)

        up_3 = self.up_3(conv_6)
        scale_idx_3 = (conv_2.shape[2] - up_3.shape[2]) // 2
        cropped_conv_2 = conv_2[:, :, scale_idx_3:-scale_idx_3, scale_idx_3:-scale_idx_3]
        up_3 = torch.cat([up_3, cropped_conv_2], dim=1)
        conv_7 = self.conv_7(up_3)

        up_4 = self.up_4(conv_7)
        scale_idx_4 = (conv_1.shape[2] - up_4.shape[2]) // 2
        cropped_conv_1 = conv_1[:, :, scale_idx_4:-scale_idx_4, scale_idx_4:-scale_idx_4]
        up_4 = torch.cat([up_4, cropped_conv_1], dim=1)
        conv_8 = self.conv_8(up_4)

        end = self.end(conv_8)
        scale_idx_5 = (end.shape[2]-x.shape[2]) // 2
        end = end[:, :, scale_idx_5:-scale_idx_5, scale_idx_5:-scale_idx_5]

        return end

 다음과 같은데 최대한 논문에서 말하는 구조를 따라하기 위해 padding을 0으로 하고 croppping하는 과정을 넣었는데 참고한 사이트에선 구현을 용이하기 위해 padding을 사용하였다. 그리고 논문에서 segmentation이 용이하게 되기 위해서는 max pooling을 하는 input의 size가 짝수여야 했는데 이미 이미지의 크기가 256*256으로 제한되어 있어서 이를 늘리거나 줄이지는 않고 홀수 일때만 padding을 추가하는 방식으로 구현하였다.
 torch.summary를 사용해 출력해보면 다음과 같다.

torch.Size([500, 256, 256])
torch.Size([500, 256, 256])
torch.Size([500, 256, 256])
torch.Size([500, 256, 256])
torch.Size([500, 256, 256])
torch.Size([156, 256, 256])
torch.Size([2656, 256, 256])

(2656, 256, 256, 3)
(2656, 256, 256, 6)
torch.Size([16, 3, 256, 256])
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 438, 438]           1,792
       BatchNorm2d-2         [-1, 64, 438, 438]             128
              ReLU-3         [-1, 64, 438, 438]               0
            Conv2d-4         [-1, 64, 436, 436]          36,928
       BatchNorm2d-5         [-1, 64, 436, 436]             128
              ReLU-6         [-1, 64, 436, 436]               0
         UNetConv2-7         [-1, 64, 436, 436]               0
         MaxPool2d-8         [-1, 64, 218, 218]               0
            Conv2d-9        [-1, 128, 216, 216]          73,856
      BatchNorm2d-10        [-1, 128, 216, 216]             256
             ReLU-11        [-1, 128, 216, 216]               0
           Conv2d-12        [-1, 128, 214, 214]         147,584
      BatchNorm2d-13        [-1, 128, 214, 214]             256
             ReLU-14        [-1, 128, 214, 214]               0
        UNetConv2-15        [-1, 128, 214, 214]               0
        MaxPool2d-16        [-1, 128, 107, 107]               0
           Conv2d-17        [-1, 256, 105, 105]         295,168
      BatchNorm2d-18        [-1, 256, 105, 105]             512
             ReLU-19        [-1, 256, 105, 105]               0
           Conv2d-20        [-1, 256, 103, 103]         590,080
      BatchNorm2d-21        [-1, 256, 103, 103]             512
             ReLU-22        [-1, 256, 103, 103]               0
        UNetConv2-23        [-1, 256, 103, 103]               0
        MaxPool2d-24          [-1, 256, 52, 52]               0
           Conv2d-25          [-1, 512, 50, 50]       1,180,160
      BatchNorm2d-26          [-1, 512, 50, 50]           1,024
             ReLU-27          [-1, 512, 50, 50]               0
           Conv2d-28          [-1, 512, 48, 48]       2,359,808
      BatchNorm2d-29          [-1, 512, 48, 48]           1,024
             ReLU-30          [-1, 512, 48, 48]               0
        UNetConv2-31          [-1, 512, 48, 48]               0
        MaxPool2d-32          [-1, 512, 24, 24]               0
           Conv2d-33         [-1, 1024, 22, 22]       4,719,616
      BatchNorm2d-34         [-1, 1024, 22, 22]           2,048
             ReLU-35         [-1, 1024, 22, 22]               0
           Conv2d-36         [-1, 1024, 20, 20]       9,438,208
      BatchNorm2d-37         [-1, 1024, 20, 20]           2,048
             ReLU-38         [-1, 1024, 20, 20]               0
        UNetConv2-39         [-1, 1024, 20, 20]               0
  ConvTranspose2d-40          [-1, 512, 40, 40]       2,097,664
           Conv2d-41          [-1, 512, 38, 38]       4,719,104
      BatchNorm2d-42          [-1, 512, 38, 38]           1,024
             ReLU-43          [-1, 512, 38, 38]               0
           Conv2d-44          [-1, 512, 36, 36]       2,359,808
      BatchNorm2d-45          [-1, 512, 36, 36]           1,024
             ReLU-46          [-1, 512, 36, 36]               0
        UNetConv2-47          [-1, 512, 36, 36]               0
  ConvTranspose2d-48          [-1, 256, 72, 72]         524,544
           Conv2d-49          [-1, 256, 70, 70]       1,179,904
      BatchNorm2d-50          [-1, 256, 70, 70]             512
             ReLU-51          [-1, 256, 70, 70]               0
           Conv2d-52          [-1, 256, 68, 68]         590,080
      BatchNorm2d-53          [-1, 256, 68, 68]             512
             ReLU-54          [-1, 256, 68, 68]               0
        UNetConv2-55          [-1, 256, 68, 68]               0
  ConvTranspose2d-56        [-1, 128, 136, 136]         131,200
           Conv2d-57        [-1, 128, 134, 134]         295,040
      BatchNorm2d-58        [-1, 128, 134, 134]             256
             ReLU-59        [-1, 128, 134, 134]               0
           Conv2d-60        [-1, 128, 132, 132]         147,584
      BatchNorm2d-61        [-1, 128, 132, 132]             256
             ReLU-62        [-1, 128, 132, 132]               0
        UNetConv2-63        [-1, 128, 132, 132]               0
  ConvTranspose2d-64         [-1, 64, 264, 264]          32,832
           Conv2d-65         [-1, 64, 262, 262]          73,792
      BatchNorm2d-66         [-1, 64, 262, 262]             128
             ReLU-67         [-1, 64, 262, 262]               0
           Conv2d-68         [-1, 64, 260, 260]          36,928
      BatchNorm2d-69         [-1, 64, 260, 260]             128
             ReLU-70         [-1, 64, 260, 260]               0
        UNetConv2-71         [-1, 64, 260, 260]               0
           Conv2d-72          [-1, 6, 260, 260]             390
================================================================
Total params: 31,043,846
Trainable params: 31,043,846
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.75
Forward/backward pass size (MB): 1773.24
Params size (MB): 118.42
Estimated Total Size (MB): 1892.42
----------------------------------------------------------------

4. 손실함수와 optimizer 정의

 우선 model과 optimizer는 쉽게 정의했는데

model = UNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

 일반적인 loss fucntion을 사용하는 것이 아닌 각 pixel에 weight를 부여하고 softmax함수와 cross entropy loss를 결합한 loss(논문에 잘 소개되어있다.)를 사용해보기 위해 함수들을 정의해보았다.

def custom_loss(outputs, labels, weights):
    # softmax 계산
    softmax_outputs = F.softmax(outputs, dim=1)

    # CPU로 이동
    labels = labels.cpu()
    weights = weights.cpu()
    softmax_outputs = softmax_outputs.cpu()

    # 0이 아닌 위치를 찾기 위한 마스크 생성
    non_zero_mask = labels != 0

    # 마스크를 사용하여 필요한 값 선택 및 계산
    selected_weights = weights.unsqueeze(1).expand_as(labels)[non_zero_mask]
    selected_softmax_outputs = softmax_outputs[non_zero_mask]

    # 손실 계산
    running_loss = (-1) * selected_weights * torch.log(selected_softmax_outputs)
    running_loss = running_loss.sum()

    running_loss /= labels.shape[0] * labels.shape[2] * labels.shape[3]

    return running_loss.to(outputs.device)

 먼저 outputs에서 softmax를 적용하고(채널이 class별 출력을 위하니 dim=1로 softmax를 계산한다.) labels에서 0이 아닌 위치를 나타내는 non_zero_mask를 정의하였다. weights의 차원은 (batch_size, height, width)이고 non_zero_mask의 차원은 (batch_size, num_classes, height, width)에서 selected_weights를 구하기 위해선 weights에서 두번째 차원을 추가하고 labels의 num_classes만큼 복사한다음 선택을 한다. 그렇게 계산한 selected_weights의 차원은 non_zero의 갯수가 N이라 하면 (N,)이다. 그리고 필는 loss function을 최소화하기 위해 -1을 붙였고 pixel-wise loss때문에 평균을 내기 위해 batch_size와 height, width로 나누어줬다. 
 그럼 weight는 어떻게 계산할까? 논문에서는

 이렇게 class의 frequency를 반영한 1차 weight에 세포와의 거리를 반영하는 추가적인 weight로 계산하였는데 구현이 엄청 어려운 것이지만 시간이 매우 오래 걸릴 것이라 판단하여 비슷한 아이디어로 정의하였다.

import torch

def find_others(labels, i, j, k, b, d):
    left = max(i - d, 0)
    right = min(i + d, 255)  # 256이 아니라 255까지
    up = max(j - d, 0)
    down = min(j + d, 255)  # 256이 아니라 255까지
    instance = labels[b, k, i, j]

    region = labels[b, k, left:right+1, up:down+1]
    other_classes = (region == 0).sum().item()
    other_instances = ((region != 0) & (region != instance)).sum().item()

    return other_classes, other_instances

def calculate_weights(masks):
    device = masks.device
    batch_size, num_classes, height, width = masks.shape
    weights = torch.zeros((batch_size, height, width), device=device)
    non_zero_counts = (masks != 0).sum(dim=(2, 3))

    for b in range(batch_size):
        non_zero_ratio = non_zero_counts[b].float() / non_zero_counts[b].sum(dim=0, keepdim=True).float()
        exp_non_zero_ratio = torch.exp(-non_zero_ratio)
        
        for k in range(num_classes):
            mask_k = masks[b, k]
            non_zero_mask = mask_k != 0
            weights[b][non_zero_mask] = exp_non_zero_ratio[k]

            for i in range(2, height, 5):
                for j in range(2, width, 5):
                    if non_zero_mask[i, j]:
                        other_classes, other_instances = find_others(masks, i, j, k, b, 2)
                        weights[b, i-2:i+3, j-2:j+3] *= (1.02)**other_classes
                        weights[b, i-2:i+3, j-2:j+3] *= (1.05)**other_instances

    return weights

 기본적인 아이디어는 먼저 class의 frequency의 비율에 exp(-x)를 적용하고 해당 pixel을 가운데로 하는 25개의 pixel에 같은 class지만 다른 instance인 pixel의 개수(a)와 다른 class인 pxiel의 개수(b)에 따라 가중치를 각각 1.05**(a), 1.02**(b) 배 해주는 방식이다. 총 2656개의 데이터의 256*256 pixels는 1억개가 넘어서(174,063,616) 다섯 pixels씩 넘어가면서 계산하도록 설계하였다. 
 matplotlib의 pyplot을 사용하여 하나의 weight를 시각해보면 다음과 같다.

 경계에서, 특히 다른 instance의 경계에서 가장 높은 weight를 갖도록 설계된 것을 알 수있다.

5. Train & Validation

num_epochs = 50  # Number of epochs
batch_size = 4
val_idx_start = len(train_loader.dataset)

for epoch in range(num_epochs):
    print(f'Epoch {epoch+1}/{num_epochs}')

    # Each epoch has a training and validation phase
    for phase in ['train', 'val']:
        if phase == 'train':
            model.train()  # Set model to training mode
            dataloader = train_loader
        else:
            model.eval()   # Set model to evaluate mode
            dataloader = val_loader

        running_loss = 0.0

        # Iterate over data with tqdm for the progress bar
        progress_bar = tqdm(enumerate(dataloader), total=len(dataloader), desc=f"{phase.capitalize()} Phase")
        for batch_idx, (inputs, labels) in progress_bar:
            inputs, labels = inputs.to(device), labels.to(device)

            # Zero the parameter gradients
            optimizer.zero_grad()

            # Forward
            with torch.set_grad_enabled(phase == 'train'):
                outputs = model(inputs)

                if phase == 'train':
                    batch_weights = weights[batch_idx * batch_size : (batch_idx + 1) * batch_size]
                else:
                    batch_weights = weights[val_idx_start + batch_idx * batch_size : val_idx_start + (batch_idx + 1) * batch_size]

                loss = custom_loss(outputs, labels, batch_weights)

                if phase == 'train':
                    loss.backward()
                    optimizer.step()

            # Statistics
            running_loss += loss.item() * inputs.size(0)
            epoch_loss = running_loss / len(dataloader.dataset)

            # Update the progress bar with the current loss value
            progress_bar.set_postfix({'loss': f'{loss.item():.4f}'})

        if phase == 'val':
            print(f'{phase.capitalize()} Loss: {epoch_loss:.4f}')
        print()

 GPU 메모리와 사용량의 한계로 epoch를 50으로 설정하였고 위에서 계산한 weight를 batch에 맞게 가져오도록 설계하였다. 훈련은 총 3시간 정도 걸렸다.

 결과를 시각화해보면

 이렇게 생각보다 output이 labels를 잘 분류하는 것을 알 수 있다. 그런데 필자의 최종 목표는 instance segmentation이기 때문에 watershed 알고리즘을 적용하였다. 아래 그림은 바로 위 그림의 첫번째 output에 적용한 결과인데

  노란색 타원으로 표시한 부분처럼 touching instances들도 구분한 것을 알 수 있다.
+ watershed 알고리즘
 객체를 분할하는데 사용되는 기법으로 지형학적 모델을 사용하여 이미지의 픽셀을 분할하는기 때문에 outputs의 각 클래스별 출력이 높이로 간주된다. 이 알고리즘은 물이 채워질 때 계곡을 따라 경계가 형성된다는 개념에서 유래됐다. 알고리즘의 단계로는 
1. 전처리 - threshold 보다 낮은 값들은 제거한다.
2. 거리 변환 -  객체 내부의 각 픽셀이 가장 가까운 배경 픽셀로부터 얼마나 떨어져 있는지 계산한다. 여기서 거리는 Manhattan distance를 의미한다. 
3. 마커 생성 - 지역 극대값을 찾아 마커로 설정한다. 마커는 객체의 중심을 나타낸다.
4. watershed 변환 - 마커에서 시작하여 물을 채워 나가, 서로 다른 마커에서 채워진 물이 만나는 지점에서 경계가 형성된다.
 가 있고 그 결과 위 그림처럼 instance segmentation이 수행된다.

 6. 마무리

 위 전체 코드를 구현한 것은 깃허브에 올려두었다.
https://github.com/ParkSeokwoo/U-Net-cancer-instance-segmentation-

 
 Watershed 알고리즘을 진작에 알았다면 loss를 설계하는 과정이 더 논문에 가까워졌을 것 같지만 이 정도만으로 만족하고 넘어가려한다. 그리고 이건 위 모델하고는 상관없는데 생각보다 kaggle notebook이 괜찮은거 같다. Colab pro를 결제해도 다 쓰는데 1주일 밖에 걸리지 않는데 kaggle notebook은 전화번호가 있는데로 쓸 수 있고 1주일마다 다시 30시간을 사용할 수 있기 때문에 앞으로 대부분의 작업을 kaggle에서 사용할 것 같다. 다음 포스트는 동아리 사람들과 diffusion 모델 스터디를 하게 되어서 스터디 준비용 포스트를 올릴 것 같다. 바이~

 이번 포스트는 U-Net을 소개하는 논문을 리뷰해보겠다. 논문은 다음 사이트에 올라와있다.

 https://arxiv.org/abs/1505.04597

 이번 역시 동아리에서 발표했던 주제라 그 때 사용한 피피티를 첨부하겠다.

Abstract

< 내용 >

 Deep nets가 성공적으로 훈련되기 위해서는 많은 annotated training samples가 필요하다고 알려져있다. 우리는 이 논무네서 이용가능한 annotated samples를 더 효율적으로 사용하기 위해 data augmentation을 활용하는 훈련 전략과 network를 소개하려 한다. 그 구조는 context를 캡처하는 contracting path와 정밀한 localization을 가능하게 하는 대칭적인 expanding path로 이루어져있다. 우리는 이 network가 end-to-end 방식으로 학습 가능한데 매우 작은 이미지더라도 기존의 방법인 sliding-window 방식을 능가하는 것을 보일 것이다.(데이터는 ISBI challenge for segmentation of neural structures in electron microscopic stacks) 같은 네트워크를 transmitted light microscopy images에도 훈련하여 ISBI cell tracking challenge 2015에서 가장 좋은 성능을 보였다. 게다가 이 네트워크는 바르다. 512*512 images를 GPU를 사용하여 inference하는데 일 초 이내의 시간이 걸린다.   

< comment >

 저자들이 소개하는 U-net은 기본적으로 FCN의 구조를 따르지만 몇가지 차이점이 존재한다. 대표적인 예시가 contracting path와 expanding path가 symmetric하다는 점인데 뒤에서 자세히 소개하도록 하고 U-Net에서 또 주목할만한 부분은 적은 이미지로도 높은 성능을 달성하기 위해 data augmentation을 사용했다는 점이다. 이는 biomedical image dataset의 특성과 관련되는데 이 dataset은 이미지의 크기가 다양하면서 대부분 크고 개수가 적다는 특징이 있다. 이를 보안하기 위해 U-Net이 어떻게 설계되었는지 차차 알아가보자.

1. Introduction

< 내용 >

 지난 2년동안 deep convolutional networks는 많은 visual recogniton tasks에서 sota를 달성했다. 이러한 성공이 오랜시간 지속되었지만 그들의 성공이 제한되었던 것은 이용가능한 training set의 사이즈와 네트워크 자체의 크기 때문이다. Krizhevsky의 돌파구로 800만 개의 매개변수를 가진 8개의 레이어로 구성된 대규모 네트워크를 ImageNet 데이터셋의 100만 개의 훈련 이미지로 supervised training한 사례가 있고 그 이후로, 더 크고 깊은 네트워크들이 훈련되었다.

 일반적으로 convolutional networks가 사용된 곳은 classification tasks이고 이는 이미지를 단일 클래스 label로 분류하는 문제이다. 그러나 많은 visual tasks, 특히 biomedical image processing은 이러한 output에 localization을 요구하였고 이는 우리가 아는 segmentation으로 각 pixel마다 class label을 assign하는 문제이다. 게다가 수천여장의 이미지가 주로 biomedical tasks에 존재한다.(확실히 적은 수치이다.) 따라서 Ciresan은 sliding-window setup을 통해 네트워크를 훈련하여 각 픽셀별 class label을 예측하고자 하였고 픽셀을 둘러싼 영역(이미지의 작은 부분인 패치)를 입력으로 삼았다. 먼저 이러한 네트워크는 localize가 가능했고 다음, patch를 입력으로 받기 때문에 training data의 개수는 images 전체를 입력받을 때보다 많았다. 그 결과 이 네트워크는 ISBI 2012 EM segmentation challenge 에서 큰 차이로 우승하였다.

 확실이 Ciresan의 전략은 두 가지 단점이 존재한다. 먼저, 네트워크가 각 패치마다 개별로 진행되어야 하기 때문에 매우 느렸고 patches 간의 overlapping으로 매우 많은 redundancy가 존재했다. 다음으로, localization accuracy와 use of context사이에 trade-off가 존재했다. 큰 패치는 많은 max-pooling layers를 요구하여 localization accuracy를 낮추었고, 반면 작은 패치는 네트워크가 작은 context만 보게 하였다. 더 최근 접근방법에는 classifier output을 multiple layers에서 설명하는 방법이 있어 good localization과 use of context가 동시에 가능하도록 한다.

 이 논문에서 우리는 더 세련된 구조인 fully convolutional network를 만들었다. 우리는 이 구조를 수정하고 확장하여 매우 작은 training images에서 precise segmentation을 수행하게 하였다. 

 FCN의 메인 아이디어(저번 포스팅에 올렸던 논문)는 일반적인 contracting network를 pooling operators를 upsampling operators로 대체한 succesive layers로 보충하는 것이다. 따라서 이 층은 output의 해상도를 중가시키니다. Localize를 위해서 contracting path의 high resolution feautres가 upsampled output과 결합한다. 이 연속적인 convolution layer은 정보의 output들을 ensemble 하여 학습한다.

 우리의 구조에서 중요한 수정은 upsampling part에서도 큰 숫자의 feature channels를 가졌다는 것이다. 이는 network가 higher resolution layers로 context information을 전달하도록 한다. 그 결과 expansive path는 거의 contracting path와 symmetric하여 u-shaped architecture을 구성하게 된다. 이 네트워크는 fully connected layers를 갖지 않고 각 합성곱의 유용한 부분만 사용한다. 즉, segmentation map은 입력 이미지에서 전체 context가 사용 가능한 pixels만 포함되는데 overlap-tile strategy를 추가로 사용하여 임의의 큰 이미지에서 매끄러운 segmentation이 가능하게 한다. 이미지의 경계에 해당하는 pixels를 예측하기 위해 missing context를 input image에 mirroring을 통해 extrapolate한다. 이 tiling strategy는 네트워크에 큰 이미지를 적용하는데 중요한데 왜나하면 이렇게 하지 않는 경우 GPU memory에 의해 제한될 수 있기 때문이다.

우리의 tasks는 매우 작은 training data가 이용가능하기 때문에 우리는 강력한 data augmentation인 elastic deformation을 training images에 수행한다. 이것은 network가 이러한 deformations에 대해 invariance를 갖게 만든다. 이는 특히 biomedical segmentation에서 아주 중요한데 왜냐하면 tissue에는 이러한 변형들이 자주 일어나고 실제의 변형들을 효과적으로 재현할 수 있기 때문이다. Unsupervised feature learning의 범위에서 learning invariance에 대한 data augmentation의 가치는 Dosovitsky에 의해 알려져있다.

 많은 cell segmentation에서 또다른 어려움은 same class의 접촉하는 objects를 분리하는 것이다. 마지막에 우리는 weigted loss를 제안하여 touching cells의 사이를 background labels로 여겨 분리하고 이 부분에 높은 가중치를 부여하는 loss function을 만들 것이다.

 결과적으로 생성된 network는 다양한 biomedical segmentation problems에 적용가능하다. 이 논문에서 우리는 EM stacks(an ongoing competition started at ISBI 2012)에서 segmentation of neuronal structures의 결과를 보여줄 것이고 우리는 이 대회에서 Ciresan의 결과를 뛰어넘었다. 뿐만 아니라, ISBI cell tracking challenge 2015의 light microscoy images cell segmentation에서의 결과를 보여줄 것이다. 우리는 2D trasmitted light datasets에서 큰 차이로 우승하였다.

< comment >

 저자들은 FCN의 구조를 확장하여 더 정확한 segementation이 가능한 구조를 만들었다. Expanding path(upsampling path)에 연속적인 layers를 보충하였고 많은 수의 채널을 갖게 설계하였다. 그 결과 두 경로가 대칭인 u-shape의 구조를 갖게 되었다. 또한 padding 없이 convolution을 수행하는 것이 특징인데 그럼 layer를 통과할 수록 크기가 줄어들어 이를 보충하기 위해 mirroring extraplation을 사용하였다. 뿐만 아니라 elastic deformation을 통한 data augmentation을 통해 작은 수의 이미지로도 충분한 성능을 내도록 하였고 instance segmentation을 수행하기 위해 touching cells 사이를 background처럼 인식하고 가중치를 부여하였다.

2. Network Architecture

< 내용 >

네트워크의 구조는 우와 같이 묘사되며 왼쪽의 contracting path와 오른쪽의 expanding path로 이루어진다. Contracting path는 일반적인 convolutional network의 구조를 따라한다. 이는 반복되는 두개의 3*3 convolutions(unpadded convolutions)로 이루어져있고 각각은 ReLU를 이어받는다. 그리고 2*2 max pooling operation with stride 2가 convolution layers 사이에 존재해 downsampling한다. 한번의 downsampling이 일어나면 feature channels는 두배가 된다. Expansive path의 모든 step에는 feature map의 채널 수를 절반으로 줄이는 2*2 합성곱 을 통해 upsampling하고 대응되는 크기로 잘라낸 contractin path의 결과와 연결된 후 두번의 3*3 합성곱(각각은 ReLU를 사용)으로 구성된다. . Cropping이 필수인데 every convolution에서 border pixels를 잃기 때문이다. 마지막 레이어의 1*1 convolution은 64개의 component feature vector을 요구되는 class의 숫자에 맞게 mapping하는데 쓰여 총 network는 23개의 convolution layers를 갖는다.

 Segmenatation map의 output이 매끄럽기 위해선 2*2 max-pooling operations에 들어가는 input tile size가 짝수여야 한다.

< comment >

구조는 contracting path와 expanding path로 이루어져 있고 segmentation이 매끄럽게 수행되기 위해선 2*2 max-pooling operations에 들어가는 tile의 사이즈가 짝수여야한다.

3. Training

< 내용 >

 Input image와 대응되는 segmentation maps는 Caffe에 구현된 stochastic gradient descent를 통해 훈련된다. Unpadded convolution 때문에 output image는 input image보다 작다.(by a constant border width) GPU memory의 overhead를 최소화하고 gpu를 최대로 사용하기 위해서 우리는 large input tiles를 large batch size보다 선호한다. 따라서 batch size를 1로 하여 단일 이미지를 사용하였다. 우리는 high momentum인 0.99를 사용하여 이전에 보았던 training samples들의 대부분이 current optimization step의 update에 사용되도록 하였다.

 Energy function은 마지막 feature map의 pixel-wise soft-max와 cross entropy loss function을 결합한 함수를 사용하였다. soft-max는 $p_k(\mathbf{x})=\textrm{exp}(a_k(\mathbf{x}))/(\sum_{k'=1}^{K}\textrm{exp}(a_k'(\mathbf{x})))$로 정의되고

$\textrm{exp}(a_k'(\mathbf{x}))$ 는 x( $\textbf{x}\in\Omega\space\text{with}\space\Omega\subset\mathbb{Z}^2$ )의 feature channel k의 activation을 의미한다. K는 classes의 number를 의미하고 $p_k(\mathbf{x})$는 approximated maximum function이다.

*approximated maximum function

일반적으로 cross-entropy loss에 log안에 들어가는 값은 실제 분포를 추정한 모델의 확률로 U-net의 마지막 layer의 출력값이다. 그러나 특정 layer의 값을 단순히 대입하는 것은 미분이 불가능하기 때문에 미분가능한 값으로 만들기 위해 approximation funtion을 사용한다.

즉 이 값이 1에 가깝다는 것은 대응되는 k에 대해 maximum activation을 갖는 것이고 다른 k에 대해서는 이 값이 0에 가깝다. Cross entropy는 $-\sum_{i=1}^np(x_i)log(q(x_i))$인데 class segmentation인 $p(x_i)$가 실제 레이블에서만 1이므로 최대화하려는 energy funtion은 다음과 같이 정의된다.

$$E=\sum_{\mathbf{x}\in\Omega}w(\mathbf{x})log(p_{l(x)}(\mathbf{x})) $$

여기서 l은 each pixel의 true label을 의미하고 w는 weight map으로 몇몇의 pixels에 더 중요하게 생각하기 위해 도입한 것이다.

 먼저 우리는 wegith map을 certatin class의 different frequency of pixels를 반영하여 보상하기 위해 pre-compute한다. 이후에 가중치를 update하여 touching cells 사이에 있는 small separation borders도 학습할 수 있게 한다.

가중치 update는 morphological operations(형태학적 작용)을 통해 계산되는데 식은 다음과 같다.

$$w(\mathbf{x})=w_c(\mathbf{x})+w_0\cdot\textrm{exp}(-\frac{(d_1(\mathbf{x})+d_2(\mathbf{x}))^2}{2\sigma^2})$$

$w_c$는 class frequencies를 balance한 weight map이고 $d_1$은 가장 가까운 셀의 경계 까지의 거리, $d_2$는 두번째 가까운 셀의 경계까지의 거리이다. 우리의 실험에는 $w_0=10$, $\sigma\approx5$ pixels이다.

 Deep networks에서 많은 convolution layers와 network를 통과하는 다양한 경로가 존재하면 weights의 초기 initialization은 매우 중요하다. 그렇지 않으면 네트워크의 일부가 과도하게 활성화되어 다른 부분은 전혀 기여하지 않을 수 있다. 이상적으로는 네트워크의 각 피처 맵이 단위 분산을 갖도록 초기 가중치를 정해야 한다. 교대로 발생하는 convolution 과 ReLU layers가 있는 우리의 구조에서는 이 목표를 달성하기 위해 표준 편차가 $\sqrt{2/N}$인 가우시안 분포에서 초기 가중치를 추출하였다. 여기서 N은 한 뉴런의 입력 노드 수를 의미한다. 예를 들어 이전 층에 64개의 채널이 있는 3*3 합성곱의 경우 N은 9*64=576이다.

3.1 Data Augmentation

Data augmentation은 매우 적은 training samples가 이용가능할 때 network에 invariance와 robustness properties를 가르치는데 필수적이다. Microscopial images의 경우 우리는 주로 여러 변형에 대한 robustness 뿐만아니라 shift와 rotation invariance, gray value varations에 대한 invariance 까지 필요하다. 특히 random elastic deformations를 training sample에 적용하는 것이 매우 적은 annotated images에서 segmentation을 수행하는 network를 훈련시킬 때 주요해보인다. 우리는 3*3 격자에서 랜덤한 변위 벡터를 사용하는 smooth deformations를 생성하였다. 변위는 표준 편차가 10 픽셀인 가우시안 분포에서 샘플링하였다. 각 픽셀의 변위는 bicubic interpolations에 의해 계산된다. 수축 경로 끝에 있는 Drop-out layers는 암묵적인 data augmentation을 추가적으로 수행한다.=

< comment >

U-Net은 biomedical image의 사이즈가 크다는 특성 때문에 patchwise training을 사용한다. 이 과정에서 큰 input tiles를 선호해 batchsize는 최소화하여 1을 사용하였다. 또한 optimize할 때 momentum을 0.99인 큰 값을 사용하여 이전 이미지가 현재에 대부분 영향을 주도록 설계하였다.

 Loss funtion은 pixel-wise softmax function과 cross-entropy loss를 결합하여 설계하였고 weight를 도입하여 class간의 frequency, touching cells 사이에 위치했는가를 반영하였다. 그리고 중요한 초기 가중치 설계를 입력 노드를 고려한 가우시안 분포에서 추출하였다.

4. Experiments

< 내용 >

< comment >

U-Net은 다양한 bio-medical image segmentation에서 좋은 성능을 보인다.

5. Conclusion

< 내용 >

 U-net 구조는 다양한 biomedical segmentation applications에서 좋은 성능을 달성하였다. Elastic deformations를 사용하는 data augmentation에 힘입어 매우 작은 annotated imges만 필요하고 타당한 training time을 소요한다. Caffe-based implementation과 trained networks를 제공한다. 우리는 u-net architecture이 다른 더 많은 tasks에서 쉽게 적용될 것을 장담한다.

< comment >

U-net은다양한 biomedical segmentation application에서 좋은 성능을 보이고 사용한 주요 techniques는 u-shaped architecture(with layers and channels), overlap-tile strategy & mirroring extrapolation, data augementation(by. elastic deformation)이다.