[Paper Review] TOTEM: TOkenized Time Series EMbeddings for General Time Series Analysis

[Background]

 

※ VQ-VAE

 

  이 논문에서는 VAE와 이산 표현을 결합한 새로운 생성모델(VQ-VAE)를 제안한다. Vector Quantisation을 사용하여 너무 큰 분산으로 생기는 어려움을 피하면서 학습하기 편하고 Posterior Collapse 문제를 회피할 수 있다고 한다. 그럼에도 연속표현을 사용하는 모델과 비등한 결과를 도출해낸다고 주장한다.

  또한 강력한 prior를 제공할 수 있는데, 예를 들어, 그림에서 변경하고 싶은 특정 부분만을 바꿀 수 있다거나, 말하는 내용을 그대로지만 목소리만 바꿀 수 있는 기능을 제공할 수 있다.

 

기본적으로 VAE에서 내부적으로 취급하는 분포는 Gaussian 분포를 따른다고 가정한다. 하지만 VQ-VAE는 여기에 이산 표현을 다루도록 한다. VQ를 사용하면서, posterior와 prior distribution은 categorical하며, 이 분포로 생성된 샘플은 Embedding table을 indexing 한다. 이 Embedding table이 Decoder의 입력으로 주어지게 된다.

 

 VQ-VAE란 데이터의 중요한 정보를 discrete codebook vector로 압축하여 학습하는 방법을 사용한다. VQ-VAE의 구조로는 Encoder, Quantization, Decoder로 구성된다.

 

 

▶ Encoder

  

  고차원 데이터를 저차원 잠재 공간으로 압축하는 역할을 한다. 입력 데이터를 Encoder에 입력한다. Encoder는 입력 데이터를 압축하고, 저차원의 잠재 벡터를 생성한다. 이 잠재 벡터는 연속적인 실수값을 가지며, quantization 과정을 통해 discrete 값으로 변환된다.

  Encoder는 입력 x를 받아  $ Z_{e}(x)$ 를 출력한다. 이산표현벡터 Z는 Embedding space에서 가장 가까운 Embedding vector를 찾는다.

 

 

위 식처럼 $ Z_{e}(x) $ 값과 Codebook의 어떤 한 벡터와의 차이 값이 최소가 되는 지점의 값을 1로 설정하고 나머지는 0으로 설정하는 원 핫 인코딩을 실행한다. 제안한 분포 q(z = k|x)는 deterministic하고, Z에 대해 단순 균등 prior를 정의함으로써 KL divergence를 상수로 얻을 수 있다.

 

▶ Quantization

 

  Encoder에서 생성한 연속적인 잠재 벡터를 discrete latent vector로 변환하는 단계이다. 먼저 Encoder에서 생성된 잠재 벡터 $  Z_{e}(x) $가 입력으로 들어온다. 이 잠재 벡터는 Codebook 내의 벡터들과 비교되어 가장 가까운 코드북 벡터로 매핑된다.

 

☆ Codebook

 

  벡터 양자화 단계에서 사용되는 이산적인 벡터의 집합으로, Encoder에서 생성한 latent vector를 discrete vector로 변환하는 데 사용된다.

  Codebook은 학습을 통해 latent space를 대표하는 K개의 이산 벡터를 포함한다. 각 벡터는 latent vector의 특정 패턴을 반영한다. Codebook은 학습 중에 지속적으로 업데이트된다. 각 latent vector가 Codebook vector로 매핑될 때, Codebook vector가 latent vector에 맞게 조정되며, Codebook이 데이터의 특징을 더 잘 표현할 수 있도록 최적화된다.

  Codebook 업데이트는 Qunatization Loss를 최소화하는 방향으로 이루어진다. 이 손실은 latent vector와 Codebook 간의 차이를 최소화하는 역할을 한다.

  위 그림에서는 Embedding space가 Codebook으로 나타나 있다. 위의 그림처럼 원래 데이터가 그림 형태라면 Codebook의 각 잠재 벡터도 2차원 벡터로 표현되어야 한다.

 

