SOM-VAE: Interpretable Discrete Representation Learning On Time Series

2021. 8. 25. 18:25Neural Networks/Time Series

논문에 대한 오역, 의역등이 다수 포함되어 있습니다. 댓글로 많은 의견 부탁드립니다

 


Author : Vincent Fortuin, Matthias Hüser, Francesco Locatello, Heiko Strathmann, Gunnar Rätsch

International Conference on Learning Representations (ICLR 2019) [paper link]

 

논문은 Tensorflow 1으로 구현되었고 [link], pytorch version으로도 implementation되어있지만 [torch link], TF 1은 다루기가 힘들고 torch version은 잘못 구현된 부분이 존재함을 확인했습니다. 따라서 저자의 구현을 바탕으로 TF 2로 migration했고, 해당 코드는 jupyter notebook으로 작성되었고 github를 통해 제공합니다 [link].

 

개인적으로 아쉬운점은 시계열 데이터를 해석하기위해 제안된 반면, 실제로는 MNIST와 fashion-MNIST에서만 실험되었다는 부분입니다. 저자는 이후 발표한 DP-SOM 등에서 SOM-VAE의 시계열 데이터의 실험을 제공하지만 역시 TF1인 관계로 기회가 닿는다면 나중에 migration해보려합니다.


1. Introduction

 

SOM-VAE는 Variational Auto-Encoder (VAE)와 Self Organized Map (SOM)을 엮은 논문입니다. 논문의 큰 골자는 VAE가 latent vector로부터 샘플링하던 방법을 SOM 알고리즘으로 대체합니다. 때문에 latent vector가 입력 데이터에 따른 연관성을 갖게되며 cluster로 표현할 수 있다는 것이 핵심입니다. Generative Adversarial Networks (GAN)은 VAE와 더불어 데이터를 latent vector로 재표현해볼 수 있는 방법 중 하나이지만, GAN의 latent vector는 continous한 특징이 있기 때문에 SOM알고리즘과 접목시키기 어려웠고 VAE를 사용하게 되었다고 합니다. 저자가 요약한 핵심 기여는 아래와 같습니다. 

  • 시계열을 해석가능하고 구분된 수치형 공간으로 재표현할 수 있는 고유 framework을 고안했습니다.
  • Representation learning 방법론에서 latent vector가 확률적인 모델은 clustering performance를 향상시키고, 해석가능한 장점을 갖습니다. 
  • 제안한 방법의 clustering performance는 벤치마크와 실제 의료 시계열 데이터셋에서 다른 방법들보다 우월했고, 제안하는 방법은 확장가능합니다. 

 

2. SOM-VAE

 

본 논문은 5개의 loss를 합쳐 하나의 모델을 학습시키게 됩니다. 한번에 수학적 notation을 모두 작성하면 설명드리기 어려울 것 같아 순서대로 작성하려합니다.

SOM-VAE architecture

저자가 제안하는 모델은 기본적으로 Auto-encoder (AE)와 같은 구조입니다. Encoder $f_\theta(x)$는 입력 데이터 $x \in \mathbb{R}^d$를 latent vector $ z_e \in \mathbb{R}^m$ (일반적으로 $m < d$)로 재표현하고, decoder $g^e_\phi(z_e)$는 latent vector $z_e$로부터 실제 입력데이터 $x$를 복원해내도록 학습합니다. 여기서 VAE는 latent vector $z$로부터 유사한 새로운 latent vector $z'$을 추출합니다. 하지만 SOM-VAE는 특이한 방법을 사용합니다. 무작위로 생성된 embedding dictionary $E$로부터 입력 데이터 $x$를 압축한 latent vector $z_e$와 가장 관련이 깊은 $e_i \in E$를 추출합니다. $z_e$와 $e_i \in E$의 관련성은 euclidean distance로 측정되며, 계산했을 때 최소화되는 $e_i$를 사용하여 $z_q$를 샘플링합니다. 위 과정을 수식으로 설명하면 아래와 같습니다. 

$$ E=\{e_1,\cdots,e_k | e_i \in \mathbb{R}^m \} $$

$$ p(z_q|z_e) = \mathbb{1}[z_q = {\arg\min}_{e \in E} \lVert z_e - e \rVert^2 ]$$

$\mathbb{1}[\cdot]$은 index를 가리키는 함수입니다. 따라서 입력 데이터 $x$를 압축한 $z_e$와 가장 유사한 index를 활용하여 embedding dictionary $E$로부터 특정 $e_i$를 선정합니다. $E$는 기본적으로 dictionary 형태로 구분되어있으며, 각 value의 차원은 $m$이기 때문에 특정 index의 $e_i$를 $z_q$로 간주할 수 있습니다. $E$로부터 샘플링된 $z_q$는 $z_e$와 가장 가까운 거리에 위치한 embedding 값에 따른 무작위 latent vector입니다. $z_q$는 $z_e$와 차원 크기가 같으므로 별도의 decoder $g^q_\phi(z_q) \in \mathbb{R}^d$를 활용하여 입력 데이터 $x$와 같은 차원으로 복구가능합니다. 한편, 입력으로부터 encoder를 지나 decoder들로부터 얻을 수 있는 복원된 $x$는 실제 입력 데이터 $x$와 같아야하므로, reconstruction loss를 사용하여 성공적인 복원이 가능한 모델을 학습합니다. 또한 $z_q$는 $z_e$로부터 sampling된 latent vector이므로 상호간의 유사성이 존재해야합니다. 이 부분을 학습하기 위해 commitment loss를 사용합니다. 

