[ICLR 2022] 1편: Efficient Transformer

‘ICLR(The International Conference on Learning Representations, 표현 학습 국제 학회)’은 2013년에 최초로 개최되어 올해로 10주년을 맞은 딥러닝 분야 세계 최고 권위 학회입니다. 인공지능, 데이터 사이언스, 비전, 음성인식 그리고 로보틱스에 이르기까지 딥러닝을 활용하는 모든 분야에 걸친 최첨단의 연구와 논문을 나누는 공유의 장이기도 합니다. 지난 4월 온라인으로 진행된 ‘ICLR 2022’에서 소개된 논문을 LG AI연구원이 소개합니다. - 1편: Efficient Transformer – Applied AI Research Lab 김정희 님 - 2편: Time Series Anomaly Detection – Applied AI Research Lab 심예슬 님 - 3편: Reinforcement Learning as a sequence modeling problem – Applied AI Research Lab 윤든솔 님

Introduction

Transformer는 구글에서 2017년에 발표한 sequence-to-sequence 모델로서 최근 Natural Language Processing (NLP), computer vision, audio 등 다양한 분야에서 활용되고 있습니다. 이번 블로그에서는 ICLR 2022에서 발표된 Transformer를 효율적으로 구현하는 세 가지 방법에 대해 소개합니다. 이를 위해 먼저 Seq-to-Seq model과 Transformer에 대해 간단히 살펴보고, 이후 Transformer에서 가장 연산을 많이 차지하는 부분인 self-attention을 효율적으로 구현하는 세 가지 방법에 대해 소개하는 것으로 블로그를 마치겠습니다.

Seq-to-Seq model Overview

그림 1. Seq-to-Seq 모델

Seq-to-Seq 모델은 Encoder-Decoder 구조를 가지면서 입력 sequence로부터 다른 도메인의 sequence를 출력하는 모델을 말합니다. Seq-to-Seq 모델링을 위해서 RNN (Recurrent Neural Network) 계열의 모델이 많이 활용되었습니다. RNN의 장점은 바로 직전까지 입력받은 시계열의 정보를 hidden state로 저장할 수 있다는 점인데, 이를 통해 시계열 데이터를 효과적으로 모델링할 수 있습니다. 그러나 현재 hidden state는 바로 이전 hidden state를 재귀적으로 받아들여 업데이트하게 되는데, 만약 sequence 사이의 거리가 먼 경우 그래디언트가 점차 줄어들어 hidden state의 과거 정보가 잘 전달이 되지 않는 long-term dependency problem을 야기하게 됩니다.

그림 2. LSTM (Long Short Term Memory) 구조

이를 해결하기 위해서 제안된 모델은 바로 LSTM (Long Short Term Memory)입니다. LSTM은 hidden state에 cell state를 추가한 구조입니다. Cell state는 일종의 컨베이어 벨트 역할을 하기 때문에 sequence가 지나감에 따라 그래디언트가 비교적 잘 전파가 되도록 합니다. 그럼에도 불구하고 long-term dependency problem이 완벽하게 해결되지 않았고, 순차적으로 연산한다는 점에서 병렬처리에 어려움이 있어서 많은 연산량으로 인해 학습 속도가 느린 단점이 있습니다. 이러한 RNN 계열의 문제점들을 해결하기 위하여 등장한 것이 바로 Transformer입니다.

Transformer Overview

그림 3. Transformer 구조

Transformer^[1]는 2017년 구글이 발표한 논문인 "Attention is all you need"에서 제안된 모델입니다. 그림 1과 같이 기존 Seq-to-Seq 모델은 인코더-디코더 구조로 구성되어 있는데 인코더는 입력 sequence를 하나의 context 벡터 표현으로 압축하고, 디코더는 이 context 벡터 표현을 통해서 출력 sequence를 만듭니다. 하지만 이러한 구조는 인코더가 입력 sequence를 하나의 벡터로 압축하는 과정에서 입력 sequence의 정보가 손실될 수 있다는 단점이 있었습니다. 이를 해결하기 위해 Transformer에서는 Encoder-Decoder 구조는 유지하면서도 self-attention만을 적용하여 기존 RNN 계열보다 우수한 성능을 보여주었습니다.

그림 4. Bert의 구조

또한 Transformer는 Bert, RoBerta 등 최신 NLP 분야에서 널리 활용되는 언어 모델의 backbone 역할을 하기도 했습니다.

그림 5. Self-attention mechanism