▶ Decoder

 

  latent vector화 된 데이터를 다시 원래의 고차원 데이터로 복원하는 과정이다. 이 단계는 Encoder의 반대 역할을 수행한다. 먼저 양자화된 벡터를 입력으로 받아 Decoder에 전달한다. Decoder는 Encoder와 반대 방향으로 작동하여, 여러 층의 역방향 컨볼루션 계층을 통해 데이터를 고차원으로 복원한다. 최종 출력은 원본 데이터와 동일한 차원을 가진 복원된 데이터이다. 목표는 원본 데이터와 최대한 유사하게 복원하는 것이다.

 

표현 $ Z_{e}(x) $ 는 다음과 같이 임베딩 e 중 가장 가까운 원소를 찾고 discretisation bottlenect으로 전달된다.

 

 

여기서 중요한 점은 이 식에서는 real gradient가 없다는 점이다. 이때 decoder input $Z_{q}(x) $를 encoder output으로 gradient를 복사하면 된다. forward 학습에서 $ Z_{q}(x) $가 decoder에게 전달되고 Back propagation에서는 gradient $\bigtriangledown _{z}L$이 Encoder로 그대로 전달된다.

 

  Encoder의 출력과 Decoder의 입력은 D 차원의 같은 공간에 존재하여, gradient가 어떻게 변화해야 하는지 정보를 줄 ㅜ 있다.

 

  전체 objective는 다음 식으로 표현된다.

 

 

1. 첫번째 Loss 항은 Reconstruction loss로 decoder와 encoder 모두를 최적화한다.

2. 임베딩 $e_{i}$는 어떤 gradient도 받지 못하기 때문에 Vector Quantisation을 사용한다.

(VQ objective는 각 e를 encoder의 출력 $ Z_{e}(x) $로 이동하도록 한다.)

3. embedding space는 무한하기 때문에 $ e_{i}$는 endover parameter만큼 빠르게 학습하지 않을 수 있다. Encoder가 embedding과 출력이 비슷할 수 있게 만들기 위해 commitmet loss를 출력한다.

 

다음은 Loss들에 대한 설명이다.

 

● 재구성 손실(Reconstruction loss)

 

  원본 데이터와 복원된 데이터 간의 차이를 최소화하는 손실 함수이다. 이는 입력 데이터를 얼마나 잘 복원했는지 측정한다.

$L_{reconstruction} = \left\| x-\hat{x}\right\|^{2}$

 

● 양자화 손실(Quantization Loss)

 

  Latent vector와 매핑된 Codebook vector 간의 차이를 최소화하는 손실이다.

 

$ L_{quantization} = \left\| Z_{e}(x) - c_{k} \right\|^{2} $  

 

● 코드북 커밋 손실(Commitment Loss)

 

  Codebook vector가 충분히 학습되도록 돕는 손실이다. 이는 Codebook vector가 latent vector에 더 잘 맞도록 조정된다.

 

$L_{commitment} = \left\| sg(Z_{e}(x)) - c_{k} \right\|^{2} $

 

  여기서 $sg(Z_{e}(x)) $는 stop-gradient 연산자로, 역전파 중 이 값에 대한 gradient가 고정된 상태로 역전파되도록 한다.

 

결국 VQ-VAE의 핵심은 데이터의 핵심 특징을 추출하여 연속적인 값으로 나타내는 것이 아니라 이산적인 값으로 잠재 벡터를 저장하는 기술이다. 이렇게 연속적인 값이 아닌 이산적인 값으로 나타내도 기존의 VAE와 성능적인 차이가 크지 않다고 한다.

 

더 자세한 설명은 VQ-VAE 본 논문을 참고하길 바란다.

 

[Summary]

 

 

 

