​신세대 어텐션 메커니즘 Lightning Attention-2: 무제한 시퀀스 길이, 일정한 컴퓨팅 성능 오버헤드, 더 높은 모델링 정확도

현재 대규모 언어 모델의 적용은 시퀀스 길이 제한으로 인해 인공 지능 분야에서의 적용이 제한됩니다. 예를 들어, 다중 회전 대화, 긴 텍스트 이해, 다중 모드 데이터 처리 및 생성에는 특정 과제가 있습니다. 이러한 제한의 근본적인 이유는 현재 대규모 언어 모델에 일반적으로 사용되는 Transformer 아키텍처가 계산 복잡성과 시퀀스 길이 사이에 2차 관계를 갖고 있기 때문입니다. 따라서 시퀀스 길이가 길어질수록 컴퓨팅 리소스에 대한 수요도 기하급수적으로 증가합니다. 따라서 긴 시퀀스를 효율적으로 처리하는 방법은 항상 대규모 언어 모델이 직면한 과제 중 하나였습니다.

과거의 접근 방식은 주로 추론 단계에서 대규모 언어 모델을 더 긴 시퀀스에 적용하는 데 중점을 두었습니다. 한 가지 접근 방식은 모델이 다양한 길이의 입력 시퀀스에 적응할 수 있도록 Alibi 또는 유사한 상대 위치 인코딩을 사용하는 것입니다. 또 다른 접근 방식은 차이에 대해 RoPE 또는 유사한 상대 위치 인코딩을 사용하고 이미 훈련된 모델을 간략하게 미세 조정하여 시퀀스 길이를 연장하는 것입니다. 이러한 방법을 사용하면 대규모 모델이 특정 긴 시퀀스 모델링 기능을 가질 수 있지만 훈련 및 추론의 오버헤드는 줄어들지 않았습니다.

OpenNLPLab 팀은 대규모 언어 모델의 긴 시퀀스 문제를 해결하도록 설계된 Lightning Attention-2라는 새로운 선형 주의 메커니즘을 오픈 소스로 공개했습니다. 이 메커니즘은 긴 시퀀스를 훈련하고 추론하는 비용을 1K 시퀀스 길이와 일관되게 유지하여 설정 후 잊어버리는 솔루션을 가능하게 합니다. 시퀀스 길이를 늘려도 메모리 병목 현상이 발생하기 전에도 모델 훈련 속도에 부정적인 영향을 미치지 않으므로 무제한 길이의 사전 훈련이 가능합니다. 또한 매우 긴 텍스트의 추론 비용도 1K 토큰에 비해 일관되거나 훨씬 낮으므로 현재 대규모 언어 모델의 추론 비용을 크게 줄입니다. 아래 그림과 같이 모델 크기가 400M, 1B, 3B인 경우 시퀀스 길이가 증가함에 따라 FlashAttention2가 지원하는 LLaMA의 학습 속도는 급격히 감소하기 시작하는 반면 Lightning Attention-2가 지원하는 TansNormerLLM의 속도는 거의 변경 없음.

그림 1

• 논문: Lightning Attention-2: 대규모 언어 모델의 무제한 시퀀스 길이 처리를 위한 무료 점심
• 논문 주소: https:// arxiv.org/pdf/2401.04658.pdf
• 오픈 소스 주소: https://github.com/OpenNLPLab/lightning-attention

Lightning Attention-2 소개

모델 다양한 시퀀스 길이에 걸쳐 사전 훈련 속도를 일관되게 유지하는 것은 불가능한 작업처럼 들립니다. 그러나 2020년 선형 주의가 출현한 이후 연구자들은 선형 주의의 실제 효율성을 이론적 선형 계산 복잡성과 일치시키기 위해 열심히 노력해 왔습니다. 2023년 중반까지 선형 주의에 대한 연구는 주로 Transformer 아키텍처와의 정확도 정렬에 중점을 둘 것입니다. 마지막으로, 개선된 선형 주의 메커니즘이 등장한 후 정확성 측면에서 최첨단 Transformer 아키텍처와 비교할 수 있습니다. 그러나 선형 주의에서 "왼쪽 곱셈에서 오른쪽 곱셈"의 가장 중요한 계산 트릭은 실제 구현에서 직접 왼쪽 곱셈 알고리즘보다 훨씬 느립니다. 이는 오른쪽 곱셈을 구현하려면 많은 수의 루프 연산을 포함하는 누적 합계(cumsum)를 사용해야 하고, 많은 수의 I/O 연산으로 인해 오른쪽 곱셈이 왼쪽 곱셈보다 훨씬 덜 효율적이기 때문입니다. 따라서 다양한 시퀀스 길이에 걸쳐 대규모 모델의 사전 학습 속도를 일관되게 유지하는 것은 여전히 ​​​​어려운 과제입니다. 연구자들은 계산 효율성을 향상하고 I/O 작업을 줄이기 위해 선형 주의 구현을 추가로 탐색하고 개선해야 합니다. 이는 다양한 시퀀스 길이의 작업 요구 사항에 더 잘 대처할 수 있도록 사전 훈련 속도의 일관성을 달성하는 데 도움이 됩니다.

