대체 맘바의 원리에 대한 심층적인 설명이 포함된 트랜스포머 리뷰(그림)

언어 모델링에서 트랜스포머의 대안

트랜스포머 아키텍처는 대규모 언어 모델(LLM)의 성공을 위한 핵심 요소입니다. 오늘날 사용되는 거의 모든 대규모 언어 모델은 다음과 같은 오픈 소스 모델에서 아키텍처를 사용합니다. 미스트랄 를 ChatGPT와 같은 비공개 소스 모델로 전환합니다.

대규모 언어 모델을 더욱 개선하기 위해 트랜스포머 아키텍처를 뛰어넘을 수 있는 새로운 아키텍처가 개발되었습니다. 이러한 접근 방식 중 하나는 Mamba한 가지 유형의 상태 공간 모델.

종이 속 맘바 맘바: 선택적 상태 공간을 사용한 선형 시간 시퀀스 모델링1 에서 발표되었습니다. 이에 대한 자세한 내용은 저장소 공식 구현 및 모델 체크포인트는 여기에서 확인하세요.

이 글에서는 언어 모델링의 맥락에서 상태 공간 모델링 분야를 소개하고 이 분야에 대한 이해를 돕기 위해 다양한 개념을 단계별로 살펴보겠습니다. 그런 다음 Mamba가 트랜스포머 아키텍처에 어떻게 도전할 수 있는지 논의할 것입니다.

이 글에서는 시각적 가이드로서 맘바 및 상태 공간 모델을 이해하는 데 도움이 되는 다양한 시각화를 살펴봅니다!

1부:트랜스포머 문제

맘바가 얼마나 흥미로운 아키텍처인지 설명하기 위해 먼저 트랜스포머를 간단히 살펴보고 그 단점 중 하나를 살펴보겠습니다.

트랜스포머는 모든 텍스트 입력을 다음과 같이 처리합니다. 토큰 구성 요소 시퀀스.

트랜스포머의 주요 장점 중 하나는 어떤 입력을 받더라도 시퀀스의 이전 토큰으로 돌아가 해당 표현을 도출할 수 있다는 것입니다.

트랜스포머의 핵심 구성 요소

트랜스포머는 텍스트를 표현하는 인코더 블록 세트와 텍스트를 생성하는 디코더 블록 세트의 두 가지 구조로 구성되어 있다는 점을 기억하세요. 이러한 구조를 결합하여 번역을 비롯한 여러 작업에 사용할 수 있습니다.

이 구조를 가지고 디코더만 사용하여 제너레이티브 모델을 만들 수 있습니다. 이 트랜스포머 기반 모델인제너레이티브 사전 교육 트랜스포머(GPT)는 디코더 블록을 사용하여 입력 텍스트의 일부를 완성합니다.

어떻게 작동하는지 살펴봅시다!

훈련의 축복...

단일 디코더 블록은 마스크된 자기 주의 메커니즘과 피드포워드 신경망이라는 두 가지 주요 구성 요소로 이루어져 있습니다.

자기 주의 메커니즘은 이 모델이 잘 작동하는 중요한 이유입니다. 이 메커니즘은 전체 시퀀스를 압축되지 않은 상태로 볼 수 있게 해주며 빠르게 훈련할 수 있습니다.

그렇다면 어떻게 작동할까요?

이 매트릭스는 각 토큰 각 이전 토큰과 비교합니다. 행렬의 가중치는 토큰 쌍의 관련성에 따라 달라집니다.

이 매트릭스는 교육 중에 한 번에 생성됩니다. "My"and"이름"주의도 사이는 계산하기 전에 계산할 필요가 없습니다."이름"and"는"주의 사이.

다음을 실현합니다.병렬화를 사용하면 훈련 속도가 크게 빨라집니다!

추론의 문제!

하지만 단점이 있습니다. 다음 토큰을 생성할 때 다음 토큰을 생성하려면전체 시퀀스토큰을 이미 생성한 경우에도 마찬가지입니다.

의 길이를 생성합니다.L이 시퀀스에는 약L²배로 증가하는데, 시퀀스의 길이가 길어지면 계산 비용이 많이 들 수 있습니다.