다음은 TOTEM의 전체 구조와 학습 및 추론 방식에 대한 개요를 설명하는 그림이다. TOTEM은 시계열 데이터를 처리하는 아키텍처로 VQ-VAE(Background 참고) 방식을 기반으로 하고 있다.

 

ⓐ : Generalist Training & Specialist Training

 

Generalist Training : 모든 데이터셋을 동시에 학습하여 일반화된 모델을 만드는 방식이다. 다양한 도메인의 시계열 데이터를 학습하여 여러 작업을 한 번에 처리할 수 있도록 한다.

 

Specialist Training : 특정 데이터셋 하나에만 집중해 학습하는 방식이다. 각 작업에 맞게 개별적인 학습을 진행한다.

 

주요 구조 :

 

1. 1D CNN Encoder $\varepsilon $

 

  입력 시계열 데이터를 1D CNN을 통해 잠재 벡터 Z로 변환한다. 이 벡터는 시간에 따른 데이터를 압축하여 중요한 정보를 추출한 것이다. 컨볼루션 계층을 통해 데이터를 처리하면서, 필터가 시계열 데이터에서 유의미한 패턴을 학습한다. 이 과정에서 입력 데이터의 크기가 줄어들며, latent space에서 Z로 표현된다.

 

2. Quantize

 

  Encoder에서 생성된 잠재 벡터 Z를 quantization하여 이산적인 벡터로 변환한다. 이때 이산 벡터로 변환하는 과정에서 Codebook을 사용하여 Codebook에서 가장 가까운 벡터로 매핑되며 코드북 벡터는 고정된 수 K개의 이산 벡터로 이루어져있다.

  이때 인코더에서 나온 잠재 벡터 Z는 코드북 내의 벡터들과 비교되어 각 잠재 벡터는 코드북 벡터 중 가장 가까운 벡터로 대체된다. 이 과정을 통해 연속적인 벡터가 이산 벡터로 변환된다.

 

3. 1D CNN Decoder $\boldsymbol{D}$

 

  다음으로 양자화된 벡터를 디코더로 복원하여 원래의 시계열 데이터와 유사한 데이터를 재구성 한다. 이때 1D 역방향 CNN을 사용하여 데이터를 다시 원래 크기의 시계열로 복원한다. 이때 디코더에서 생성된 데이터는 원본 시계열 데이터와 비교되어 재구성 손실을 측정한다. 이 과정에서 재구성 손실을 최소화하도록 모델이 학습된다.

 

이제 Total loss를 구하는 방법에 대해 알아보자.

 

Total loss는 다음과 같다.

 

$\mathit{L}_{total} = \mathit{L}_{reconstruction} + \alpha\mathit{L}_{commitment}$

 

$ \mathit{L}_{reconstruction} = \frac{1}{E\cdot S}\sum _{i}\left\|x_{i}-\hat{x}_i \right\|^{2}$

 

$ \mathit{L}_{commitment} = \alpha\left\|Z_{e}(x)-sg[c_{k}] \right\|^{2}$

 

  TOTEM은 VQ-VAE모델을 기반으로 하지만 총손실값 계산에서는 약간의 다른 부분이 존재한다. VQ-VAE의 총손실은 위의 식에 Quantization Loss 값을 추가하여 총 손실을 계산하지만 TOTEM은 해당 loss값 계산을 하지 않는다.(이유는 게재되어 있지 않음) 여기서 $\alpha$ 값은 이 손실의 중요도를 조정하는 스칼라 파라미터값이다. 계산한 총 손실 값을 바탕으로 Codebook이 해당 데이터셋을 잘 나타낼 수 있도록 학습한다.

 

4. Freeze Codebook $\boldsymbol{C}$

 

  학습 중 일정 시점에서 코드북을 고정시켜 특정 작업에서 더 나은 성능을 발휘하도록 조정한다. 고정된 Codebook을 바탕으로 다음 작업을 수행한다.

 

 

ⓑ : In Domain Inference & Zero-Shot Inference

 

