Lightning Attention-2: 무제한 시퀀스 길이, 일정한 컴퓨팅 성능 비용 및 더 높은 모델링 정확도를 달성하는 차세대 Attention 메커니즘

Lightning Attention-2는 긴 시퀀스의 훈련 및 추론 비용을 1K 시퀀스 길이와 일치시키는 새로운 선형 주의 메커니즘입니다.

대형 언어 모델의 시퀀스 길이 제한으로 인해 다중 회전 대화, 긴 텍스트 이해, 다중 모드 데이터 처리 및 생성 등과 같은 인공 지능 분야에서의 적용이 크게 제한되었습니다. 이러한 제한의 근본적인 이유는 현재 대규모 언어 모델에 사용되는 Transformer 아키텍처가 시퀀스 길이에 비해 2차 계산 복잡성을 갖기 때문입니다. 이는 시퀀스 길이가 증가함에 따라 필요한 컴퓨팅 리소스가 기하급수적으로 증가한다는 것을 의미합니다. 긴 시퀀스를 효율적으로 처리하는 방법은 항상 대규모 언어 모델의 과제 중 하나였습니다.

이전 방법은 추론 단계에서 대규모 언어 모델을 더 긴 시퀀스에 적응하는 방법에 초점을 맞추는 경향이 있었습니다. 예를 들어 Alibi 또는 유사한 상대 위치 코딩 방법을 사용하여 모델이 다양한 입력 시퀀스 길이에 적응할 수 있도록 하거나 RoPE와 같은 유사한 상대 위치 코딩 방법을 사용하여 차이를 수행할 수 있으며 모델에 대해 추가 교육을 수행할 수 있습니다. 훈련이 완료된 시퀀스 길이를 증폭시키는 목적을 달성하기 위해 단기 미세 조정이 수행됩니다. 이러한 방법을 사용하면 대규모 모델이 특정 긴 시퀀스 모델링 기능만 가질 수 있지만 실제 교육 및 추론 오버헤드는 줄어들지 않습니다.

OpenNLPLab 팀은 대규모 언어 모델의 긴 시퀀스 문제를 완전히 해결하려고 노력합니다. 그들은 긴 시퀀스의 훈련 및 추론 비용을 1K 시퀀스 길이와 일치시키는 새로운 선형 주의 메커니즘인 Lightning Attention-2를 제안하고 오픈 소스로 제공했습니다. 메모리 병목 현상이 발생하기 전에 시퀀스 길이를 무한정 늘려도 모델 학습 속도에 부정적인 영향을 미치지 않습니다. 이를 통해 무제한 길이의 사전 훈련이 가능해졌습니다. 동시에 매우 긴 텍스트의 추론 비용도 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년에 선형 주의가 출현한 이후[https://arxiv.org/abs/2006.16236], 연구자들은 이론적 선형 계산 복잡성과 일치하는 선형 주의의 실제 효율성을 만들기 위해 열심히 노력해 왔습니다. 2023년 이전에는 대부분의 작업이 Transformers와 정확도를 맞추는 데 초점을 맞춘 선형 주의에 대한 작업이었습니다. 마침내 2023년 중반에 개선된 선형 주의 메커니즘[https://arxiv.org/abs/2307.14995]이 정확성 측면에서 최첨단 Transformer 아키텍처와 정렬될 수 있습니다. 그러나 선형 주의에서 계산 복잡도를 선형으로 변경하는 가장 중요한 "왼쪽 곱셈에서 오른쪽 곱셈" 계산 트릭(아래 그림 참조)은 실제 왼쪽 직접 곱셈 알고리즘보다 훨씬 느립니다. 구현. 그 이유는 오른쪽 곱셈을 구현하려면 많은 수의 루프 연산을 포함하는 누적 합계(cumsum)를 사용해야 하기 때문입니다. IO 연산 수가 많기 때문에 왼쪽 곱셈에 비해 오른쪽 곱셈의 효율성이 훨씬 낮습니다.

^{그림 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의 핵심 아이디어 중 하나는 계산 비용이 많이 드는 소프트맥스 연산자를 제거하는 것입니다. 주의사항 계산식은 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) 매개변수 모델 -2에서 Lightning Attention을 비교했습니다. 번개 Attention-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배 이상으로 뛰어난 추론 속도 이점을 보여줍니다.

Summary

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

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