인공 지능

BlackMamba: 상태공간 모델을 위한 전문가의 혼합

게재

1 달 전

2024 년 3 월 26 일

디코더 전용 변환기 모델로 구축된 LLM(대형 언어 모델)의 개발은 자연어 처리(NLP) 영역을 변화시키고 다음을 포함한 다양한 딥 러닝 애플리케이션을 발전시키는 데 중요한 역할을 했습니다. 강화 학습, 시계열 분석, 이미지 처리 등이 있습니다. 그러나 확장성과 강력한 성능에도 불구하고 디코더 전용 변압기 모델로 구축된 LLM은 여전히 심각한 단점에 직면해 있습니다. 비록 표현적이지만 변환기 파생 LLM의 주의 메커니즘은 추론과 훈련 모두에서 높은 계산 리소스를 요구하므로 시퀀스 길이와 2차 FLOP에 대한 상당한 메모리가 필요합니다. 이러한 높은 계산 요구 사항은 변환기 모델의 컨텍스트 길이를 제한하여 자동 회귀 생성 작업의 규모에 비례하여 비용이 많이 들고 지속적인 데이터 스트림으로부터의 학습과 진정한 무제한 시퀀스 처리 기능을 방해합니다.

최근에, 상태공간 모델 (SSM)은 시퀀스 길이와 선형 시간의 함수로 메모리 복잡성을 달성하면서 대규모 모델링 벤치마크에서 변압기 아키텍처 모델과 경쟁하면서 놀라운 기능과 성능을 입증했습니다. 또한 최근 출시된 State Space Model인 Mamba는 다양한 언어 모델링 및 긴 시퀀스 처리 작업에서 뛰어난 성능을 보여주었습니다. 동시에, MoE(Mixture of Expert) 모델은 더 큰 메모리 공간을 희생하더라도 지연 시간과 추론 계산 비용을 크게 줄이면서 인상적인 성능을 보여주었습니다. Mamba 및 MoE 모델을 기반으로 하는 이 기사에서는 Mamba 상태 공간 모델과 MoE 모델을 결합하여 두 프레임워크가 제공하는 이점을 활용하는 새로운 아키텍처인 BlackMamba에 대해 설명합니다. BlackMamba에 대한 실험에서는 FLOP 교육 및 추론 모두에서 기존 Mamba 프레임워크 및 변환기 기준을 능가하는 성능을 입증했습니다. BlackMamba 프레임워크의 뛰어난 성능은 Mamba와 MoE 프레임워크의 기능을 효과적으로 결합하여 Mamba의 선형 복잡성 생성과 함께 MoE의 빠르고 비용 효율적인 추론을 제공할 수 있음을 보여줍니다.

이 문서의 목적은 BlackMamba 프레임워크를 심층적으로 다루는 것입니다. 우리는 프레임워크의 메커니즘, 방법론, 아키텍처를 탐색하고 최신 이미지 및 비디오 생성 프레임워크와 비교합니다. 시작하자.

BlackMamba: 상태 공간 모델을 위한 MoE 소개

LLM(대형 언어 모델), 특히 디코더 전용 변환기 아키텍처를 기반으로 하는 모델의 발전은 눈에 띄게 영향을 미쳤습니다. 자연 언어 처리 (NLP) 강화 학습, 시계열 분석, 이미지 처리 등을 포함한 다양한 딥 러닝 애플리케이션으로 확장되었습니다. 그럼에도 불구하고 확장성과 강력한 성능에도 불구하고 이러한 디코더 전용 변환기 기반 LLM은 주목할만한 문제에 직면합니다. Transformer 기반의 핵심 기능인 어텐션 메커니즘 LLM추론과 훈련 모두를 위해 광범위한 계산 리소스가 필요합니다. 여기에는 시퀀스 길이에 따라 증가하는 메모리와 2차적으로 증가하는 계산 연산(FLOP)이 필요합니다. 이러한 집약적인 계산 요구 사항은 모델의 컨텍스트 길이를 제한하고, 모델이 확장됨에 따라 자동 회귀 생성 작업 비용을 높이며, 지속적인 데이터 스트림에서 학습하거나 길이가 무제한인 시퀀스를 효율적으로 처리하는 모델의 능력을 방해합니다.

