CSDI 논문 리뷰(Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation)

Time-series diffusion model 중에서도 인용수가 높다는 CSDI 논문을 리뷰해보고자 한다. (현 기준 인용수 350회) 이 논문은 저자들이 매우 유명한 사람들인데, Yusuke Tashiro, Jiaming Song, Yang Song, Stefano Ermon(나의 cs236 강의 교수님)이 21년도에 뉴립스에 낸 논문이다. github 주소를 보니, 이들을 Ermongroup이라고 부르는 것 같다. 멋진 사람들..

 

Abstract

CSDI(Condition Score-based Diffusion Models for Imputation)은 observed data를 조건으로 score-based diffusion model을 이용한 것이다. 기존의 score-based approach와 다르게 conditional diffusion model은 imputation을 직접 학습할 수 있으며, 관측된 데이터들 간의 correlation을 사용할 수 있다. probabilistic imputation method와 deterministic imputation에서 CSDI는 모두 좋은 성적을 거두었으며, time series interpolation과 probaiblistic forecasting에서도 좋은 결과를 가져온다.

 

Introduction 

CSDI는 conditional distribution을 conditional score-based diffusion model을 통해 직접 학습하여 probabilistic imputation을 용이하게 한다. imputation과정은 random noise에서 시작해서 reverse process $p_{\theta}$를 통해 plausible time series로 변환하는 과정으로 기존에 관측된 정보의 활용을 통해 denoising한다. 여기서 attention mechanism을 사용해 temporal dependency와 feature dependency를 모두 포착하고자 설계하였다.

conditional diffusion model을 학습하기 위해서는 observed value와 ground-truth missing values가 필요하다. 하지만 ground-truth를 모르거나 training data에 missing value가 포함되지 않은 것이 일반적이므로, 이 논문에서는 self-supervised training method를 사용해 observed value를 conditional information과 imputation targets로 나누어 학습한다.

 

Related works

Time series imputations with deep learning : For deterministic imputation - RNN(LSTMs, GRUs), RNN + Gan / self-training, RNN + attention mechanisms / For probabilistic imputation - GP-VAE

Score-based generative model : TimeGrad(RNN기반으로 imputation 불가)

 

Background

 

※ DDPM에 대한 설명은 더보기를 참고하세요!

Conditional score-based diffusion model for imputation (CSDI)

* notation이 너무 복잡해서 논문을 최대한 활용해 정리하고자 한다. 

CSDI에서는 우선 ${\epsilon}_{\theta}$를 위와 같이 확장해 ${\mu}_{\theta}$와 ${\sigma}_{\theta}$를 확장했다. 달라진 것은 ${\epsilon}_{\theta}$밖에 없으므로 DDPM의 학습 과정을 그대로 차용할 것이다. noisy target $x_{ta_t}  $은$x_{ta_t} = \sqrt{\alpha_t} x_{ta_0} + (1 - \alpha_t)\epsilon$으로 샘플링할 수 있으며,  ${\epsilon}_{\theta}$는 아래 loss function으로 학습할 수 있다.(imputation target $x^{ta}_0$와 동일한 차원) 

논문의 (7)번 식

하지만 위와 같은 학습 방식에는 문제가 있는데, 실제로 누락된 값의 정답을 알 수 없기 때문에 학습 샘플 $x_0$를 어떻게  $x^{ta}_0$와  $x^{co}_0$로 선택할 지가 명확하지 않다. 이 문제를 해결하기 위해 self-supervised learning 방법을 사용한다. 

 

self-supervised learning 

샘플 $x_0$가 주어지면 이를 두 부분으로 나누고, 뒤에서 논의할 target choice strategy를 이용해  일부분을 $x^{ta}_0$로, 나머지는 $x^{co}_0$으로 설정한다. 그런 다음 노이즈가 포함된 imputation target $x^{ta}_t$를 샘플링하면서 위 loss function을 통해 ${\epsilon}_{\theta}$을 학습한다. 

target choice strategy

imputation target을 어떻게 고를 것인지에 대한 네 가지 선택 전략은 다음과 같다.

1. Random strategy : 누적된 패턴에 대해 알지 못할 때 사용되며, 관찰된 값의 일정 비율을 무작위로 선택하여 imputation target으로 설정한다.

2. Historical strategy : 학습 데이터셋의 누락된 패턴을 활용하는 전력으로 주어진 샘플 $x_0$와 다른 샘플 $\tilde{x}_0$에 대해 $x_0$의 관찰된 인덱스와 $\tilde{x}_0$의 누락된 인덱스의 교집합을 imputation target으로 설정한다. 

이 전략은 학습 데이터셋과 테스트 데이터셋에서 누락 패턴이 높은 상관관계를 가질 때 좋은 성능을 가져다준다.

3. Mix strategy : Random strategy + Historical strategy로 Historical strategy만을 사용했을 때 생길 수 있는 overfitting을 방지해준다. 

4. Test pattern strategy : 테스트 데이터셋에서 누락된 패턴을 알고 있을 때, 이 패턴을 그대로 보완 대상으로 설정한다. 예를 들어 time-series forecasting을 수행할 경우 누락된 패턴은 주어진 미래시점으로 고정되어 있으므로 이 전략을 사용할 수 있다. 

 

지금까지의 설명은 일반적인 imputation에 모두 해당하는 설명이였다. 이제부터는 time series imputation 중심으로 설명하고자 한다. 

 

Implementation of CSDI for time series imputation

