[논문 리뷰] Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context

참고 사항 - 아직 작성 중인 글 입니다!!

0. 요약

Transformer-XL(Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context, 2019) 논문은 segment-level recurrence mechanism을 도입하여 입력 문맥 길이를 확장함으로써, 고정된 길이를 넘어서는 장거리 의존성을 학습 가능하게 한 논문이라고 요약할 수 있습니다. 여기에, 하나의 문장을 여러개로 분할하는 segmentation 과정에서 발생할 수 있는 위치 정보에 대한 손실을 어느정도 해결 해 줄 수 있는 새로운 positional encoding 방법을 제안합니다. 논문의 저자들에 의하면, RNN 모델 대비 80% , 기존 Vanilla Transformer 모델 대비 450% 정도의 장거리 의존성 학습 향상이 되었다고 합니다. 또한 Vanilla Transformer 보다 evaluation task에서 더 빠르다고 합니다.

 

본 논문은 Character-Level Language Modeling with Deeper Self-Attention 논문과 관련이 깊다.

1. Problem : 기존 모델(RNN/Vanilla Transformer)의 한계?

RNN 모델은 Long Term Dependency 문제로 인해 LSTM이나 GRU 모델과 같이 다양한 파생 모델이 제안되었습니다. 하지만, 제안된 모델도 gradient vanishing 과 gradient explosion 문제를 완전히 해결하기에는 힘들었습니다. 한편, Transformer 모델은 RNN 모델의 문제점 중 하나인 Long Term Dependency(장거리 의존성 문제)를 Attention mechanism을 도입하여 어느정도 해결 한 모델로 평가받습니다. 하지만, Transformer 모델은 한정된 길이의 토큰을 입력으로 받기 때문에, 이로 인해서 몇가지 문제가 발생합니다.

첫 번째 문제는 고정된 입력 길이로 인해 입력 문장이 분리되는(잘리는) 경우에, 고정된 길이를 넘어서는 장거리 의존성에 대해 모델이 학습할 수 없습니다. 두 번째로, 고정된 입력 길이로 인해, 입력 segment를 생성할 때 문장이나 다른 어떤 의미적 경계(boundary)를 고려하지 않은 채, 연속된 기호(단어나 문자 등)를 선택하여 생성합니다. 이로 인해 모델은 초기 몇개의 기호를 정확하게 예측하는  데 필요한 문맥의 정보가 부족해지기 때문에 최적화의 비효율성 문제와 성능 저하의 문제가 생기게 됩니다. 본 논문의 저자는 이러한 문제를 'Context Fragmentation(문맥 단절)' 이라고 부릅니다. 저자들은 앞서 언급한 문제를 해결하기 위해 Transformer-XL 모델을 제안합니다.

2. 문제를 해결하기 위한 전략 세우기

Transformer나 self-attention을 언어 모델로 적용하기 위해서는 어떻게 Transformer 모델에 효과적으로 임의의 긴 문장을 고정된 표현으로 인코딩(Encoding)할 것인가 입니다. 만약 메모리의 크기가 무한하다면 매우 긴 문장을 그냥 입력하면 간단하지만, 유한한 컴퓨팅 자원으로 인해 유한한 길이의 문장을 입력 할 수 밖에 없습니다. 이 문제에 대한 아주 간단한 해결 방법은 전체 corpus를 관리 가능한 작은 segment로 분리하고, 분리된 segment 각각을 모델 학습에 사용하는 방법입니다. 하지만, 이러한 방법은 (1) 최대 의존 길이가(참고할 수 있는 최대 길이가) 몇 백개의 문자 수준정도 이며, (2)문장이나 다른 의미 경계 측면에

3. Context Fragmentation을 해결하기 위한 노력 : Segment-level Recurrence Mechanism 도입

논문의 저자들은 고정된 컨텍스트 길이의 한계를 해결하기 위해 Recurrence Mechanism을 제안하게 됩니다. 여기서 Recurrence mechanism이란, 이전 segment에서 계산된 hidden state sequence를 고정하고 따로 저장해서 다음에 새로 들어오는 segment를 계산할 때 다시 사용하는 방법입니다. 그림으로 표현하면 다음과 같습니다.

Figure 2: Illustration of the Transformer-XL model with a segment length 4.

그림은 segment length가 4인 경우에 대해 recurrence mechanism을 도입한 Transformer-XL의 학습 방법에 대해 표현하고 있습니다. 기존 vanilla Transformer는 새롭게 입력된 new segment가 이전에 계산된 segment의 hidden state를 사용하고 있음을 알 수 있습니다.

 

저자가 제안한 방법을 수학적으로 표현하면 다음과 같습니다.