Transformer가 Seq-to-Seq 모델링에서 우수한 성능을 갖는 이유는 바로 self-attention 때문입니다. 예를 들면 ＂The animal didn＇t cross the street because it was too tired.＂라는 문장이 있다고 했을 때, it이 무엇을 의미하는지 알아내는 것은 문장 해석의 중요한 포인트입니다. 이때 self-attention을 사용하면 문장 내에서 it에 대응하는 단어들을 모두 고려하여 가장 가능성이 높은 'animal'로 학습하게 됩니다. 따라서 self-attention을 통해 모델이 문장 내부 단어의 의미를 더 잘 이해할 수 있게 됩니다.

그림 6. Self-attention의 연산량

Self-attention에서는 입력 sequence의 표현을 Query (Q), Key (K), Value (V)로 나누어 학습하고 Query, Key, Value의 행렬 곱셈을 통해 attention weight를 계산하게 되는데, 이 attention weight를 통해 전체 입력 sequence 중 어느 부분에 주목할지를 결정하게 됩니다. 그러나 그림 6에서 보는 것과 같이 self-attention은 $O\left(N^2\right)$의 연산량을 갖기 때문에 NLP, 수요 예측 등 long sequence의 입력이 필요한 다양한 응용에 대해서는 적용하기 어렵다는 단점을 갖고 있습니다. 여기서 N은 입력 sequence의 길이입니다.

그림 7. Self-attention을 효율적으로 대체하는 방법^[4]

따라서 self-attention을 효율적으로 구현하기 위해 여러 가지 방법이 제안되었는데, 크게 2가지로 나눌 수 있습니다. 첫 번째 방법은 self-attention을 approximation 하는 방법입니다. Attention weight는 값의 크기가 대부분 0에 가까운 값을 갖기 때문에 모두 계산하지 않고 중요하고 필요한 attention weight만 계산하는 방법이 제안되었습니다. 대표적으로 Sparse attention ^[2], low lank approximation ^[3] 등이 제안되었습니다.

두 번째 방법으로는 self-attention을 효율적인 구조로 대체하는 방법입니다. Self-attention의 matrix multiplication을 addition으로 대체하거나 self-attention을 단일 query만 계산하는 등 attention free 방법^[4]이 제안되었습니다.

이번 블로그에서는 ICLR 2022에서 발표된 self-attention을 효율적으로 구현하는 세 가지 방법 (Pyraformer^[5], cosFormer^[6], PoNet^[7]) 에 대해 소개하겠습니다. Pyraformer와 cosFormer는 self-attention을 approximation 하는 방법이며, PoNet은 pooling mechanism을 통해 self-attention을 효율적인 구조로 대체하는 방법입니다.

Pyramial Attention based Transformer (Pyraformer)

그림 8. 다양한 Seq-to-Seq 모델의 signal traversing path

먼저 Pyraformer에 대해 소개해 드리겠습니다. Seq-to-Seq 모델에서는 signal traversing path가 짧을수록 sequence의 dependency가 잘 파악이 됩니다. 예를 들면, 기존 CNN, RNN에서는 그림 8 (b), (c)와 같이 sequence의 길이가 N이라고 할 때 입력 신호의 dependency를 모두 파악하기 위해서는 O(N)의 signal traversing path가 필요합니다. 한편 Transformer(full attention)는 layer마다 모든 입력 sequence 사이에서 self-attention을 수행하기 때문에 그림과 같이 maximum signal traversing path가 O(1)이 됩니다. 만약 Transformer에 compact multi-resolution 구조를 활용하면 maximum signal traversing path는 O(1)을 유지하면서도 Transformer의 연산량을 줄일 수 있지 않을까 하는 생각에서 Pyraformer가 제안되었습니다.

그림 9. Pyraformer의 구조

Pyraformer의 구조는 그림 9와 같습니다. Pyraformer와 기존 Transformer와 다른 점은 coarser scale construction module (CSCM)과 Pyramidal Attention Module (PAM) 입니다.

그림 10. Coarser Scale Construction Module (CSCM)과 피라미드 구조

Pyraformer에서는 이름 그대로 compact multi-resolution representation을 만들기 위해 피라미드 구조를 활용했습니다. 이를 위해 CSCM에서는 그림 10과 같이 stride가 C인 Convolution을 반복 적용하여 C-ary Tree 구조, 즉 피라미드 구조를 생성합니다. 이후 PAM에서는 그림과 같이 CSCM에 의해 생성된 피라미드 구조에서 부모 노드  ${\mathrm{\mathbb{P}}}_{\mathcal{l}}^{\left(s\right)}$, 같은 scale에 있는 주변 노드 ${\mathbb{A}}_{\mathcal{l}}^{\left(s\right)}$, 자식 노드 ${\mathrm{ℂ}}_{\mathcal{l}}^{\left(s\right)}$의 제한된 범위의 attention 연산을 하게 됩니다. 구체적으로 노드 $n_{\mathcal{l}}^{\left(s\right)}$의 attention weight는 다음 식과 같이 계산합니다.