전체 시퀀스를 다시 계산해야 하는 이 문제는 트랜스포머 아키텍처의 주요 병목 현상 중 하나입니다.

"고전적인" 기술인 순환 신경망(RNN)이 이 느린 추론 문제를 어떻게 해결하는지 살펴보겠습니다.

RNN이 해결책일까요?

순환 신경망(RNN)은 시퀀스 기반 네트워크입니다. 각 시간 단계, 즉 시간 단계마다 두 개의 입력을 받습니다.t입력 및 이전 시간 단계t-1를 사용하여 다음 숨겨진 상태를 생성하고 출력을 예측합니다.

RNN에는 한 단계에서 다음 단계로 정보를 전달할 수 있는 루핑 메커니즘이 있습니다. 이 시각화 프로세스를 '확장'하여 보다 명확하게 만들 수 있습니다.

출력을 생성할 때 RNN은 이전 숨겨진 상태와 현재 입력만 고려하면 됩니다. 따라서 트랜스포머에 필요한 이전의 모든 숨겨진 상태를 다시 계산하는 문제를 피할 수 있습니다.

즉, RNN은 시퀀스 길이에 따라 선형적으로 확장되기 때문에 빠른 추론이 가능합니다! 이론적으로는컨텍스트 길이 무제한.

이를 설명하기 위해 앞서 사용한 입력 텍스트에 RNN을 적용해 보겠습니다.

각 숨겨진 상태는 이전의 모든 숨겨진 상태를 집계한 것으로, 일반적으로 압축된 보기입니다.

하지만 여기에 문제가 있습니다 ......

이름을 생성할 때 "Maarten"마지막 숨겨진 상태가 더 이상 단어에 대한 정보를 포함하지 않을 때 "안녕하세요"RNN은 마지막 상태만 고려하기 때문에 시간이 지나면 정보를 잊어버리는 경향이 있습니다.

RNN은 훈련과 추론 측면에서 더 빠르지만, Transformer 모델에 비해 정확도가 떨어집니다.

따라서 우리는 상태 공간 모델을 연구하여 RNN(때로는 컨볼루션)을 효율적으로 사용합니다.

2부:상태 공간 모델링(SSM)

트랜스포머나 RNN과 같은 상태 공간 모델(SSM)은 텍스트나 신호와 같은 정보 시퀀스를 처리합니다. 이 섹션에서는 SSM의 기본 개념과 텍스트 데이터와의 관계에 대해 소개합니다.

상태 공간이란 무엇인가요?

상태 공간은 시스템을 완전히 설명하는 데 필요한 최소한의 변수를 포함합니다. 상태 공간은 시스템의 가능한 상태를 정의하여 문제를 수학적으로 표현하는 방식입니다.

단순화해 봅시다. 우리가 미로를 여행하고 있다고 상상해 봅시다. "상태 공간"가능한 모든 위치(상태)를 표시한 지도입니다. 각 지점은 미로에서 고유한 위치를 나타내며 출구까지의 거리와 같은 구체적인 세부 정보가 표시됩니다.

"상태 공간 표현"는 이 지도에 대한 간단한 설명입니다. 현재 위치(현재 상태), 이동할 수 있는 다음 위치(가능한 미래 상태), 다음 위치로 이동하는 방법(오른쪽 또는 왼쪽 이동)을 표시합니다.

상태 공간 모델은 방정식과 행렬을 사용하여 이러한 동작을 추적하지만, 실제로는 현재 위치, 이동 가능한 위치, 이동 방법을 추적하는 방법일 뿐입니다.

변수는 상태(이 경우 X 및 Y 좌표)와 출구까지의 거리를 나타내며, 이는 "상태 벡터".

익숙한 단어인가요? 임베딩 또는 벡터는 언어 모델에서도 입력 시퀀스의 '상태'를 설명하는 데 자주 사용되기 때문입니다. 예를 들어 현재 위치를 설명하는 벡터(상태 벡터)는 다음과 같이 보일 수 있습니다:

신경망에서 '상태'는 일반적으로 시스템의 숨겨진 상태를 의미하며, 이는 대규모 언어 모델에서 새로운 토큰을 생성하는 데 있어 가장 중요한 측면 중 하나입니다.

상태 공간 모델이란 무엇인가요?