이러한 한계를 극복하기 위해 지난 몇 년 동안 상당한 노력이 이루어졌으며, 가장 유망한 후보 아키텍처인 SSM 및 MoE 모델을 사용하여 정식 밀도 주의 변압기 모델에 대한 아키텍처 대안을 고안하는 쪽으로 관심이 옮겨졌습니다. 변환기 아키텍처 모델보다 상태 공간 모델을 선호함으로써 얻을 수 있는 주요 이점은 변환기가 제공하는 2차 복잡성과 달리 SSM이 제공하는 입력 시퀀스 길이에 대한 선형 계산 복잡성입니다. 이론적으로 입력 시퀀스 길이와 관련된 선형 계산 복잡성을 통해 상태 공간 모델은 주어진 FLOPS 또는 초당 부동 소수점 작업 예산에 대해 변환기 아키텍처 모델보다 더 큰 시퀀스를 처리하고 KV 캐시 없이 계산에서 자동 회귀 생성 상수를 렌더링할 수 있습니다. Mamba, RetNet 및 기타 몇 가지를 포함하여 최근 개발된 상태 공간 모델은 유사한 스케일링 속성을 가진 변환기에 대한 경쟁력 있는 언어 모델링 작업 성능과 함께 효율적인 긴 시퀀스 추론 및 훈련을 보여주었습니다. 반면, 전문가 모델 혼합 아키텍처는 밀도가 높은 모델과 비슷한 품질을 달성하는 데 필수적인 추론 및 훈련 FLOP를 크게 줄일 수 있기 때문에 밀도가 높은 변환기의 대안으로 인기를 얻고 있습니다. MoE(Mixture of Experts) 모델은 단일 정방향 전달 동안 전체 매개변수의 희박한 선택만 활성화하여 작동합니다. 그들은 라우팅 기능을 활용하여 주어진 상황에 따라 어떤 '전문가'가 조치를 취해야 하는지 결정합니다. 이 접근 방식은 추론의 계산 비용과 총 매개변수 수를 분리하여 매개변수 수가 증가하고 메모리 요구 사항이 커지더라도 고정된 추론 예산 내에서 성능을 향상시킬 수 있습니다.

이러한 아키텍처의 발전은 기존 변압기에 비해 주목할만한 이점을 제공하며 향후 개발을 위한 흥미로운 방향을 제시합니다. 우리는 이러한 향상된 기능을 결합된 Mamba-MoE 모델에 통합하면 표준 변환기 모델보다 언어 모델링 기능과 효율성을 크게 가속화할 수 있다고 가정합니다. 기존의 고밀도 변압기 모델과 비교하여 Mamba-MoE 아키텍처의 예상되는 이점은 다음과 같습니다.

맘바: 훈련 및 추론 단계 모두에 대해 입력 시퀀스 길이에 비해 선형 계산 복잡성을 달성합니다. 이를 통해 일정한 시간 프레임과 일정한 메모리 사용으로 자동 회귀 생성이 발생할 수 있습니다.

환경부: 더 작고 조밀한 기준 모델에 필적하는 추론 속도와 훈련 계산 효율성을 제공하는 동시에 더 조밀한 버전과 동일한 수의 매개변수를 가진 모델에 필적하는 모델 품질 수준을 유지합니다.

즉, 변환기 아키텍처 모델이 여전히 최첨단이며 언어 모델링 작업 및 시퀀스 처리 작업에서 일관되고 놀라운 강력한 성능을 입증했다는 점을 언급하는 것이 중요합니다. 핵심적으로 변환기 아키텍처는 시퀀스에 포함된 서로 다른 토큰 간의 내적 유사성에 대한 2차 전체 대 전체 비교를 수행하고 출력 벡터에 대한 선형 맵을 수행하는 self-attention을 사용합니다. 변환기 모델은 주어진 활성화 함수를 가진 2층 MLP로 추가로 구성된 MLP 또는 다층 퍼셉트론 블록 사이에 쌓인 self-attention 블록으로 구성됩니다.

BlackMamba : 아키텍처 및 방법론

상태공간 모델

상태 공간 모델은 입력 시퀀스의 길이와 관련하여 선형 복잡도를 갖는 시퀀스 모델 그룹에 속합니다. 상태 공간 모델의 아키텍처는 주의 기반 아키텍처보다는 Recurrent Neural Networks 및 Convolutional Neural Networks에 더 잘 부합하며 암시적 잠재 공간을 통해 1차원 함수를 매핑하는 연속 동적 시스템에서 영감을 받았습니다. 선형 동적 시스템은 연관 스캔이나 컨볼루션 스캔을 사용하여 병렬 계산을 효율적으로 만듭니다. 실제 시나리오에서 상태 공간 모델의 반복적 특성은 GPU와 같은 고도로 병렬화된 AI 하드웨어에 여전히 채택되는 이유였습니다. 그러나 RWKV와 같은 SSM의 등장과 맘바 병렬 스캔 커널을 사용하여 반복 작업을 GPU에 효율적으로 매핑함으로써 변환기 모델에서 달성한 것과 비슷한 효율성으로 새로운 아키텍처의 교육을 촉진했습니다.