In Domain Inference : TOTEM이 학습한 데이터 도메인 내에서의 추론을 의미한다. 모델이 학습한 데이터셋에 대해 결측값 대체(imputation), 이상 탐지(anomaly detection), 예측(forecasting)의 작업을 수행한다.

 

Zero-Shot Inference : TOTEM은 학습되지 않은 새로운 데이터셋에 대해서도 추론할 수 있다. 이 경우에도 학습된 Codebook을 활용하여 일반화된 작업을 수행한다.

 

주요 구조 :

 

1. Tokenize

 

  입력 데이터를 코드북 벡터로 변환하여 이산적인 토큰 시퀀스로 변환한다. 각 토큰은 시계열 데이터의 중요한 패턴을 나타낸다.

 

2. Inference Tasks

 

  다음은 TOTEM이 학습된 코드북으로 추론할 수 있는 작업들이다.

 

Imputation : 데이터의 결측된 부분을 채우는 작업                                       

Anomaly Detection : 시계열 데이터에서 비정상적인 패턴을 탐지하는 작업

Forecasting : 시계열 데이터를 바탕으로 미래의 값을 예측하는 작업          

 

  Imputation이나 Anomaly Detection은 다른 구조 없이 TOTEM의 VQ-VAE 구조를 사용하면 즉시 해결할 수 있다. 하지만 Forecasting은 다른 모델 없이 구현해 낼 수 없다. 이 논문에서는 transformer 구조를 사용한다.

 

 

1. Tokenize with Pretrained Codebook

 

  먼저 시계열 데이터의 각 시간 스텝은 사전 학습된 Codebook을 사용해 이산화된다. 이 과정에서 시계열 데이터는 토큰화된 형태로 변환된다. 토큰화과정으로 이산화된 토큰 시퀀스는 이후 Transformer의 입력으로 제공된다.

 

2. Add Positional Embedding

 

  시계열 데이터의 시간 정보를 Transformer에게 제공하기 위해 사용된다. Transformer는 기본적으로 입력 간의 순서를 고려하지 않기 때문에 시계열 데이터처럼 시간 순서가 중요한 데이터에서는 Positional Encoding을 통해 시간 정보를 추가해준다.

  이 과정을 통해 Transformer는 각 토큰 시퀀스가 시계열 데이터의 어느 시간 스텝에 해당하는지를 학습할 수 있다.

 

3. Transformer Encoder

 

  이 토큰 시퀀스를 처리하여 시계열 데이터의 장기적인 패턴을 학습한다. Self-Attention 메커니즘을 사용하여 토큰 간의 관계를 학습하고, 이를 바탕으로 미래 시점의 데이터를 예측하는 데 필요한 정보를 추출한다. Transformer는 이산화된 토큰 간의 상호작용을 분석하여 시계열 데이터의 패턴을 예측하게 된다.

 

4. Linear Layers

 

  Transformer 인코더의 출력은 선형 계층을 통해 후처리된다. 여기서 미래 시계열 값의 평균과 표준편차를 예측하게 된다. 이 선형 계층은 Transformer에서 학습된 패턴을 기반으로 미래 시점의 시계열 데이터를 예측하는데, 이때 시계열 데이터의 예측값에 대한 분포(평균과 표준편차)를 반환한다. Transformer 계층에서 평균과 분산에 대한 학습은 이루어지지 않기에 선형 모델을 통해 평균과 분산을 예측한다.

 

최종 예측값은 Transformer 모델을 통해 예측한 시계열 데이터 값에 표준편차를 곱하고 평균을 더하여 정규화된 데이터르르 복원하는 방식으로 계산된다. 당연히 이 경우에도 예측값과 실제값 사이의 손실을 최소화하기 위해 역전파학습이 이루어진다. 

 

이제 전체 과정을 정리해보자.

 

먼저 Codebook 학습과정과 실제 데이터 예측 과정 두 분류로 나누어서 설명하겠다.

 