$y_i=\sum _{\mathcal{l}\in {\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}}\frac{\exp \left(\frac{{\mathit{q}}_{\mathit{i}}{\mathit{k}}_{\mathit{l}}^{\mathit{T}}}{\sqrt{{\mathit{d}}_{\mathit{K}}}}\right){\mathbf{v}}_{\mathcal{l}}}{\sum _{\mathcal{l}\in {\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}}\exp \left(\frac{{\mathit{q}}_i{\mathit{k}}_{\mathcal{l}}^T}{d_K}\right)}$ 

여기서 ${\mathrm{\mathbb{N}}}_{\mathcal{l}}^{\left(s\right)}={\mathrm{\mathbb{P}}}_{\mathcal{l}}^{\left(s\right)}\cup {\mathbb{A}}_{\mathcal{l}}^{\left(s\right)}\cup {\mathrm{ℂ}}_{\mathcal{l}}^{\left(s\right)}$입니다. 이를 통해 Pyraformer는 기존 Transformer에 비해 적은 개수의 Query-Key pair를 연산하게 되며, 최종적으로 O(N)의 연산량을 갖게 됩니다.

 

실험 결과에서 보면 Q-K pair 기준으로 보면 기존 방법에 비해 Pyraformer는 5~6배 정도 적은 pair만 계산하게 됩니다. 그럼에도 불구하고 Pyraformer의 성능은 다소 향상된 성능을 보임을 확인할 수 있습니다.

cosFormer

 

그림 11. cosFormer의 선형화된 self-attention

두 번째 방법은 cosFormer입니다. 기존 self-attention은 그림 11과 같이 1) Q (query)와 K (key)의 행렬 곱셈, 2) softmax 함수를 적용, 3) V (value)와 행렬 곱셈을 통해 attention weight를 계산하게 되며, 최종적으로 O(N²)의 연산량이 필요합니다. 한편 cosFormer에서는 그림과 같이 decomposable similarity function $\varphi \left(·\right)$ 을 활용하여 self-attention의 선형화 시켰습니다. 따라서 K (key)와 V (value)를 먼저 연산을 하고 나중에 Q (query)와 연산하여 연산량을 O(N)으로 감소시켰습니다.

 

먼저 cosFormer에서는 self-attention의 선형화를 위해 여러 decomposable linear function $\varphi \left(·\right)$과 softmax의 성능을 비교했습니다. 먼저 decomposable linear function $\varphi \left(·\right)$으로 identity 함수, leaky ReLU 함수, ReLU 함수을 고려해 보겠습니다. Loss와 세 가지 데이터셋 (QQP, SST-2, MNLI)에 대한 accuracy로 성능을 평가한 결과, 세 가지 함수 중 ReLU 함수가 Loss와 accuracy 관점에서 가장 좋은 성능을 보였습니다. 이 결과를 해석해보면 ReLU 함수는 연산 결과인 attention weight이 non-negativity를 만족하지만, identity 함수와 leaky ReLU 함수는 non-negativity를 만족하지 않습니다. 이를 통해 attention weight의 non-negativity가 성능에 큰 영향을 준다는 것을 유추할 수 있습니다. 또한 ReLU 함수와 softmax 함수의 성능 비교를 통해 softmax 함수의 non-linear re-weighting도 성능에 중요한 영향을 미치는 것을 확인할 수 있습니다.

위와 같은 관찰을 통해 cosFormer에서는 attention weight의 non-negativity는 ReLU 함수를 통해, non-linear re-weighting은 cos 함수 기반의 re-weighting을 활용하여 self-attention의 선형화를 해결하고자 했습니다. 구체적으로 먼저 다음과 같이 ReLU 함수를 통해 간단하게 key와 value를 먼저 계산할 수 있습니다.

