Dual Chunk Attention
- Related Project: Private
- Category: Paper Review
- Date: 2024-07-14
Training-Free Long-Context Scaling of Large Language Models
- url: https://arxiv.org/abs/2402.17463
- pdf: https://arxiv.org/pdf/2402.17463
- html: https://arxiv.org/html/2402.17463v2
- model: chunk-llama
- abstract: The ability of Large Language Models (LLMs) to process and generate coherent text is markedly weakened when the number of input tokens exceeds their pretraining length. Given the expensive overhead of finetuning large-scale models with longer sequences, we propose Dual Chunk Attention (DCA), which enables Llama2 70B to support context windows of more than 100k tokens without continual training. By decomposing the attention computation for long sequences into chunk-based modules, DCA manages to effectively capture the relative positional information of tokens within the same chunk (Intra-Chunk) and across distinct chunks (Inter-Chunk), as well as integrates seamlessly with Flash Attention. In addition to its impressive extrapolation capability, DCA achieves performance on practical long-context tasks that is comparable to or even better than that of finetuned models. When compared with proprietary models, our training-free 70B model attains 94% of the performance of gpt-3.5-16k, indicating it is a viable open-source alternative. All code and data used in this work are released at url.
Contents
TL;DR
- 본 논문은 대규모 언어모델의 컨텍스트 윈도우를 확장하기 위한 학습 없는 새로운 프레임워크, Dual Chunk Attention(DCA)을 소개합니다.
- DCA는 긴 컨텍스트를 효과적으로 처리하며, 이전에 학습된 범위를 넘어서는 입력 길이에서도 낮은 Perplexity(PPL)를 유지합니다.
- 본 모델은 70B 크기의 모델에서 언어 모델링, 키워드 검색 및 실세계의 긴 컨텍스트 애플리케이션을 포함한 다양한 작업에서 향상된 성능을 보여줍니다.
1. 서론
본 논문에서는 대규모 언어모델(LLMs)의 긴 컨텍스트 정보 처리 능력의 중요성을 강조합니다. 최근 연구에 따르면, 짧은 컨텍스트 모델을 긴 텍스트 시퀀스로 추가 학습하는 것이 긴 컨텍스트 능력을 향상시킬 수 있음을 보여주고 있습니다. 이런 접근 방식은 Llama2 Long 모델에서 두드러지게 나타나며, 이는 긴 텍스트 데이터와 원래 Llama2의 사전 학습 코퍼스를 혼합하여 학습되었습니다. 그러나, training dataset의 접근성 제한과 긴 컨텍스트의 파인튜닝 비용 때문에 현재의 오픈 소스 모델들은 성능면에서 종종 제한을 받고 있습니다.
2. 배경
2.1 위치 인코딩
트랜스포머 모델에서의 원래 위치 임베딩은 절대 위치 인덱스를 특징 공간으로 매핑하며, 입력 계층에 이 정보를 통합합니다. $\mathbf{x}_i = \mathbf{x} + f(i)$는 위치 인덱스 $i$와 연관된 입력이며, $f: \mathbb{N} \to \mathbb{R}^d$는 위치 임베딩 함수입니다. Rotary Positional Encoding (RoPE)는 이런 위치 정보를 입력 계층이 아닌 어텐션 계층에 직접 통합하는 방식으로, 위치 인덱스를 이용하여 쿼리 및 키 벡터를 수정합니다. 이를 통해 쿼리와 키 사이의 상대적 위치 정보만을 포함하는 내적을 계산할 수 있습니다.
2.2 RoPE의 확장
최근 연구에서는 RoPE가 사전 학습 단계에서 본 적 없는 상대 위치에 대한 강인한 길이 확장 능력이 부족하다는 것을 보여주었습니다. 이 문제를 해결하기 위해, 위치 인덱스를 조정하여 사전 학습 중 관찰된 컨텍스트 길이 내에 떨어지도록 하는 방법들이 제안되었습니다. 예를 들어, Position Interpolation(PI)은 위치 인덱스를 스케일링하여 상대 위치 행렬을 조정합니다.
3. 방법
3.1 내부 청크 어텐션 (Intra-Chunk Attention)
내부 청크 어텐션은 같은 청크 내의 쿼리와 키 사이의 내적을 계산하는 방법으로 긴 시퀀스를 여러 청크로 분할하고, 각 청크 내의 위치 인덱스가 청크 크기를 초과하지 않도록 합니다. 예를 들어, 12개의 토큰으로 구성된 시퀀스가 청크 크기 6인 2개의 청크로 분할됩니다. 쿼리와 키의 위치 인덱스는 청크 크기 내에서 조정됩니다. 수식은 다음과 같습니다.
\[\mathbf{P}_{\text{qIntra}} = \mathbf{P}_k = [0, 1, \ldots, l-1] \mod s\]위치 인덱스 $i$와 $j$가 같은 청크에 있는 경우, 즉 $\lfloor i/s \rfloor = \lfloor j/s \rfloor$이면, 쿼리와 키 사이의 상대적 위치는 다음과 같이 계산됩니다.
\[\mathbf{M}[i][j] = \mathbf{P}_{\text{qIntra}}[i] - \mathbf{P}_k[j]\]이때, $\mathbf{q}_i^\top \mathbf{k}_j$의 계산은 다음과 같이 이루어집니다.
\[\mathbf{q}_i^\top \mathbf{k}_j = f(\mathbf{q}, \mathbf{P}_{\text{qIntra}}[i])^\top f(\mathbf{k}, \mathbf{P}_k[j])\]3.2 청크 간 어텐션 (Inter-Chunk Attention)
다른 청크 간의 정보를 집약하기 위해 청크 간 어텐션를 도입합니다. 기본적으로 쿼리의 위치 인덱스는 키의 위치 인덱스보다 커야 합니다. 예를 들어 청크 크기가 $s$일 때, 청크 0의 마지막 키 위치는 $s-1$이고 청크 1의 첫 번째 쿼리 위치는 $c-1$로 설정하여 청크 간의 연속성을 유지합니다. 상대 위치 행렬 $\mathbf{M}$은 다음과 같이 계산됩니다.
\[\mathbf{M}[i][j] = \mathbf{P}_{\text{qIntra}}[i] - \mathbf{P}_k[j] = c-1 - \mathbf{P}_k[j] \geq c-s\]이는 $\mathbf{P}_{\text{qInter}}$의 모든 요소가 $c-1$로 설정되어 있음을 의미합니다.
3.3 연속 청크 어텐션 (Successive-Chunk Attention)
인접 청크 간의 국소성을 유지하기 위해 연속 청크 어텐션가 도입되며, 특히 두 청크가 인접할 때 적용됩니다. 청크 0의 마지막 키와 청크 1의 첫 번째 쿼리 사이의 상대적 거리는 \(\mathbf{P}_{\text{qInter}}[6] - \mathbf{P}_k[5] = 4\)로 계산되며, 이는 절대 거리와 일치하지 않습니다. 연속 청크 어텐션를 사용하여 이런 문제를 해결하고, 인접 토큰 사이의 국소성을 유지합니다. \(\mathbf{P}_{\text{qSucc}}\)는 다음과 같이 설정되고,
\[\mathbf{P}_{\text{qSucc}} = [s, s+1, \ldots, s+w-1, c-1, \ldots, c-1]\]이를 통해 인접 청크 간의 연속성을 유지하면서 $\mathbf{M}[i][j]$를 계산할 수 있습니다.
3.4 정규화 및 Flash Attention
내적 계산은 소프트맥스 함수를 사용하여 정규화됩니다. 이는 다음과 같이 정의됩니다.
\[\mathbf{p}_i = \text{softmax}([\mathbf{q}_i^\top \mathbf{k}_0, \mathbf{q}_i^\top \mathbf{k}_1, \ldots, \mathbf{q}_i^\top \mathbf{k}_i])\]Flash Attention을 통해 DCA는 Llama의 원래 self-attention과 비교할 때 유사한 GPU 메모리 사용량 및 인퍼런스 속도를 달성합니다. 결과는 Figure 3에서 확인할 수 있습니다. 이런 구조적 접근 방식은 LLMs의 긴 컨텍스트 이해력을 향상시키는 데 중요한 역할을 합니다.
4. 실험
다양한 벤치마크에서 DCA의 성능을 평가한 결과, 언어 모델링, 키워드 검색, 긴 컨텍스트 애플리케이션에서 우수한 결과를 보여줍니다. DCA는 학습 없이도 4k context window을 32k 이상으로 확장할 수 있음을 입증하며, 기존의 방법들이 8k 이상의 컨텍스트 길이에서 성능 저하를 경험하는 반면, 본 모델은 100k 토큰을 초과하는 컨텍스트 크기를 처리할 수 있는 향상된 확장 능력을 보여줍니다.
[참고자료 1] 예시로 이해하기
Dual Chunk Attention (DCA)은 대규모 언어모델의 긴 컨텍스트 처리 능력을 향상시키기 위한 방법으로, 내부 청크 어텐션, 청크 간 어텐션, 그리고 연속 청크 어텐션 세 가지 주요 구성 요소를 포함합니다. 이 세 가지 요소는 모두 특정한 위치 인코딩 방식을 사용하여 시퀀스 내의 상대적 위치 정보를 보존하면서 효율적인 어텐션 기반 연산을 가능하게 합니다. 결론적으로 기존의 어텐션 메커니즘보다 더 정교하게 긴 컨텍스트의 정보를 캡처할 수 있도록 설계합니다.
1. Intra-Chunk Attention
12개 토큰의 시퀀스가 주어지고, 청크 크기를 6으로 설정했다고 가정하면 이 시퀀스는 두 개의 청크, 각각 [0, 1, 2, 3, 4, 5]과 [6, 7, 8, 9, 10, 11]으로 나뉩니다.
각 청크 내에서 쿼리와 키의 위치 인덱스는 [0, 1, 2, 3, 4, 5]로 설정됩니다. 이때 청크 내에서 첫 번째 쿼리와 네 번째 키의 내적을 계산하려면 위치 인덱스가 각각 0과 3인 쿼리와 키를 사용합니다. 즉, $\mathbf{q}_0^\top \mathbf{k}_3$를 계산하며, 이들 간의 상대적 위치는 0 - 3 = -3이 됩니다.
2. Inter-Chunk Attention
청크 간 어텐션은 다른 청크 간의 정보를 통합합니다.
위의 예에서, 청크 0의 마지막 키 위치는 5이고, 청크 1의 첫 번째 쿼리 위치는 가상적으로 10으로 설정할 수 있습니다. 이 설정은 $\mathbf{q}_6$ (청크 1의 첫 번째 쿼리)과 $\mathbf{k}_5$ (청크 0의 마지막 키) 사이의 상대 위치를 계산할 때 사용됩니다. 만약 $\mathbf{q}_6$의 위치를 인위적으로 9로 설정하고, $\mathbf{k}_5$의 위치가 5라면, 그들 사이의 상대적 위치는 9 - 5 = 4가 됩니다.
3. Successive-Chunk Attention
연속 청크 어텐션은 인접한 청크 간의 정보를 더 정확하게 연결하기 위해 사용됩니다.
이전 예제에서, 청크 1의 첫 번째 쿼리 위치를 실제 청크 위치보다 크게 설정하여, 청크 0의 마지막 키와 연결하는 방식입니다. 이 경우, 청크 0의 마지막 키 위치를 5, 청크 1의 첫 번째 쿼리 위치를 6으로 설정하여, $\mathbf{q}_6$과 $\mathbf{k}_5$ 사이의 상대적 위치를 6 - 5 = 1로 계산할 수 있습니다. 이렇게 조정함으로써, 모델은 두 인접 청크 사이에서 정보 흐름을 더 잘 유지할 수 있습니다.
4. 정규화 및 Flash Attention
계산된 내적 값은 소프트맥스 함수를 통해 확률값으로 변환하여 각 쿼리에 대해 가장 관련 있는 키를 선택하는 데 도움을 줍니다.
예를 들어, 내적 값이 [-2, 4, 1, -3]라면, 소프트맥스 함수를 적용하여 이 값들을 [0.02, 0.88, 0.07, 0.02]와 같은 확률 분포로 변환할 수 있습니다. 이 분포는 각 쿼리가 어텐션를 기울여야 할 키를 표현합니다.