그림 2

Lightning Attention-2의 아이디어를 더 잘 이해하기 위해 먼저 전통적인 소프트맥스 어텐션의 계산 공식을 검토해 보겠습니다. O=softmax ((QK^T) M_) V, 여기서 Q, K, V, M 및 O는 각각 쿼리, 키, 값, 마스크 및 출력 행렬입니다. 여기서 M은 단방향 작업(예: GPT)의 하삼각 All-1 행렬입니다. Bert와 같은 양방향 작업은 무시할 수 있습니다. 즉, 양방향 작업에 대한 마스크 매트릭스가 없습니다.

저자는 설명을 위해 Lightning Attention-2의 전반적인 아이디어를 다음 세 가지로 요약합니다.

1. Linear Attention의 핵심 아이디어 중 하나는 계산 비용이 많이 드는 소프트맥스 연산자를 제거하여 Attention의 계산 공식을 O=((QK^T)⊙M_) V로 작성할 수 있다는 것입니다. 그러나 단방향 작업에 마스크 행렬 M이 존재하기 때문에 이 형식은 여전히 ​​왼쪽 곱셈 계산만 수행할 수 있으므로 O(N)의 복잡도를 얻을 수 없습니다. 그러나 양방향 작업의 경우 마스크 행렬이 없으므로 선형 주의 계산 공식은 O=(QK^T) V로 더욱 단순화될 수 있습니다. Linear Attention의 미묘함은 단순히 행렬 곱셈의 결합 법칙을 사용하여 계산 공식을 다음과 같이 변환할 수 있다는 것입니다. 이 계산 형식을 오른쪽 곱셈이라고 하며 해당 전자는 다음과 같습니다. 왼쪽. 그림 2에서 우리는 Linear Attention이 양방향 작업에서 매력적인 O(N) 복잡성을 달성할 수 있다는 것을 직관적으로 이해할 수 있습니다!

2. 그러나 디코더 전용 GPT 모델이 점차 LLM의 사실상 표준이 되면서 단방향 작업을 가속화하기 위해 Linear Attention의 올바른 곱셈 기능을 어떻게 사용하는 것이 시급한 문제가 되었습니다. 이 문제를 해결하기 위해 이 기사의 저자는 주의 행렬의 계산을 대각 행렬과 비대각선 행렬의 두 가지 형태로 나누고 서로 다른 방식을 사용하는 "분할 정복" 아이디어를 제안했습니다. 그들을 계산하는 방법. 그림 3에서 볼 수 있듯이 Linear Attention-2는 컴퓨터 분야에서 일반적으로 사용되는 Tiling 아이디어를 사용하여 Q, K 및 V 행렬을 동일한 수의 블록으로 나눕니다. 그 중 블록 자체(블록 내)의 계산은 마스크 행렬의 존재로 인해 여전히 왼쪽 곱셈 계산 방법을 유지하며, 블록(인터-블록)의 계산은 O(N^2)입니다. 블록)에는 마스크 행렬이 없으므로 올바른 곱셈 계산 방법을 사용하여 O(N)의 복잡성을 즐길 수 있습니다. 두 개를 개별적으로 계산한 후 직접 추가하여 i번째 블록에 해당하는 선형 주의 출력 Oi를 얻을 수 있습니다. 동시에 KV의 상태는 Cumsum을 통해 누적되어 다음 블록 계산에 사용됩니다. 이처럼 Lightning Attention-2 전체의 알고리즘 복잡도는 블록 내에서는 O(N^2), 블록 간 Trade-off에서는 O(N)이다. 더 나은 절충점을 얻는 방법은 타일링의 블록 크기에 따라 결정됩니다.

3. 주의깊은 독자라면 위의 프로세스가 Lightning Attention-2의 알고리즘 부분일 뿐이라는 것을 알게 될 것입니다. 이름이 Lightning인 이유는 저자가 GPU 하드웨어 실행 프로세스에서 알고리즘 프로세스의 효율성을 충분히 고려했기 때문입니다. . FlashAttention 시리즈 작업에서 영감을 받아 GPU에서 실제로 계산을 수행할 때 작성자는 분할된 Q_i, K_i, V_i 텐서를 GPU 내부의 더 큰 용량을 가진 느린 HBM에서 더 작은 용량의 더 빠른 SRAM으로 옮겼습니다. 시스템을 사용하여 대량의 메모리 IO 오버헤드를 줄입니다. 블록이 Linear Attention 계산을 완료한 후 출력 결과 O_i가 HBM으로 다시 이동됩니다. 모든 블록이 처리될 때까지 이 과정을 반복합니다.