SSM(상태 공간 모델)은 이러한 상태 표현을 설명하고 특정 입력에 기반한 예측을 통해 다음 상태를 예측하는 데 사용되는 모델 클래스입니다.

전통적으로 시간이 지나면 tSSM:

• 입력 순서는 다음과 같습니다. x(t)(예: 미로에서 왼쪽 및 아래로 이동)을 잠재적 상태 표현에 매핑합니다. h(t)(예: 출구까지의 거리 및 x/y 좌표)
• 를 사용하여 예측된 출력 시퀀스를 도출합니다. y(t)(예: 출구로 더 빨리 가려면 다시 왼쪽으로 이동)

그러나 불연속 시퀀스(예: 왼쪽으로 한 번 이동)를 사용하는 대신 SSM은 연속 시퀀스를 입력으로 받아들이고 출력 시퀀스를 예측합니다.

SSM은 동적 시스템(예: 3D 공간에서 움직이는 오브젝트)이 시간적으로 t 의 상태와 이를 예측하는 두 가지 방정식입니다.

이러한 방정식을 풀면 관찰된 데이터(입력 시퀀스 및 이전 상태)를 기반으로 시스템의 상태를 예측하는 통계적 원리를 밝힐 수 있다는 가설을 세웁니다.

목표는 이 상태 표현을 찾는 것입니다. h(t)를 사용하여 입력에서 출력 시퀀스로 이동할 수 있습니다.

이 두 가지 방정식은 상태 공간 모델의 핵심입니다.

이 두 가지 방정식은 가이드 전반에 걸쳐 참조할 것입니다. 보다 직관적인 이해를 돕기 위해색상 코딩을 사용합니다.로 이동하여 빠르게 인용할 수 있습니다.

상태 방정식 는 입력에 따라 상태가 어떻게 변경되는지 설명합니다( 매트릭스 B) 충격 상태( 매트릭스 A)를 클릭하고 변경합니다.

앞서 언급했듯이h(t) 는 주어진 시간을 나타냅니다. t 잠재적 상태 표현은x(t) 는 특정 입력을 나타냅니다.

출력 방정식 는 상태 전달 방식을 설명합니다. 매트릭스 C 를 출력으로 변환하고, 입력이 출력으로 변환되는 방식은 매트릭스 D 출력에 영향을 줍니다.

다음 사항에 유의하십시오.: 매트릭스 A및B및C 노래로 응답 D 흔히 매개변수학습이 가능하기 때문입니다.

이 두 가지 방정식을 시각화하면 다음과 같은 아키텍처를 얻을 수 있습니다:

이러한 행렬이 학습 과정에 어떤 영향을 미치는지 단계별로 살펴보겠습니다.

몇 가지 입력 신호가 있다고 가정해 보겠습니다. x(t)신호는 먼저 매트릭스 B 곱하기.매트릭스 B 입력이 시스템에 어떤 영향을 미치는지 설명합니다.

업데이트된 상태(신경망의 숨겨진 상태와 유사)는 환경에 대한 핵심 '지식'을 담고 있는 잠재적 공간입니다. 이 상태를 매트릭스 A 이 행렬을 곱하면 모든 내부 상태 간의 연관성을 설명하며 시스템의 기본 역학을 나타냅니다.

눈치채셨겠지만매트릭스 A 상태 표현이 생성되기 전에 적용하고 상태 표현이 업데이트된 후에 다시 적용합니다.

그런 다음 매트릭스 C 를 사용하여 상태가 출력으로 변환되는 방식을 설명합니다.

마지막으로 매트릭스 D 입력에서 출력으로 직접 신호를 제공합니다. 이는 일반적으로 점프 연결.

때문에 매트릭스 D 점프 연결과 유사하게 SSM은 점프 연결을 포함하지 않는 다음과 같은 형태로 간주되는 경우가 많습니다.

단순화된 관점으로 돌아가서 이제 매트릭스에 집중할 수 있습니다. A및B 노래로 응답 C 를 SSM의 핵심으로 삼고 있습니다.

이전과 마찬가지로 원래 방정식을 업데이트하고 멋진 색상을 추가하여 각 행렬의 용도를 표시할 수 있습니다.