변환기 내 시퀀스 길이와 관련된 고유한 2차 복잡성은 매우 긴 컨텍스트에 대한 추론과 이해를 방해하는 잘 알려진 제한 사항입니다. 최근 혁신을 통해 컨텍스트 길이를 확장하는 아이디어가 도입되어 추론 중에 훨씬 더 긴 컨텍스트에 적용되기 전에 변환기를 실행 가능한 규모로 훈련할 수 있습니다. 이러한 발전에도 불구하고 추론 프로세스는 특히 KV(키-값) 캐시를 유지 관리하기 위해 여전히 상당한 양의 계산 리소스와 메모리를 요구하므로 리소스 집약적인 작업이 됩니다. 최근 연구 노력은 어텐션 메커니즘에서 발견되는 QKV(쿼리, 키, 값) 행렬과 유사한 입력 종속 게이팅 메커니즘을 통합하여 상태공간 모델의 표현 기능을 향상시키는 데 중점을 두었습니다.

이러한 노력은 상태공간 재귀의 본질적인 선형 진행을 보존하여 컨볼루션 또는 선택적 스캔 프로세스를 통해 효율적인 실행을 허용하는 것을 목표로 합니다. 이 접근 방식은 실제 응용 분야에서 변압기와의 성능 차이를 크게 줄입니다. 이러한 발전 중에서 Mamba는 이전 연구의 목표를 반영하는 상태 공간 모델로 두각을 나타내며 최대 2.8억 매개변수 규모에서 변압기에 필적하는 인상적인 성능 수준을 보여줍니다. 이는 상태 공간 모델(SSM) 재귀의 입력에 입력 종속 게이팅을 적용하는 동시에 맞춤형 선택적 스캔 커널을 사용하여 효율적인 계산을 보장함으로써 이를 달성합니다.

전문가 모델의 혼합

MoE(Mixture of Expert) 모델은 순방향 전달 중에 매개변수를 선택적으로 활성화하여 추론 비용과 총 매개변수 수를 분리합니다. 모든 매개변수를 사용하는 대신 이 모델은 토큰을 특정 MLP(Multilayer Perceptron) 전문가에게 전달합니다. 이상적으로 각 전문가는 본질적으로 소형 신경망인 라우팅 메커니즘을 통해 특정 유형의 입력을 처리하도록 맞춤화되어 각 토큰에 가장 적합한 전문가를 결정합니다. 이 접근 방식은 더 조밀한 구성에서 동일한 수의 매개 변수를 사용하여 모델의 포괄적인 표현력을 유지하면서 계산 요구 사항을 상당히 줄이는 것을 목표로 합니다. 일반적으로 라우터는 토큰에서 전문가 인덱스로 선형 레이어를 매핑하며, 각 전문가는 단순히 표준 변환기인 다층 퍼셉트론입니다. 그러나 전문가 할당 문제는 미분 불가능하고 Mixture of Expert 모델은 하드웨어 효율성을 위해 여러 전문가 간의 로드 밸런싱 및 교육 안정성에 어려움을 겪는 경우가 많기 때문에 개발자는 아직 라우터에 대한 최적의 교육 방법을 파악하지 못했습니다.

아키텍처

BlackMamba는 기본적으로 인터리빙된 MLP 블록과 잔여 스트림을 따라 순차적으로 추가되는 어텐션 블록으로 구성된 표준 변환기 모델을 사용합니다. 이제 대부분의 Mixture of Expert 모델은 다층 퍼셉트론 블록을 라우팅된 전문가 계층으로 대체합니다. 반면 BlackMamba 프레임워크는 변환기의 다층 퍼셉트론 블록을 라우팅된 전문가 레이어로 대체할 뿐만 아니라 Attention 레이어를 Mamba State Space Model 레이어로 대체합니다. BlackMamba 프레임워크의 아키텍처는 다음 그림에 나와 있습니다.