$$ \hat{x}_e = g^e_\phi(z_e)\ and \ \hat{x}_q = g^q_\phi(z_q)$$

$$\mathcal{L}_{\mathrm{reconstruction}} = \lVert x-\hat{x}_e \rVert^2 + \lVert x-\hat{x}_q \rVert^2$$

$$\mathcal{L}_{\mathrm{commitment}} = \lVert z_e(x) - z_q(x)\rVert^2 $$

마지막으로 SOM-VAE는 latent space에서 SOM algorithm을 구현하기 위해 불연속적(discrete) 공간상의 정보를 loss로 학습합니다. 앞서 $z_e$는 입력 데이터 $x$를 압축시킨 값이고 $z_e$와 가장 가까운 거리의 embeding 값을 $z_q$로 정의했습니다. 그런데 샘플링된 $z_q$는 embedding dictionary $E$ 중 하나이기 때문에 사실상 $e_i = z_q$와 같습니다. Embedding dictionary $E$의 주소가 2차원상으로 표현가능하다면, 그리고 2차원상 주소들 중 하나에 $e_i$가 위치한다면, $e_i$의 주변에 위치한 값들도 알 수 있을 것입니다. 입력 데이터를 압축했을 때와 가장 가까운 embedding 값은 commitment loss가 학습시켜줍니다. 하지만 embedding space를 기준으로 본다면 특정 embedding $z_q$만 학습되고 $z_q$의 주변 값들은 고정되어 있기 때문에 $z_q$와 그 주변의 embedding 값들 $N(z_q)$는 연관성을 갖지 못하게 됩니다. 선정된 $z_q$가 주변에 대해서도 영향을 끼친다면, latent vector는 무작위 공간이 아니라, 입력에 따른 제한된 공간을 보장할 수 있을것입니다. SOM loss는 이것을 반영하기 위해 사용됩니다. 

$$ \mathcal{L}_{\mathrm{SOM}} = \sum_{ \tilde{e} \in N(z_q)} \lVert \tilde{e} - SG[z_e] \rVert^2 $$

$SG[\cdot]$ 연산자는 위 로스로 발생하는 역전파를 encoder에는 학습시키지 않는다는 stopping gradient 연산자입니다. Commitment loss가 반영한 $z_e$와 $e_i$의 관계를 $e_i$의 주변부와 엮는 과정에서 제외시켜, encoder를 두번 학습시키지 않으면서 주변부에 대해서는 $e_i$와 유사한 의미를 띄도록 학습하게 됩니다. 저자는 앞서 살펴본 세가지 loss들을 묶어 $\mathcal{L}_{\mathrm{SOM-VAE}}$라고 칭합니다. 

$$\mathcal{L}_{\mathrm{SOM-VAE}} = \mathcal{L}_{\mathrm{reconstruction}}(x,\hat{x}_q,\hat{x}_e) + \alpha \mathcal{L}_{\mathrm{commitment}}(x)+\beta \mathcal{L}_{SOM}(x)$$

 

3. SOM-VAE with time series

 

이번 리뷴의 서두에서 시계열에 관해 언급했지만 지금까지 살펴본 SOM-VAE는 시계열의 특징을 살린다고 보기 어려운 방법입니다. 이후 markov model을 활용하여 시계열의 의미를 담을수 있는 추가 loss들을 $\mathcal{L}_{\mathrm{SOM-VAE}}$에 더해주고 최종 loss로 사용합니다. 시계열이 입력되었을 때 손실함수는 아래와 같이 씁니다. 

$$\mathcal{L}( x^{t-1}, x^t, \hat{x}^t_q, \hat{x}^t_e) = \mathcal{L}_{\mathrm{SOM-VAE}}(x^t,\hat{x}^t_q,\hat{x}^t_e)+\gamma \mathcal{L}_{\mathrm{transitions}}(x^{t-1}) + \tau \mathcal{L}_{\mathrm{smoothness}}(x^{t-1},x^t)$$

먼저 $\mathcal{L}_{\mathrm{transitions}}$는 Markov model 기반의 시계열의 의미를 반영하기 위한 loss입니다. 시간이 진행됨에 따라 $z_q$ state가 변화할 확률을 계산하며 logloss로 학습합니다.

$$\mathcal{L}_{\mathrm{transitions}}= - \log P_M( z_q( x^{t-1} ) \rightarrow z_q(x^t) )$$

한편 $\mathcal{L}_{\mathrm{smoothness}}$는 SOM loss와 유사하게 이전 state의 z_q가 다음 state의 z_q뿐만 아니라 z_q의 이웃에도 반영되도록 유도하는 loss입니다. SOM loss와의 차이점이라면 SG 연산자가 encoder부분에서 제외됩니다. 

$$ \mathcal{ L }_{ \mathrm{ smoothness } } = \mathbb{ E }_{ P_M( z_q( x^{ t-1 } ) \rightarrow \tilde{ e } ) } [ \lVert \tilde{ e } - z_e(x^t) \rVert^2 ]$$

 

4. Experiments

 

실험 결과는 상단 논문 링크를 참고해주세요. 이번 논문을 MNIST에 대해 직접 실험한 결과는 아래와 같고, github의 jupyter notebook으로 재현가능합니다. 

SOM-VAE cluster map with MNIST
MNIST reconstruction test with SOM-VAE