이 두 가지 방정식은 함께 관측된 데이터로부터 시스템의 상태를 예측하는 것을 목표로 합니다. 입력이 연속적이므로 SSM의 주요 표현은 다음과 같습니다. 연속 시간 척도.

연속 신호에서 불연속 신호로

연속 신호가 있는 경우 상태 표현을 찾습니다. h(t) 구문 분석이 어렵습니다. 또한 일반적으로 불연속 입력(예: 텍스트 시퀀스)이 있기 때문에 모델을 불연속화하려고 합니다.

이를 위해 제로 오더 유지 기술. 먼저 불연속 신호를 수신할 때마다 새로운 불연속 신호를 수신할 때까지 해당 값을 유지합니다. 이 과정을 통해 SSM에 대한 연속 신호가 생성됩니다:

학습 가능한 새로운 매개변수로 시간 값을 유지합니다. 심박 조율기 ∆ 표시합니다. 입력의 해상도를 나타냅니다.

이제 입력에 대한 연속 신호를 생성했으므로 이제 연속 출력을 생성하고 입력의 시간 단계에 따라 이러한 값을 샘플링할 수 있습니다.

이 샘플링된 값은 이산화의 결과입니다!

수학적으로 영차 보류를 적용하면 다음과 같이 계산할 수 있습니다:

이 두 가지를 함께 사용하면 연속형 SSM에서 개별형 SSM으로 전환할 수 있습니다. 이 시점에서 모델은 더 이상 기능 대 기능 x(t) → y(t)대신 시퀀스 간 xₖ → y_ₖ::

여기서 매트릭스 A 노래로 응답 B 이제 모델의 이산화된 매개변수를 나타냅니다.

당사는 다음을 사용합니다. k 대신 t를 사용하여 연속형 SSM과 불연속형 SSM을 구분할 수 있습니다.

주의: 교육 중에는 여전히 매트릭스 A 를 이산화된 버전이 아닌 연속적인 형태로 사용합니다. 훈련 중에는 연속형 표현이 이산화됩니다.

이제 이산 표현에 대한 공식을 얻었으므로 실제로 어떻게 카운트 모델.

재귀적 표현

이산화된 SSM을 사용하면 연속적인 신호가 아닌 특정 시간 단계로 문제를 구성할 수 있습니다. 앞서 RNN에서 살펴본 것처럼 재귀적 방법은 여기서 매우 유용합니다.

연속 신호 대신 불연속 시간 단계를 고려하면 시간 단계를 사용하여 문제를 다시 공식화할 수 있습니다:

각 시간 단계에서 현재 입력(Bxₖ) 이전 상태에 어떤 영향을 미치는지(Ahₖ₋₁)를 클릭한 다음 예상 출력(Chₖ).

이 표현은 이미 여러분에게 익숙할 수 있습니다! 앞서 RNN을 처리할 때와 마찬가지로 분석할 수 있습니다.

다음과 같이 확장(또는 일련의 시간 단계로 확장)할 수 있습니다:

RNN의 기본 접근 방식을 사용하여 이산화된 버전을 사용할 수 있다는 점에 유의하세요.

컨볼루션 표현(수학)

기존 이미지 인식 작업에서는 필터(커널)를 사용하여 집계된 특징을 추출합니다:

이미지가 아닌 텍스트를 다루기 때문에 일차원적인 보기를 사용해야 합니다:

이 '필터'를 표현하는 데 사용하는 커널은 SSM 공식에서 파생된 것입니다:

이 커널이 실제로 어떤 기능을 하는지 살펴보겠습니다. 컨볼루션과 마찬가지로 SSM 커널을 사용하여 각 토큰 집합을 순회하고 출력을 계산할 수 있습니다:

이것은 또한 패딩이 출력에 미칠 수 있는 영향을 보여줍니다. 시각화를 개선하기 위해 패딩의 순서를 변경했지만 일반적으로 문장의 마지막에 패딩을 적용합니다.

다음 단계에서는 커널이 한 번 이동하여 다음 단계의 계산을 수행합니다:

마지막 단계에서는 커널의 전체 효과를 확인할 수 있습니다:

SSM을 컨볼루션으로 표현할 때의 주요 이점 중 하나는 컨볼루션 신경망(CNN)처럼 병렬로 훈련할 수 있다는 점입니다. 그러나 커널 크기가 고정되어 있기 때문에 추론 속도가 RNN만큼 빠르지 않고 제약이 없습니다.

세 가지 유형의 표현

이 세 가지 표현은 다음과 같습니다.진행 상황및재귀(계산) 노래로 응답 컨볼루션 각기 다른 장단점이 있습니다:

흥미로운 점은 이제 효율적인 추론을 위해 재귀적 SSM을 사용하는 동시에 병렬 학습을 위해 컨볼루션 SSM을 사용할 수 있다는 점입니다.

이러한 표현을 사용하면 작업에 따라 표현을 선택하는 영리한 트릭을 사용할 수 있습니다. 훈련 중에는 병렬화가 가능한 컨볼루션 표현을 사용하고, 추론 중에는 효율적인 재귀 표현을 사용합니다:

이 모델은 선형 상태 공간 레이어(LSSL). 2

이러한 표현은 다음과 같은 중요한 속성을 공유합니다.선형 시간 불변(수학)（LTI는 SSM의 매개변수를 나타냅니다. A및B 노래로 응답 C 는 모든 시간 단계에서 고정되어 있습니다. 즉, 매트릭스 A및B 노래로 응답 C 는 생성된 모든 토큰에 대해 동일합니다.

즉, 어떤 시퀀스를 SSM에 제공하든 상관없습니다.A및B 노래로 응답 C 모든 값의 값은 동일하게 유지됩니다. 콘텐츠를 신경 쓰지 않는 정적 표현이 있습니다.

맘바가 이 문제를 어떻게 해결하는지 살펴보기 전에 퍼즐의 마지막 조각인매트릭스 A.

행렬 A 중요성

SSM 공식의 가장 중요한 측면 중 하나는 바로 매트릭스 A. 앞서 재귀적 표현에서 살펴본 것처럼, 재귀적 표현은 다음과 같은 정보를 캡처합니다. 이전 항목 상태에 대한 정보로 갱신 상태.

본질적으로.매트릭스 A 숨겨진 상태를 생성합니다:

따라서 매트릭스 A 는 이전 토큰 몇 개만 기억할 수 있는지 아니면 지금까지 본 모든 토큰을 캡처할 수 있는지를 결정할 수 있습니다. 특히 재귀적 표현의 맥락에서 보면 는 다음과 같이 회상합니다. 이전 상태.

많은 메모리를 유지하는 방식으로 작성하는 방법(컨텍스트 크기) 매트릭스 A?

헝그리 헝그리 히포! 또는 HiPPO3 실현하기 위해귀하의(경칭)계단 방법(어느 정도)학기(수학) 의 손에 자신을 맡기십시오.영화 행운operator.HiPPO는 지금까지 본 모든 입력 신호를 계수 벡터로 압축하려고 시도합니다.

다음을 사용합니다. 매트릭스 A 를 사용하여 최근 토큰을 더 잘 포착하고 오래된 토큰을 소멸시키는 상태 표현을 구성할 수 있으며, 공식은 다음과 같이 표현할 수 있습니다:

정사각형이 있다고 가정해 보겠습니다. 매트릭스 A이것은 우리에게 줍니다:

HiPPO로 구축하기 매트릭스 A 가 무작위 행렬로 초기화하는 것보다 훨씬 낫다는 것이 입증되었습니다. 따라서 이 방식은 업데이트 시그널링(가장 가까운 토큰) 측면이 더 중요합니다. 이전 신호(초기 토큰)가 더 정확합니다.

HiPPO 매트릭스의 핵심 아이디어는 히스토리를 기억하는 숨겨진 상태를 생성하는 것입니다.

수학적으로는 다음을 추적하여 이를 수행합니다. 레전드 다항식 계수를 사용하여 모든 과거 기록에 근사치를 구할 수 있습니다.4

그런 다음 앞서 살펴본 재귀 및 컨볼루션 표현에 HiPPO를 적용하여 장거리 종속성을 처리합니다. 결과는 다음과 같습니다. 시퀀스를 위한 구조화된 상태 공간(S4)SSM은 긴 시퀀스를 효율적으로 처리할 수 있는 SSM 클래스입니다.5