훈련 및 데이터세트

BlackMamba 모델은 사용자 정의 데이터 세트에서 300억 개 이상의 토큰에 대해 훈련되었으며 전문 다층 퍼셉트론에 SwiGLU 활성화 기능을 사용합니다. 프레임워크는 8명의 전문가로 교육하며, 개발자는 이 숫자가 모델의 메모리 공간과 추론 비용 사이에서 적절한 균형을 이루고 균형을 이루는 것으로 나타났습니다. BlackMamba 프레임워크를 훈련하는 데 사용되는 사용자 정의 데이터 세트는 Starcoder, SlimPajama, Pile 등을 포함한 기존 오픈 소스 데이터 세트의 혼합으로 구성됩니다. 다음 표는 BlackMamba 프레임워크 교육에 사용되는 각 데이터 세트의 가중치를 보여줍니다. 전체적으로 데이터세트에는 1.8조 XNUMX천억 개의 토큰이 있습니다.

BlackMamba : 결과

Mamba와 BlackMamba 간의 공정한 비교를 보장하기 위해 개발자는 동일한 훈련 데이터에 대해 동일한 훈련 매개변수를 사용하여 두 모델을 훈련했습니다. BlackMamba 프레임워크는 추론 시 동일한 순방향 전달 모델 크기와 초당 부동 소수점 연산 교육에 대해 Mamba 및 변환기 모델 모두보다 뛰어난 성능을 발휘할 수 있습니다. 다음 그림은 초기 단일 토큰 프롬프트에서 주어진 길이의 시퀀스를 자동 회귀적으로 생성하는 데 걸리는 시간을 시퀀스 길이의 함수로 보여줍니다.

또한 전문가 혼합 모델과 Mamba 모델의 지연 시간 이점이 BlackMamba 프레임워크에 결합되어 변환기 모델, 순수 Mamba 모델 및 MoE 모델과 비교할 때 추론 시간이 훨씬 빨라집니다. 또한 BlackMamba 프레임워크의 추론 이점은 시퀀스 길이에 정비례하므로 BlackMamba는 긴 시퀀스 생성에 매우 효과적입니다. 다음 그림은 각각 340억 640천만 개 및 XNUMX억 XNUMX천만 개의 매개 변수를 사용하여 BlackMamba 모델에 할당된 토큰 수를 보여줍니다. 보시다시피, 대부분의 레이어는 BlackMamba 모델에 의해 구현된 향상된 Sinkhorn 알고리즘의 결과로 높은 수준의 전문가 균형을 보여줍니다.

다음 표에는 다양한 오픈 소스 사전 훈련된 언어 모델과 비교한 BlackMamba 프레임워크의 평가 점수가 나와 있습니다. 관찰할 수 있듯이 BlackMamba 프레임워크는 모든 기준에 걸쳐 대부분의 프레임워크와 경쟁하고 성능을 능가할 수 있습니다. 또한 BlackMamba를 능가하는 모델의 매개변수 수가 훨씬 더 많고 성능 차이도 최소화된다는 점은 주목할 가치가 있습니다. 이는 더 적은 매개변수를 사용하여 BlackMamba 프레임워크의 능력을 나타냅니다.

최종 생각

이 기사에서는 Mamba 상태 공간 모델과 전문가 모델의 혼합을 결합하여 두 프레임워크가 제공하는 이점을 모두 활용하는 새로운 아키텍처인 BlackMamba에 대해 설명했습니다. BlackMamba에 대한 실험에서는 FLOP 교육 및 추론 모두에서 기존 Mamba 프레임워크 및 변환기 기준보다 뛰어난 성능을 발휘하는 것으로 나타났습니다. BlackMamba 프레임워크의 탁월한 성능은 MoE의 저렴하고 빠른 추론과 Mamba의 선형 복잡성 생성을 결합하므로 Mamba 및 MoE 프레임워크의 기능을 매우 잘 상속하고 결합할 수 있음을 보여줍니다. 우리는 BlackMamba 프레임워크의 아키텍처가 어떻게 강력하게 훈련된 대형 언어 모델, 기존 Mamba 프레임워크 및 전문가 혼합 모델을 FLOP 훈련 및 추론 비용 측면에서 능가할 수 있는지에 대해 이야기했습니다. 또한 BlackMamba 프레임워크는 Mixture of Expert 모델과 Mamba 프레임워크 모두에서 세대 FLOP와 감소된 교육을 동시에 상속합니다.