segment length = \( L \) 인 \(\tau\)번째 segment는 token \(x\)에 대해 다음과 같이 적을 수 있습니다.

\(s_{\tau} = [x_{\tau,1},..., x_{\tau,L}]\)

\( s_{\tau+1} = [x_{\tau+1,1},..., x_{\tau+1,L}]\)

이에 대해 \(n\)번째 layer의 hidden state를 \( h^{n}_{\tau} \in \mathbb{R}^{L \times d}\) (d = hidden dimension)라고 하였을때, \(s_{\tau+1}\)에 의해 생성되는 n번째 hidden state는 다음과 같습니다.

\( \tilde{\mathbf{h}}_{\tau+1}^{n-1} = \left[\text{SG}\left(\mathbf{h}_\tau^{n-1}\right) \circ \mathbf{h}_{\tau+1}^{n-1}\right] \) 

위 수식에 대해 첨부하면 

• SG(⋅): Stop-Gradient 연산, 이전 segment의 hidden state를 고정하여 gradient가 흐르지 않도록 합니다.
• \(\circ\): Concatenation 연산, 두 개의 hidden sequence를 sequence length dim을 통해 연결합니다

 

\(\mathbf{q}_{\tau+1}^n,\mathbf{k}_{\tau+1}^n,\mathbf{v}_{\tau+1}^n = \mathbf{h}_{\tau+1}^{n-1} \mathbf{W}_q,\tilde{\mathbf{h}}_{\tau+1}^{n-1} \mathbf{W}_k,\tilde{\mathbf{h}}_{\tau+1}^{n-1} \mathbf{W}_v \)

여기서 눈여겨 봐야 할 점은 Key와 Value는 이전 hidden state를 concatenation한 hidden state를 \(\tilde{h}\) 를 사용하지만, Query는 현재 세그먼트에서만 생성된다는 점입니다. 이렇게 계산된 q,k,v를 이용하여 최종적으로 n-th layer의 hidden state를 계산합니다.

\( \mathbf{h}_{\tau+1}^n = \text{Transformer-Layer}(\mathbf{q}_{\tau+1}^n, \mathbf{k}_{\tau+1}^n, \mathbf{v}_{\tau+1}^n) \)

 

이 수식을 표현은 그림 Figure 2에서 녹색 선으로도 확인 할 수 있습니다.

 

저자들이 제안한 이전 hidden state를 다시 사용하는 특별한 메커니즘은 오직 연속된 2개의 segment에서만 적용됩니다. 또한 기존 RNN기반의 언어 모델처럼 하나의 layer를 여러번 다시 사용하는 것과 달리, 두 개의 segment 사이에서 layer가 하나 아래로 이동하면서 학습하는 방식에서 차이가 있습니다. 이러한 구조의 recurrence scheme는 특히 평가에 있어서 더 빠르게 동작할 수 있으며, 기존 Vanilla Transformer보다 1,800배 더 빠릅니다.

 

마지막으로, 메모리 측면에서 저자들이 제안한 방법은 이전 segment에 대한 정보를 따로 저장해서 학습에 사용해야 하므로 저장공간이 많이 필요할 것으로 보입니다. 하지만, 저자들은 \( M \)개의 hidden state만을 저장하는 방식으로 기억해야 하는 hidden state의 양을 조절하는 방법을 채택하였습니다. 따라서 모델은 오래된 정보는 잊는 대신 계속해서 최신 정보를 얻을 수 있게 됩니다.

4. Context Fragmentation을 해결하기 위한 노력 : Relative Positional Encoding 도입

저자들이 제안한 Recurrence scheme는 이전 hidden state를 재사용하는 새로운 방법입니다. 하지만, 여기서 중요한 문제는 재 사용하는 hidden state의 위치 정보를 어떻게 유지할 것인가입니다. Vanilla Transformer 모델은 sinusoidal positional encoding을 word-embedding에 element-wise하게 더하는 방법을 사용하였습니다. 이 방법을 그대로 활용하여 저자들이 제안한 Recurrence mechanism에 적용하여 표현하면 다음과 같습니다.

\( h_{\tau+1} = f(h_{\tau},E_{s_{\tau+1}} + U_{1:L})\)

\( h_{\tau} = f(h_{\tau-1},E_{s_{\tau}} + U_{1:L})\)

여기서, \(f\)는 transformation funtion,  E는 Word Embedding , U는 positional encoding을 의미합니다.

수식에서 주목할만한 점은 hidden state \(h_{\tau+1}\) 과 \(h_{\tau}\) 에서 같은 Positional encoding을 사용한다는 점입니다. 이렇게 되면, 모델은 두 segment의 i번째 token순서(\(x_{\tau+1,i}\)와\( x_{\tau,i}\))를 구별할 만한 정보가 없기 때문에 성능 하락의 요인이 될 수 있습니다. 따라서 이러한 부분을 해결하기 위해서는 hidden state에 '상대적인' 위치 정보를 주어야 합니다.

 