세 부분으로 구성되어 있습니다:

• 상태 공간 모델
• 처리용 HiPPO 장거리 의존성
• 이산화는 다음을 생성하는 데 사용됩니다. 재귀(계산) 노래로 응답 컨볼루션 표시

이러한 유형의 SSM은 선택한 표현 방식(재귀와 컨볼루션)에 따라 몇 가지 장점이 있습니다. 또한 긴 텍스트 시퀀스를 처리할 수 있고 HiPPO 행렬을 기반으로 하여 메모리를 효율적으로 저장할 수 있습니다.

다음 사항에 유의하십시오.HiPPO 행렬을 계산하고 자신만의 S4 모델을 구축하는 방법에 대해 더 깊이 이해하고 싶다면 다음을 읽어보시길 적극 추천합니다. 주석이 달린 S4.

3부:맘바 - 선택적 상태 공간 모델

드디어 맘바의 특별한 점을 이해하는 데 필요한 모든 기본 사항을 다루었습니다. 상태 공간 모델을 사용하여 텍스트 시퀀스를 모델링할 수 있지만, 여전히 피하고 싶은 단점이 많이 있습니다.

이 섹션에서는 맘바의 두 가지 주요 기능에 대해 설명합니다:

1. 한 종류의 선택적 스캔 알고리즘를 사용하여 모델이 (관련 없는) 정보를 선별할 수 있도록 합니다.
2. 한 종류의 하드웨어 인식 알고리즘by 병렬 스캔및커널 퓨전 노래로 응답 다시 계산 를 사용하여 (중간) 결과를 효율적으로 저장할 수 있습니다.

이 두 가지를 함께 사용하면 선택적 상태 공간 모델링 어쩌면 S6 자기 주의처럼 사용할 수 있는 모델입니다. 맘바 블록.

이 두 가지 주요 기여 사항을 살펴보기 전에 먼저 이러한 기능이 필요한 이유를 살펴보겠습니다.

해결하려는 문제는 무엇인가요?

상태 공간 모델, 심지어 S4(구조화된 상태 공간 모델)도 언어 모델링 및 생성의 특정 주요 작업에서 다음과 같은 성능이 저하됩니다. 특정 입력에 집중하거나 무시할 수 있는 기능.

다음과 같은 두 가지 합성 작업을 통해 이를 설명할 수 있습니다. 선택적 복사 노래로 응답 인덕터 헤드.

존재 선택적 복사 작업에서 SSM의 목표는 입력된 부분을 복사하여 순서대로 출력하는 것입니다:

그러나 SSM은 선형 시간 불변(명목식 형태로 사용됨)이 작업을 제대로 수행하지 못한다는 것을 알 수 있습니다. 앞서 살펴본 것처럼 매트릭스 A및B 노래로 응답 C 생성된 각 토큰은 SSM에 대해 동일합니다.

결과적으로 SSM은 다음을 실행할 수 없습니다. 콘텐츠 인식 추론이는 SSM이 입력(힌트)에 대해 추론하기를 원하기 때문에 발생하는 문제입니다. A, B, C의 고정 행렬로 인해 각 토큰을 동등하게 취급하기 때문에 발생하는 문제입니다.

SSM은 또 다른 작업, 즉 인덕터 헤드는 입력에서 발견된 패턴을 재현하는 것을 목표로 합니다:

위의 예에서는 기본적으로 모델에 각 "Q."를 제공한 후 "A." 응답을 반환합니다. 그러나 SSM은 시간에 변하지 않기 때문에 기록에서 어떤 이전 토큰을 불러올지 선택할 수 없습니다.

이를 위해 다음 사항에 집중해 보겠습니다. 매트릭스 B 를 예로 들어 이 점을 설명합니다. 입력에 관계없이 x 내용입니다.매트릭스 B 는 항상 동일하므로 x 관련 없음:

마찬가지로.A 노래로 응답 C 도 입력과 무관하게 항상 고정된 상태로 유지됩니다. 이는 지금까지 살펴본 SSM의 작동하지 않는 상태 특성.