$\mathrm{S}\left(Q,K\right)=s\mathit{(}{\varphi }_{linear}\mathit{\left(}Q\mathit{\right)}\mathit{,}{\varphi }_{linear}\mathit{\left(}K\mathit{\right)}\mathit{)}\mathit{=}s\mathit{(}\mathrm{ReLU}\mathit{\left(}Q\mathit{\right)}\mathit{,}\mathrm{ReLU}\mathit{\left(}K\mathit{\right)}\mathit{)}$
$o_i=\frac{\sum _{j=1}^N\left(\mathrm{ReLU}\left(Q_i\right)\mathrm{ReLU}{\left(K_j\right)}^T\right)V_j}{\sum _{j=1}^N\left(\mathrm{ReLU}\left(Q_i\right)\mathrm{ReLU}{\left(K_j\right)}^T\right)}=\frac{\mathrm{ReLU}\left(Q_i\right)\sum _{j=1}^N\mathrm{ReLU}{\left(K_j\right)}^T{V}_j}{\mathrm{ReLU}\left(Q_i\right)\sum _{j=1}^N\mathrm{ReLU}{\left(K_j\right)}^T}$

또한 다음과 같이 cos 함수의 re-weight을 통해 softmax 함수의 re-weighting을 대체하려고 했습니다. 여기서 cos 함수를 사용한 이유는 cos 합공식에 의해 두 개의 summation으로 나눌 수 있어서 선형화가 가능하기 때문입니다. 구체적으로 최종적인 cos 기반 re-weighting은 다음과 같이 적용하게 됩니다.

$s\left(Q_i^{\text{'}},K_j^{\text{'}}\right)=Q_i^{\text{'}}{K_j^{\text{'}}}^{T}cos\left(\frac{\pi }{2}\times \frac{i-j}{M}\right)$    $Q^{\text{'}}=\mathrm{ReLU}\left(Q\right), K^{\mathit{\text{'}}}\mathit{=}\mathrm{ReLU}\left(K\right)$

위의 식에서 거리가 먼 sequence 보다 가까운 sequence에 대해 더 많은 weight를 주기 위해 (i-j)의 함수로 되어 있습니다. 이를 cos 합 공식을 이용하여 더 전개해 보면 다음과 같습니다.

$$\begin{gathered} Q_i^{\text{'}}{K_j^{\text{'}}}^{T}cos\left(\frac{\pi }{2}\times \frac{i-j}{M}\right)=Q_i^{\text{'}}{K_j^{\text{'}}}^T\left(cos\left(\frac{\pi i}{2M}\right)cos\left(\frac{\pi j}{2M}\right)+sin\left(\frac{\pi i}{2M}\right)sin\left(\frac{\pi j}{2M}\right)\right) \\ =\left(Q_i^{\text{'}}cos\left(\frac{\pi i}{2M}\right)\right){\left(K_j^{\text{'}}cos\left(\frac{\pi j}{2M}\right)\right)}^T+\left(Q_i^{\text{'}}sin\left(\frac{\pi i}{2M}\right)\right){\left({K_j^{\text{'}}}^{T}sin\left(\frac{\pi j}{2M}\right)\right)}^T \end{gathered}$$ 

따라서 최종적으로 다음과 같이 attention weight를 구할 수 있습니다.

$$\begin{gathered} O_i=\frac{\sum _{j=1}^{ \\ N}{Q}_i^{\cos }\left({\left(K_j^{\cos }\right)}^T{V}_j\right)+\sum _{j=1}^N{Q}_i^{\sin }\left({\left(K_j^{\sin }\right)}^T{V}_j\right)}{\sum _{j=1}^N{Q}_i^{\cos }{\left(K_j^{\cos }\right)}^T+\sum _{j=1}^N{Q}_i^{\sin }{\left(K_j^{\sin }\right)}^T} \end{gathered}$$

여기서 $Q_i^{\cos }=Q_i^{\text{'}}\cos \left(\frac{\pi i}{2M}\right)$, $Q_i^{\sin }=Q_i^{\text{'}}\sin \left(\frac{\pi i}{2M}\right)$, $K_j^{\cos }=K_j^{\text{'}}\cos \left(\frac{\pi j}{2M}\right)$, $K_j^{\sin }=K_j^{\text{'}}\sin \left(\frac{\pi j}{2M}\right)$입니다.

 

그림 12. Transformer와 cosFormer의 attention weight의 visualization

그림 12는 Transformer와 cosFormer의 Attention weight를 visualization 한 결과를 보여주고 있습니다. 그림을 통해 1) cos re-weighting matrix를 통해 diagonal 근처에 더 많은 weight를 주고 있는 것을 확인할 수 있으며 이를 통해 거리가 먼 sequence 보다 거리가 가까운 sequence에 더 많은 weight를 준다는 것을 확인할 수 있습니다. 2) 또한, cosFormer를 통해 구한 attention weight는 기존의 Transformer와 비슷한 weight를 갖는 것을 확인할 수 있습니다. 3) 마지막으로, re-weighting을 하지 않은 것에 비해 기존 Transformer의 weight과 더 유사함을 확인할 수 있습니다. 이를 통해 ReLU 함수를 통한 decomposable similarity function과 cos 기반의 re-weighting이 softmax 함수를 대체할 수 있음을 실험적으로 확인할 수 있습니다.

 