Codebook 학습과정

 

1. 여러 데이터셋 또는 하나의 데이터셋을 입력으로 준비한다.

2. 데이터셋을 Encoder를 통해 codebook의  latent vector에 가장 가까운 벡터로 인덱싱된다.

3. 모든 데이터셋이 인덱싱된 후에 해당 table을 디코더를 통해 원래 데이터셋으로 복원시킨다.

4. 두 데이터셋의 손실값을 계산하여 손실값을 줄이는 방향으로 학습된다.

5. Codebook의 학습이 완료되었다고 생각되면 Freeze한다.

 

학습된 Codebook으로 예측하기

 

1. 먼저 예측하고자 하는 데이터셋을 입력한다.

2. 위의 과정에서 학습한 Codebook으로 입력 데이터셋을 tokenize한다.

3. Transformer는 시간 순서에 상관없이 학습하므로 positional embedding을 추가한다.

4. 입력한 데이터에 대해 transformer 모델이 표준화된 값을 예측하여 출력한다.

5. 또한 입력 데이터셋의 평균과 분산을 선형 모델을 통해 예측한다.

6. 4와 5단계에서 얻은 값으로 원래의 데이터셋으로 복원한다.

 

[Conclusion]

  논문을 읽으면서 궁금했던 부분은 시계열 예측 부분이므로 Imputation과 Anomaly Detection 부분은 건너뛰도록 하겠다.

 

 

위 Table은 특정 데이터셋을 학습시킨 후 예측한 결과이다. TOTEM 모델은 28.6% 정도의 예측 성공률을 달성했다.

 

 

위 테이블은 Generalist Forecasting 방법으로 GPT2와 비교했을때  A. In-Domain performance에서는  67.9%, B. zero-shot performance에서는 90.0%의 예측성공률을 보여주었다.

 

이러한 TOTEM의 구조는 특성상 많은 시계열 데이터셋을 학습시키도록 되어있다. 그러한 점에서 여러 윤리적인 문제가 발생할 수 있다. 

프라이버시 우려 : 시계열 데이터는 건강 분야와 같이 개인에 관한 민감한 정보를 포함할 수 있다.

 

오용 가능성 : 시계열 예측 모델은 오용될 수 있다. 예를 들어 주식 가격이나 시장 변동을 예측하는 경우, 이는 내부자 거래나 기타 불법 금융 활동에 악용될 수 있다.

 

경제적 영향 : 시계열 모델을 기반으로 한 자동화된 예측과 결정은 산업과 노동 시장에 긍정적 또는 부정적 영향을 미칠 수 있다.

 

  논문을 읽으면서 TOTEM이 과연 좋은 모델인가에 대한 고민을 계속하게 되었다. TOTEM을 활용하기 위해 수많은 시계열 데이터로 Codebook을 학습시켜야 한다. 심지어 시계열 데이터셋은 다른 데이터셋보다 정형화되어 있지 않고 데이터의 가짓수도 상대적으로 적은 편이다. 이런 시계열 데이터의 한계점에서 이 모델은 시계열 데이터의 단점을 부각시키는 모델이라는 생각이 들었다. 만약 학습시킨 데이터셋이 충분치 않아, Codebook에 학습된 잠재벡터와 입력값의 이산화된 벡터가 손실 값이 큰 경우 예측 성공률이 낮아질 수 있다. Codebook을 사용한다는 점이 아무래도 제약사항이 많지만 VQ-VAE 방식을 사용하여 시계열 데이터를 예측한다는 점은 시계열 예측 분야에 좋은 시도라고 생각한다.

 

[Reference]

 

Xiyuan, Ranak, et al. "Large Language Models for Time Series: A Survey" arXiv:2402.01801v3(2024)

 

Sabera, Yisong, Georgia "TOTEM: TOkenized Time Series EMbeddings for General Time Series Analysis" arXiv:2402.16412v1(2024)