반면, 트랜스포머는 입력 순서를 기반으로 하기 때문에 이러한 작업이 비교적 간단합니다. 동적(과학) 주의 전환. 시퀀스의 다른 부분을 선택적으로 '보기' 또는 '집중'할 수 있습니다.

이러한 작업에서 SSM의 성능 저하는 시간 불변 SSM의 잠재적인 문제점을 보여줍니다. A및B 노래로 응답 C 정적 특성 콘텐츠 센스 문제를 해결했습니다.

정보의 선택적 보존

SSM의 재귀적 표현은 전체 히스토리를 압축하기 때문에 매우 효율적인 작은 상태를 생성합니다. 그러나 트랜스포머 모델에 비해 트랜스포머 모델은 주의 매트릭스를 통해 히스토리를 압축하지 않으므로 성능이 더 뛰어납니다.

맘바는 두 가지 장점을 모두 갖추는 것을 목표로 합니다. 트랜스포머 국가만큼 강력한 작은 국가:

위에서 언급했듯이, 데이터를 선택적으로 압축하는 방식으로 이를 수행합니다. 입력 문장이 있을 때 일반적으로 의미가 없는 정지 단어와 같은 일부 정보가 있습니다.

정보를 선택적으로 압축하기 위해서는 입력에 따라 매개변수가 달라져야 합니다. 이를 위해 먼저 훈련 중 SSM의 입력과 출력의 차원을 살펴보겠습니다:

구조화된 상태 공간 모델(S4)에서 행렬 A및B 노래로 응답 C 의 치수는 입력과 독립적이기 때문입니다. N 노래로 응답 D 는 정적이며 변경되지 않습니다.

대신 Mamba는 입력 시퀀스 길이와 배치 크기를 결합하여 행렬을 만듭니다. B 노래로 응답 C심지어 심박 조율기 ∆입력에 따라 달라집니다:

즉, 각 입력 토큰에 대해 이제 각기 다른 B 노래로 응답 C 매트릭스를 통해 콘텐츠 인식 문제를 해결합니다!

다음 사항에 유의하십시오.: 매트릭스 A 는 상태 자체가 정적으로 유지되기를 원하기 때문에 변경되지 않지만 영향을 받는 방식은 ( B 노래로 응답 C)는 동적입니다.

함께 선택적으로 숨겨진 상태로 유지할 항목과 입력에 따라 달라지므로 무시할 항목을 선택합니다.

작은 심박 조율기 ∆ 는 이전 문맥을 위해 특정 단어를 무시하도록 유도하고, 큰 심박 조율기 ∆ 대신 문맥보다는 입력 어휘에 더 중점을 둡니다:

스캔 작업

이러한 행렬은 이제 동적(과학) 를 가정하는 컨볼루션 표현으로는 계산할 수 없습니다. 고정 를 사용할 수 없습니다. 재귀적 표현만 사용할 수 있으므로 컨볼루션이 제공하는 병렬화 이점을 잃게 됩니다.

병렬화를 달성하기 위해 재귀를 사용하여 출력을 계산하는 방법을 살펴보겠습니다:

각 상태는 이전 상태( A)에 현재 입력값을 곱한 값( B)의 및. 이를 스캔 작업이는 for 루프를 사용하여 쉽게 계산할 수 있습니다.

반대로 병렬화는 이전 상태를 사용할 수 있어야만 각 상태를 계산할 수 있기 때문에 불가능해 보입니다. 하지만 맘바는 다음과 같은 방법으로 각 상태의 병렬화를 가능하게 만들었습니다. 병렬 스캔 알고리즘이 이를 가능하게 합니다.

결합의 법칙의 속성을 활용하여 연산을 수행하는 순서는 중요하지 않다고 가정합니다. 따라서 시퀀스를 부분적으로 계산한 다음 반복적으로 결합할 수 있습니다:

동적 매트릭스 B 노래로 응답 C 및 병렬 스캐닝 알고리즘을 결합하여 선택적 스캔 알고리즘를 사용하여 재귀적 표현의 역동적이고 빠른 특성을 표현합니다.

하드웨어 인식 알고리즘

최신 GPU의 단점 중 하나는 작지만 효율이 높은 SRAM과 크지만 효율이 약간 떨어지는 DRAM 간의 전송(IO) 속도가 제한된다는 점입니다. SRAM과 DRAM 간의 잦은 정보 복사는 병목 현상이 될 수 있습니다.