또한 두가지 자연어 관련 실험 (Bidirectional finetune과 LRA (long-range-arena)) 을 통해 cos 함수 기반의 re-weighting의 효과를 검증했습니다. 위 테이블에서 확인할 수 있는 것과 같이 cos 함수의 re-weighting은 locality를 강조하며 이는 자연어 task에 큰 효과가 있다고 확인할 수 있습니다.

Pooling Network (PoNet)

그림 13. Pooling Network (PoNet)의 구조

마지막으로 self-attention의 구조를 효율적인 pooling network로 대체한 PoNet에 대해 알아보겠습니다. 구체적으로 PoNet에서는 Global aggregation (GA), Segment max-pooling (SMP), Local max-pooling (LMP) 라는 3가지 형태의 multi-granularity pooling mechanism을 활용하여 연산량을 크게 줄였습니다. GA, SMP, LMP 모듈은 pooling network를 통해 서로 다른 level의 contextual information을 추출하게 됩니다.

먼저 Global aggregation (GA)에서는 그림 13에서 표현된것과 같이 average pooling을 통해 전체 입력 신호의 평균 표현을 단일 query로 추출합니다. 단일 query만을 이용하여 attention을 연산하기 때문에 attention weight를 계산할 때 O(N)의 연산량만 필요하게 됩니다.

$\mathit{g}=\frac{1}{N}\sum _{n=1}^N{\mathit{h}}_Q$
${\mathit{g}}^{\text{'}}=\mathrm{Attention}\left\{\mathit{g},{\mathit{H}}_K,{\mathit{H}}_Q\right\}$

다음은 Segment max-pooling (SMP)입니다. GA에서 하나의 query만으로 self-attention을 수행했기 때문에 입력 sequence의 정보를 많이 잃어버린다는 단점이 있습니다. 이를 보완하기 위해 여러 sequence를 하나의 segment로 보고 segment level에서 중요한 정보를 추출하기 위해 segment 단위로 max-pooling을 연산합니다. 특히 SMP에서는 행렬의 곱셈 연산이 필요 없기 때문에 추가되는 연산량이 없습니다.

SMP와 비슷하게 Local max-pooling (LMP)에서는 window를 sliding 하면서 sequence와 인접 sequence 사이에 max-pooling을 적용하여 sequence level에서 중요한 정보를 추출하게 됩니다. 여기서도 SMP와 마찬가지로 max-pooling을 사용했기 때문에 추가되는 연산량이 없다는 장점이 있습니다.

마지막으로 GA, SMP, LMP에서 추출한 정보를 합쳐서 최종적으로 P의 결과를 도출하게 됩니다. 특히 GA에서 추출된 정보 g^'와 SMP에서 추출된 정보 S_k(n)는 원래의 sequence H_on와 element-wise multiplication을 하여 다양한 정보가 추출하여 학습할 수 있도록 합니다.

${\mathit{G}}_n={\mathit{g}}^{\text{'}}\circ {\mathit{H}}_{on}$
${\mathit{S}}_{\mathbf{n}}^{\mathbf{\text{'}}}={\mathit{S}}_{k\left(n\right)}\circ {\mathit{H}}_{on}$
$\mathit{P}={\mathit{G}}_n+{\mathit{S}}_n^{\text{'}}+\mathit{L}$

여기에서도 element-wise multiplication으로 인해 O(N)의 연산량만 필요하게 되고, 따라서 PoNet의 각 layer에서 요구되는 연산량은 O(N)이 됩니다.

 

실험 결과에서 보면 기존 Transformer에 비해 max-pooling mechanism을 활용한 PoNet이 학습 속도에서는 최대 9배, 메모리 사용량에서는 최대 10배가 절감된 것을 확인할 수 있습니다.

Conclusion

이번 블로그에서는 Transformer에 대해 간단히 살펴보고, ICLR 2022에 발표된 Transformer의 self-attention을 효율적으로 구현하는 세 가지 방법에 대해 알아보았습니다. Transformer는 다양한 응용에서 활용될 수 있는데 특히 수요예측 시스템의 성능은 과거 수요 신호를 얼마나 길게 보느냐에 따라 크게 달라지게 됩니다. 따라서 이러한 효율적인 Transformer를 수요예측 시스템에 적용하면 기존 Transformer의 seq-to-seq 모델링의 우수한 성능은 유지하면서도 효율적으로 구현할 수 있어서 많은 장점이 있다고 생각됩니다.