time-series imputation을 수행하게 되면 ${\epsilon}_{\theta}$가 input으로 받는 $x^{ta}_0$와  $x^{co}_0$가 달라진다는 문제가 생긴다. 이 문제를 해결하기 위해  ${\epsilon}_{\theta}$의 입력을 고정된 샘플 공간 $\mathbb{R}^{K \times L}$로 한정한다. 즉, zero padding을 추가해 $x^{ta}_0$와  $x^{co}_0$의 shape을 $\mathbb{K \times L}$로 고정한다는 말이다. 위 그림에서 하얀 색 부분을 zero로 패딩한다고 이해하면 된다.

zero padding된 부분을 표현하기 위해 논문에서는 conditional mask $m_{co} \in \{0, 1\}^{K \times L}$를 ${\epsilon}_{\theta}$의 추가적인 input으로 설정한다. $m_{co}$는 $x^{co}_0$와 대응되며,  $m_{co}$의 원소는 $x^{co}_0$의 원소가 관측되면 1을, 아니면 0을 부여받는다. 

또한 출력을 쉽게 다루기 위해 output shape 역시 제로 패딩을 사용해  $\mathbb{R}^{K \times L}$에 속하도록 한다. 그러면 conditional denoising function $\epsilon_{\theta}(x_{ta_t}, t \mid x_{co_0}, m_{co})$은 $\epsilon_{\theta} : (\mathbb{R}^{K \times L} \times \mathbb{R} \mid \mathbb{R}^{K \times L} \times \{0, 1\}^{K \times L}) \rightarrow \mathbb{R}^{K \times L}$와 같이 나타낼 수 있다. 

 

샘플링 시 $x^{co}_0$는 모두 관측 가능하므로 $m_{co} = M$이고 $x_{co_0} = m_{co} \odot X$이다. 

학습 시  $x^{ta}_0$와  $x^{co}_0$는 target choice strategy에 의해 결정되며, $x^{co}_0$를 사용해 $m^{co}$를 설정한다. 그러면 $x^{co}_0 = m^{co} \odot X$, $x^{ta}_0 = (M-m^{co}) \odot X$으로 쓸 수 있다. 

 

마지막으로 ${\epsilon}_{\theta}$는 DiffWave의 구조를 기본으로 채택했으며, multiple residual layers with residual channel C로 구성된다. 다만 DiffWave의 구조에서 개선한 부분도 있는데 Appendix E에서 detail을 확인할 수 있다. 

 

Attention mechanism

multivariate time series의 temporal / feature dependency를 포착하기 위해 각 residual layer에서 convolution architecture 대신 2차원 attention mechanism을 사용한다. 아래 그림과 같이 1-layer transformer encoder인 temporal transformer layer과 feature transformer layer를 도입한다. temporal transformer layer 층은 각 feature에 대한 텐서를 입력으로 받아 temporal dependency를 학습하며, feature transformer layer는 각 시점의 텐서를 입력으로 받아 temporal dependency를 학습한다. 

시계열의 길이 $L$는 시계열에 따라 다를 수 있으나, attention mechanism을 통해 모델은 다양한 길이를 처리할 수 있다. batch training의 경우, 시퀀스의 길이가 동일하도록 각 시퀀스에 zero padding을 적용한다. 

 

side information

시간 의존성 학습을 위해 사용하는 time embedding $s = \{s_{1:L}\}$은 128차원의 시간 임베딩을 사용한다. 또한 K개의 특성에 대해 categorical feature embedding을 활용하며, 이는 16차원을 가진다.

 

 

Experimental results

time series imputation - probabilistic imputation

✔️ 2가지 데이터 사용 : 의료 데이터셋(80%가 누락된 값으로 ground-truth X, random strategy 채택)과 대기질 데이터셋(13%가 누락된 값으로 historical strategy 채택)

✔️  비교한 모델 : Multitask GP, GP-VAE, V-RIN

✔️  비교 지표 : CRPS

✔️  측정 방법 : 누락된 값에 대한 확률분포를 근사하기 위해 100개의 샘플을 생성하고, 모든 누락된 값에 대한 CRPS의 normalized averge를 측정)

* CRPS(Continuous ranked probability score) : 추정된 확률분포와 실제 관측치 간의 호환성을 측정함

✔️  결과 :

time series imputation - deterministic imputation

✔️ 측정 방법 : 100개의 생성된 샘플의 중앙값 활용

✔️ 비교 지표 : MAE

✔️  결과 : 

누락 비율이 작을 때 특히 CSDI가 잘 작동한다.

Interpolation of irregularly sampled time series

✔️ 앞서 사용한 의료 데이터셋 + 불규칙하게 샘플링된 시계열로 처리

✔️ ground-truth가 없기 때문에 임의의 10/50/90%의 시간 지점을 선택하고, 해당 시간 지점에서의 관찰된 값을 테스트 데이터의 정답으로 이용(random strategy 채택)

✔️결과 :

Time series Forecasting

✔️ 5개의 데이터셋 사용 / test pattern strategy 채택

✔️ 비교 지표 : CRPS-sum(모든 K개의 특성에 걸친 시계열 합계의 분포에 대한 CRPS)

✔️ 결과 :

imputation보다 성능이 좋진 않았다.

 

Conclusion

CSDI의 계산 효율성은 ODE solver 등과 결합해 더 빨라질 수 있을 것이며, 분류와 같은 하위 작업들에 대해서도 적용할 수 있을 것이다. 또한 시계열 외의 다른 modality로 확장해는 것도 흥미로운 작업이 될 것이다.