맘바는 DRAM에서 SRAM으로 이동하는 횟수를 제한한다는 점에서 플래시 어텐션과 유사합니다. 다음과 같은 방법으로 이를 수행합니다. 커널 퓨전 이를 구현하면 모델에서 중간 결과 작성을 방지하고 계산이 완료될 때까지 계속 계산을 수행할 수 있습니다.

Mamba의 기본 아키텍처를 시각화하여 DRAM 및 SRAM 할당의 구체적인 예를 확인할 수 있습니다:

여기에서는 다음이 하나의 커널에 통합되어 있습니다:

• 재량화 단계는 다음과 동일합니다. 단계 크기 ∆
• 선택적 스캔 알고리즘
• 와 함께 C 곱셈

하드웨어 인식 알고리즘의 마지막 부분은 다음과 같습니다. 재계산.

중간 상태는 저장되지 않지만 리버스 패스에서 그라데이션을 계산하는 데 필요합니다. 대신 작성자는 리버스 패스 중에 이러한 중간 상태를 다시 계산합니다.

이는 비효율적으로 보일 수 있지만 상대적으로 느린 DRAM에서 이 모든 중간 상태를 읽는 것보다 훨씬 저렴합니다.

이제 아래 기사에서 아키텍처의 모든 구성 요소를 다루었습니다:

선택적 SSM. 검색 출처: 구, 알버트 및 트라이 다오. "맘바: 선택적 상태 공간을 사용한 선형 시계열 모델링." arXiv 사전 인쇄물 arXiv:2312.00752 (2023).

이 아키텍처는 종종 선택적 SSM 어쩌면 S6 모델은 기본적으로 선택적 스캐닝 알고리즘을 사용하여 계산된 S4 모델이기 때문입니다.

맘바 모듈

지금까지 살펴본 내용 선택적 SSM 는 디코더 모듈에서 자체 주의를 표현할 수 있는 것처럼 모듈로 구현할 수 있습니다.

디코더처럼 여러 개의 맘바 모듈을 쌓아서 그 출력을 다음 맘바 모듈의 입력으로 사용할 수 있습니다:

입력 임베딩을 확장하기 위해 선형 투영으로 시작합니다. 그런 다음 선택적 SSM 독립적인 토큰 계산을 방지하기 위해 이전에 컨볼루션을 적용했습니다.

선택적 SSM 에는 다음과 같은 특징이 있습니다:

• 통과(청구서 또는 검사 등) 불연속 생성 재귀적 SSM
• 매트릭스 A 계속 진행 HiPPO 캡처를 위한 초기화 장거리 종속성
• 선택적 스캔 알고리즘 를 클릭해 정보를 선택적으로 압축합니다.
• 하드웨어 인식 알고리즘 계산 속도를 높이려면

코드 구현을 살펴보면 이 아키텍처를 더욱 확장하여 엔드투엔드 예제가 어떤 모습일지 살펴볼 수 있습니다:

출력 토큰을 선택하기 위한 정규화 레이어와 소프트맥스 추가 등 몇 가지 변경 사항에 유의하세요.

이 모든 것을 종합하면 무한한 컨텍스트에서도 빠른 추론과 학습이 가능합니다. 저자들은 이 아키텍처를 사용하여 같은 크기의 트랜스포머 모델과 성능이 비슷하고 때로는 그 성능을 능가한다는 사실을 발견했습니다!

평결에 도달하기

이것으로 상태 공간 모델과 선택적 상태 공간 모델을 사용한 놀라운 맘바 아키텍처에 대한 탐구를 마칩니다. 이 포스팅을 통해 상태 공간 모델과 특히 맘바에 대해 더 잘 이해하셨기를 바랍니다. 트랜스포머를 대체할 수 있을지는 모르겠지만, 현재로서는 이렇게 다른 아키텍처가 주목받는 것을 보는 것만으로도 놀랍습니다!

빅 언어 모델링과 관련된 더 많은 비주얼리제이션을 보고 이 뉴스레터를 지원하려면 제가 Jay Alammar와 공동 저술한 책을 확인하세요.