더 자세한 내용을 알고 싶은 독자들은 이 글의 알고리즘 1과 알고리즘 2는 물론, 논문의 자세한 도출 과정도 꼼꼼히 읽어보시면 됩니다. 알고리즘과 파생 프로세스 모두 Lightning Attention-2의 정방향 프로세스와 역방향 프로세스를 구별하여 독자가 더 깊은 이해를 할 수 있도록 도와줍니다.

그림 3

Lightning Attention-2 정확도 비교

연구진은 먼저 소규모(400M) 매개변수에서 Lightning Attention-2와 L을 비교했습니다. 모델 번개 주의 -1 정확도 차이는 아래 그림과 같이 둘 사이에는 거의 차이가 없습니다.

그런 다음 연구원들은 Lightning Attention-2가 지원하는 TransNormerLLM(TNL-LA2)을 동일한 코퍼스에서 1B 및 3B에서 FlashAttention2가 지원하는 다른 고급 비-Transformer 아키텍처 네트워크 및 LLaMA와 비교했습니다. 아래 그림에서 볼 수 있듯이 TNL-LA2와 LLaMA는 비슷한 추세를 유지하며 손실 성능이 더 좋습니다. 이 실험은 Lightning Attention-2가 언어 모델링에 있어 최첨단 Transformer 아키텍처와 뒤지지 않는 정확도 성능을 가지고 있음을 보여줍니다.

대형 언어 모델 작업에서 연구원들은 비슷한 크기의 대형 모델에 대한 공통 벤치마크에서 TNL-LA2 15B와 Pythia의 결과를 비교했습니다. 아래 표에서 볼 수 있듯이, 동일한 토큰을 먹는 조건에서 TNL-LA2는 상식 추론에 대한 Softmax 주목과 객관식 종합 능력을 기반으로 하는 Pythia 모델보다 약간 더 높습니다.

Lightning Attention-2 속도 비교

연구원들은 Lightning Attention-2와 FlashAttention2의 단일 모듈 속도와 메모리 사용량을 비교했습니다. 아래 그림에서 볼 수 있듯이 Lightning Attention-1 및 FlashAttention2와 비교하여 Lightning Attention-2는 시퀀스 길이에 비해 속도가 엄격한 선형 증가를 보여줍니다. 메모리 사용량 측면에서는 세 가지 모두 유사한 추세를 보이지만 Lightning Attention-2는 메모리 사용량이 더 적습니다. 그 이유는 FlashAttention2와 Lightning Attention-1의 메모리 사용량도 거의 선형이기 때문입니다.

저자는 이 글의 주요 초점이 선형 주의 네트워크의 훈련 속도를 해결하고 모든 길이의 긴 시퀀스에 대해 1K 시퀀스와 유사한 훈련 속도를 달성하는 것임을 알아냈습니다. 추론 속도 측면에서는 소개가 많지 않습니다. 추론 과정에서 선형 어텐션(Linear Attention)이 손실 없이 RNN 모드로 변환되어 비슷한 효과를 얻을 수 있기 때문입니다. 즉, 단일 토큰을 추론하는 속도가 일정하기 때문입니다. Transformer의 경우 현재 토큰의 추론 속도는 이전 토큰 수와 관련이 있습니다.

저자는 Lightning Attention-1이 지원하는 TransNormerLLM-7B와 일반 7B 모델 간의 추론 속도 비교를 테스트했습니다. 아래 그림에서 볼 수 있듯이 대략적인 매개변수 크기에서 Lightning Attention-1의 처리 속도는 Baichuan의 4배, ChatGLM의 3.5배 이상으로 뛰어난 추론 속도 이점을 보여줍니다.

요약

Lightning Attention-2는 선형 주의 메커니즘의 주요 발전을 나타내며 정확성과 속도 측면에서 기존 Softmax 주의를 완벽하게 대체하여 더 많은 것을 위한 플랫폼을 제공합니다. 앞으로 더 많은 관심을 받을 것입니다. 대형 모델은 지속 가능한 확장 기능을 제공하고 무한히 긴 시퀀스를 더 효율적으로 처리할 수 있는 방법을 제공합니다. OpenNLPLab 팀은 현재의 메모리 장벽 문제를 해결하기 위해 앞으로 선형 주의 메커니즘을 기반으로 하는 순차 병렬 알고리즘을 연구할 예정입니다.

위 내용은 ​신세대 어텐션 메커니즘 Lightning Attention-2: 무제한 시퀀스 길이, 일정한 컴퓨팅 성능 오버헤드, 더 높은 모델링 정확도의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!