positional encoding은 개념적으로 모델에게 어떻게 정보들을 모아야 하는지에 대한 일시적인 '단서'를 제공하는 개념입니다. 저자들은 이러한 '단서'를 Embedding에 element-wise하게 더하지 않고, 각 layer마다 attention score에 '단서'를 제공하는 방식을 도입하였습니다. Attention Score를 계산할때 Query는 일시적인 Key vector의 순서를 구별하기 위해 key의 절대적인 위치정보를 사용하는 대신, \(q_{\tau,i}\)와 \(k_{\tau,j}\) 사이의 \(i-j\)만큼의 상대적인 거리 정보만을 사용합니다. 이 방식을 이용하여 relateive positional encoding (\(R \in \mathbb{R}^{L_{max} \times d}\))을 만들 수 있습니다. 여기서 \(i\)번째 row인 \( R_i \)는 \( q_i \) 사이의 위치를 상대적으로 나타냅니다.

 

우선, Standard Transformer 모델의 Attention Score는 다음과 같은 수식으로 표현할 수 있습니다.

하나의 segment에 포함되어 있는 token에서, Query \(q_i\) 와 Key \(k_j\)의 Attention score는

\(A^{abs}_{i,j} = \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{E}_{x_j} 
+ \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{U}_{j} 
+ \mathbf{U}_{i}^\top \mathbf{W}_q^\top\mathbf{W}_{k} \mathbf{E}_{x_j} 
+ \mathbf{U}_{i}^\top \mathbf{W}_q^\top\mathbf{W}_k \mathbf{U}_{j}\) 입니다

---

잠시 쉬어가며, 어떻게 Attention score가 위와 같이 decomposition 될 수 있는지 살펴보겠습니다.

일반적으로 Attention score는 다음 수식으로 기억하는 분들도 있을 것 입니다.

\(Attention Score(Q,K) =  \left( \frac{QK^\top}{\sqrt{d_k}} \right)\)

이제 이 식을 query의 \(i\) 번째 원소 \(q_i\)와 key의 \(j\)번째 원소 \(k_j\)에 대해 다시 작성하면

\(Attention Score(q_i,k_j) = \left( \frac{q^\top_{i}k_j}{\sqrt{d_k}} \right)\) 가 됩니다.

여기서 같은 segment에 포함되는 두 개의 token을  각각 \(x_i,x_j\)라고 하고, 이들의 word embeddding을 각각 \(E_{x_i} ,E_{x_j}\), Positional Encoding을 \(U_i,U_j\)라고 하였을 때, Vanilla Transformer에서 입력은 word embedding + positional encoding의 합이므로 다음과 같이 적을 수 있습니다.

\(q_i = W_{q}(E_{x_i}+U_i) = W_qE_{x_i} + W_qU_i \)

\(k_j = W_{k}(E_{x_j}+U_j) = W_kE_{x_j} + W_kU_j  \) 

이때, Attention score은 상수 \(d_k\)를 제외하면 \(A_{i,j} = q^\top_{i}k_j \) 이므로 다음과 같이 적을 수 있습니다.

\(A_{i,j} = q^\top_{i}k_j \)

\(A_{i,j} = (W_qE_{x_i} + W_qU_i)^\top(W_kE_{x_j} + W_kU_j) \)

이때, 두 Matrix \(A,B\)에 대해 \( (AB)^\top = B^\top A^\top \)와 \((A+B)^\top = A^\top+B^\top \)이 성립하므로 다음과 같이 작성할 수 있습니다.

\(A_{i,j} = (E_{x_i}^\top W_q^\top + U_i^\top W_q^\top)(W_kE_{x_j} + W_kU_j)\)

\(A_{i,j} = \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{E}_{x_j} 
+ \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{U}_{j} 
+ \mathbf{U}_{i}^\top \mathbf{W}_q^\top\mathbf{W}_{k} \mathbf{E}_{x_j} 
+ \mathbf{U}_{i}^\top \mathbf{W}_q^\top\mathbf{W}_k \mathbf{U}_{j}\) 

---

다음으로, 이렇게 분

\(A_{i,j} = \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{E}_{x_j}
+ \mathbf{E}_{x_i}^\top \mathbf{W}_q^\top \mathbf{W}_k \mathbf{R}_{i-j} 
+ \mathbf{u}^\top \mathbf{W}_k \mathbf{E}_{x_j} 
+ \mathbf{v}^\top \mathbf{W}_k \mathbf{R}_{i-j}\)