소식

트랜스포머의 최강 경쟁자인 맘바를 이해하기 위한 한 편의 글

2024-08-19

한어Русский языкEnglishFrançaisIndonesianSanskrit日本語DeutschPortuguêsΕλληνικάespañolItalianoSuomalainenLatina



기계 심장 보고서

편집자: 팬더

Mamba는 훌륭하지만 아직 개발 단계입니다.

딥러닝 아키텍처는 많지만 최근 몇 년 동안 가장 성공적인 아키텍처는 여러 응용 분야에서 우위를 확고히 한 Transformer입니다.

이러한 성공의 핵심 동인은 주의 메커니즘입니다. 이를 통해 Transformer 기반 모델은 입력 시퀀스의 관련 부분에 집중하여 더 나은 상황별 이해를 얻을 수 있습니다. 그러나 Attention 메커니즘의 단점은 계산 오버헤드가 높고 입력 크기에 따라 2차적으로 증가하여 매우 긴 텍스트를 처리하기 어렵다는 것입니다.

다행스럽게도 얼마 전에 구조화된 상태 공간 시퀀스 모델(SSM)이라는 큰 잠재력을 지닌 새로운 아키텍처가 탄생했습니다. 이 아키텍처는 시퀀스 데이터의 복잡한 종속성을 효율적으로 캡처할 수 있으므로 Transformer의 강력한 상대가 됩니다.

이러한 유형의 모델 디자인은 고전적인 상태 공간 모델에서 영감을 얻었습니다. 우리는 이를 순환 신경망과 컨볼루션 신경망의 융합 모델로 생각할 수 있습니다. 루프 또는 컨볼루션 작업을 사용하여 효율적으로 계산할 수 있으므로 계산 오버헤드를 시퀀스 길이에 따라 선형 또는 거의 선형으로 확장할 수 있으므로 계산 비용이 크게 절감됩니다.

보다 구체적으로 SSM의 가장 성공적인 변형 중 하나인 Mamba는 Transformer와 유사한 모델링 기능을 제공하면서 시퀀스 길이에 따라 선형 확장성을 유지합니다.

Mamba는 먼저 입력을 기반으로 SSM을 다시 매개변수화하는 간단하면서도 효과적인 선택 메커니즘을 도입하여 모델이 관련 없는 정보를 필터링하는 동시에 필요하고 관련 있는 데이터를 무기한 유지할 수 있도록 합니다. 그런 다음 Mamba에는 컨볼루션 대신 스캔을 사용하여 모델을 반복적으로 계산하는 하드웨어 인식 알고리즘도 포함되어 있어 A100 GPU에서 계산 속도를 3배까지 높일 수 있습니다.

그림 1에서 볼 수 있듯이, 복잡한 긴 시퀀스 데이터를 모델링하는 강력한 능력과 선형에 가까운 확장성을 갖춘 Mamba는 기본 모델로 등장했으며 컴퓨터 비전, 자연어 처리, 의료 등 다양한 연구 개발 분야에 혁명을 일으킬 것으로 예상됩니다. 관리.



따라서 Mamba의 연구와 응용에 관한 문헌은 급속히 늘어나고 있으며 현기증이 나고 있으며 종합적인 검토 보고서는 큰 도움이 될 것입니다. 최근 홍콩 폴리테크닉 대학교(Hong Kong Polytechnic University) 연구팀이 arXiv에 기고한 내용을 발표했습니다.



  • 논문 제목: A Survey of Mamba
  • 논문 주소: https://arxiv.org/pdf/2408.01129

이 검토 보고서는 Mamba를 다양한 관점에서 요약하여 초보자가 Mamba의 기본 작동 메커니즘을 배우는 데 도움이 될 뿐만 아니라 숙련된 실무자가 최신 진행 상황을 이해하는 데도 도움이 됩니다.

Mamba는 인기 있는 연구 방향이므로 많은 팀에서 리뷰 보고서를 작성하려고 합니다. 이 기사에서 소개한 것 외에도 상태 공간 모델이나 시각적 Mamba에 초점을 맞춘 다른 리뷰가 있으므로 자세한 내용은 해당 논문을 참조하세요. :

  • Mamba-360: 장 시퀀스 모델링을 위한 변환기 대안으로서의 상태 공간 모델 조사: 방법, 응용 및 과제. arXiv:2404.16112
  • 변압기에 대한 신세대 네트워크 대안을 위한 상태 공간 모델: 조사. arXiv:2404.09516
  • Vision Mamba: 포괄적 조사 및 분류법. arXiv:2405.04404
  • 비전 맘바에 대한 조사: 모델, 응용 및 과제. arXiv:2404.18861
  • 비주얼 맘바에 대한 조사. arXiv:2404.15956

예비 지식

Mamba는 RNN(Recurrent Neural Network)의 순환 프레임워크, Transformer의 병렬 컴퓨팅 및 주의 메커니즘, SSM(State Space Model)의 선형 특성을 결합합니다. 그러므로 Mamba를 완전히 이해하기 위해서는 먼저 이 세 가지 아키텍처를 이해해야 합니다.

순환 신경망

순환 신경망(RNN)은 내부 메모리를 유지하는 능력으로 인해 시퀀스 데이터를 처리하는 데 능숙합니다.

구체적으로, 각 이산 시간 단계 k에서 표준 RNN은 이전 시간 단계의 은닉 상태와 함께 벡터를 처리한 다음 다른 벡터를 출력하고 은닉 상태를 업데이트합니다. 이 Hidden State는 과거에 보았던 입력 정보를 유지할 수 있는 RNN의 메모리로 사용될 수 있습니다. 이 동적 메모리를 통해 RNN은 다양한 길이의 시퀀스를 처리할 수 있습니다.

즉, RNN은 Hidden State에 저장된 과거 지식을 활용하여 시간적 패턴을 효과적으로 포착하는 비선형 순환 모델입니다.

변신 로봇

Transformer의 셀프 어텐션 메커니즘은 입력 간의 전역 종속성을 캡처하는 데 도움이 됩니다. 이는 다른 위치에 비해 중요성을 기준으로 각 위치에 가중치를 할당함으로써 수행됩니다. 보다 구체적으로, 원래 입력은 먼저 입력 벡터의 시퀀스 x를 쿼리 Q, 키 K 및 값 V의 세 가지 유형의 벡터로 변환하기 위해 선형 변환됩니다.

그런 다음 정규화된 주의 점수 S가 계산되고 주의 가중치가 계산됩니다.

단일 Attention 기능을 수행할 수 있는 것 외에도 다중 헤드 Attention을 수행할 수도 있습니다. 이를 통해 모델은 다양한 유형의 관계를 포착하고 다양한 관점에서 입력 시퀀스를 이해할 수 있습니다. Multi-head attention은 여러 세트의 self-attention 모듈을 사용하여 입력 시퀀스를 병렬로 처리합니다. 이러한 각 헤드는 독립적으로 작동하며 표준 self-attention 메커니즘과 동일한 계산을 수행합니다.

이후 각 헤드의 Attention 가중치를 집계하고 결합하여 값 벡터의 가중치 합을 얻습니다. 이 집계 단계를 통해 모델은 여러 헤드의 정보를 사용하고 입력 시퀀스에서 다양한 패턴과 관계를 캡처할 수 있습니다.

상태공간

상태 공간 모델(SSM)은 시간 경과에 따른 시스템의 동적 동작을 설명하는 데 사용할 수 있는 전통적인 수학적 프레임워크입니다. 최근 몇 년 동안 SSM은 사이버네틱스, 로봇공학, 경제학 등 다양한 분야에서 널리 사용되고 있습니다.

기본적으로 SSM은 "상태"라는 숨겨진 변수 세트를 통해 시스템 동작을 반영하여 시간 데이터의 종속성을 효과적으로 캡처할 수 있습니다. RNN과 달리 SSM은 연관 속성을 갖는 선형 모델입니다. 구체적으로, 고전적인 상태 공간 모델은 N차원 은닉 상태 h(t)를 통해 현재 시간 t에서의 입력 x와 출력 y 사이의 관계를 모델링하기 위해 두 가지 핵심 방정식(상태 방정식과 관찰 방정식)을 구성합니다.

  • 이산화

기계 학습의 요구 사항을 충족하기 위해 SSM은 연속 매개 변수를 이산 매개 변수로 변환하는 이산화 프로세스를 거쳐야 합니다. 일반적으로 이산화 방법의 목표는 연속 시간을 가능한 한 동일한 적분 면적을 갖는 K개의 이산 간격으로 나누는 것입니다. 이 목표를 달성하기 위해 SSM에서 채택한 가장 대표적인 솔루션 중 하나는 Δ = [_{−1}, _ ] 구간의 함수 값이 일정하게 유지된다고 가정하는 ZOH(Zero-Order Hold)입니다. Discrete SSM은 순환 신경망과 유사한 구조를 가지므로 Discrete SSM은 Transformer 기반 모델보다 추론 프로세스를 더 효율적으로 수행할 수 있습니다.

  • 컨볼루션 계산

이산 SSM은 연관 속성을 갖는 선형 시스템이므로 컨벌루션 계산과 원활하게 통합될 수 있습니다.

RNN, Transformer 및 SSM의 관계

그림 2는 RNN, Transformer 및 SSM의 계산 알고리즘을 보여줍니다.



한편, 기존 RNN은 비선형 순환 프레임워크를 기반으로 작동하며 각 계산은 이전 숨겨진 상태와 현재 입력에만 의존합니다.

이 형식을 사용하면 RNN이 자동 회귀 추론 중에 출력을 빠르게 생성할 수 있지만 RNN이 GPU의 병렬 컴퓨팅 성능을 완전히 활용하기 어려워 모델 교육 속도가 느려집니다.

반면 Transformer 아키텍처는 여러 쿼리-키 쌍에 대해 병렬로 행렬 곱셈을 수행하며, 행렬 곱셈을 하드웨어 리소스에 효율적으로 할당할 수 있어 어텐션 기반 모델을 더 빠르게 훈련할 수 있습니다. 그러나 Transformer 기반 모델이 응답이나 예측을 생성하려는 경우 추론 프로세스에 매우 많은 시간이 소요될 수 있습니다.

한 가지 유형의 계산만 지원하는 RNN 및 Transformer와 달리 이산 SSM은 선형 특성으로 인해 매우 유연하며 루프 계산과 컨볼루션 계산을 모두 지원할 수 있습니다. 이 기능을 통해 SSM은 효율적인 추론뿐만 아니라 병렬 교육도 달성할 수 있습니다. 그러나 가장 일반적인 SSM은 시불변입니다. 즉, A, B, C 및 Δ는 모델 입력 x에 독립적입니다. 이로 인해 상황 인식 모델링 기능이 제한되어 SSM이 선택적 복사와 같은 일부 특정 작업에서 성능이 저하됩니다.

맘바

위와 같은 기존 SSM의 단점을 해결하고 상황 인식 모델링을 달성하기 위해 Albert Gu와 Tri Dao는 Universal Sequence 기본 모델의 백본 네트워크로 사용할 수 있는 Mamba를 제안했습니다. Machine Heart 보고서 "5배 처리량, Transformer를 포괄적으로 둘러싼 성능: 새로운 아키텍처 Mamba는 AI 서클을 폭파합니다."

그 후 두 사람은 구조화된 우주 상태 이중성(SSD/Structured Space-State Duality)이 구조화된 SSM을 다양한 형태의 주의와 연결하는 강력한 이론적 프레임워크를 구축한 Mamba-2를 추가로 제안했습니다. Transformer to SSM을 위해 원래 개발된 알고리즘 및 시스템 최적화 기술은 Heart of the Machine 보고서 "Fighting Transformer Again!"을 참조할 수도 있습니다. 원작자가 이끄는 Mamba 2가 등장했고, 새로운 아키텍처의 훈련 효율성이 크게 향상됐다."

Mamba-1: 하드웨어 인식 알고리즘을 사용한 선택적 상태 공간 모델

Mamba-1은 구조화된 상태 공간 모델을 기반으로 하는 세 가지 주요 혁신 기술, 즉 메모리 초기화, 선택 메커니즘, HiPPO(고차 다항식 투영 연산자) 기반의 하드웨어 인식 컴퓨팅을 도입합니다. 그림 3과 같습니다. 이러한 기술의 목표는 SSM의 장거리 선형 시계열 모델링 기능을 향상시키는 것입니다.



특히, 초기화 전략은 장거리 메모리를 효과적으로 촉진하기 위해 일관된 숨겨진 상태 매트릭스를 구성합니다.

그런 다음 선택 메커니즘을 통해 SSM은 인식 가능한 콘텐츠의 표현을 얻을 수 있습니다.

마지막으로 훈련 효율성을 높이기 위해 Mamba에는 병렬 연관 스캔(Parallel Associative Scan)과 메모리 재계산(Memory Recomputation)이라는 두 가지 하드웨어 인식 컴퓨팅 알고리즘도 포함되어 있습니다.

Mamba-2: 상태 공간 이중성

Transformer는 매개변수 효율적인 미세 조정, 치명적인 망각 완화, 모델 양자화 등 다양한 기술 개발에 영감을 주었습니다. 상태 공간 모델이 원래 Transformer용으로 개발된 이러한 기술의 이점을 활용하기 위해 Mamba-2는 SSD(Structured State Space Duality)라는 새로운 프레임워크를 도입했습니다. 이 프레임워크는 이론적으로 SSM과 다양한 형태의 주의를 연결합니다.

본질적으로 SSD는 Transformer에서 사용되는 주의 메커니즘과 SSM에서 사용되는 선형 시불변 시스템이 모두 반분리 가능한 행렬 변환으로 볼 수 있음을 보여줍니다.

또한 Albert Gu와 Tri Dao는 선택적 SSM이 반분리형 마스크 매트릭스를 사용하여 구현된 구조화된 선형 주의 메커니즘과 동일하다는 것을 입증했습니다.

Mamba-2는 블록 분해 행렬 곱셈 알고리즘을 사용하여 하드웨어를 보다 효율적으로 사용하는 SSD 기반의 컴퓨팅 방법을 설계합니다.

특히, Mamba-2는 이러한 행렬 변환을 통해 상태공간 모델을 반분리 가능한 행렬로 처리함으로써 이 계산을 행렬 블록으로 분해할 수 있으며, 여기서 대각선 블록은 블록 내 계산을 나타냅니다. 비대각선 블록은 SSM의 숨겨진 상태 분해를 통한 블록 간 계산을 나타냅니다. 이 방법을 사용하면 Mamba-2의 훈련 속도가 Mamba-1의 병렬 상관 스캔보다 2~8배 더 빠르며 성능은 Transformer와 비슷합니다.

맘바 블록

Mamba-1과 Mamba-2의 블록 디자인을 살펴보겠습니다. 그림 4에서는 두 아키텍처를 비교합니다.



Mamba-1의 설계는 SSM 중심입니다. 여기서 선택적 SSM 계층의 작업은 입력 시퀀스 X에서 Y로의 매핑을 수행하는 것입니다. 이 설계에서는 처음에 X의 선형 투영을 생성한 후 (A, B, C)의 선형 투영이 사용됩니다. 그런 다음, 입력 토큰과 상태 행렬은 병렬 상관을 사용하여 선택적 SSM 장치를 통해 스캔되어 출력 Y를 얻습니다. 이후 Mamba-1은 기능 재사용을 장려하고 모델 훈련 중에 자주 발생하는 성능 저하 문제를 완화하기 위해 건너뛰기 연결을 채택합니다. 마지막으로 Mamba 모델은 이 모듈을 표준 정규화와 잔여 연결을 교대로 쌓아서 구성됩니다.

Mamba-2의 경우 [X, A, B, C]에서 Y로의 매핑을 생성하기 위해 SSD 계층이 도입되었습니다. 이는 표준 주의 아키텍처가 Q, K, V 투영을 병렬로 생성하는 방법과 유사하게 블록 시작 부분에서 단일 투영을 사용하여 [X, A, B, C]를 동시에 처리함으로써 달성됩니다.

즉, Mamba-2 블록은 Mamba-1 블록을 기반으로 시퀀스의 선형 투영을 제거하여 단순화됩니다. 이를 통해 SSD 아키텍처는 Mamba-1의 병렬 선택 스캔보다 더 빠르게 계산을 수행할 수 있습니다. 또한 훈련 안정성을 향상시키기 위해 Mamba-2는 건너뛰기 연결 뒤에 정규화 계층도 추가합니다.

Mamba 모델이 진화하고 있습니다

상태 공간 모델과 Mamba는 최근 빠르게 발전하여 큰 잠재력을 지닌 기본 모델 백본 네트워크 선택이 되었습니다. Mamba는 자연어 처리 작업에서는 잘 수행되지만 Transformer 기반 언어 모델에 비해 메모리 손실, 다양한 작업으로 일반화하기 어려움, 복잡한 패턴에서의 성능 저하 등 여전히 몇 가지 문제가 있습니다. 이러한 문제를 해결하기 위해 연구 커뮤니티에서는 Mamba 아키텍처에 대한 많은 개선 사항을 제안했습니다. 기존 연구는 주로 수정 블록 설계, 스캔 패턴 및 메모리 관리에 중점을 두고 있습니다. 표 1은 관련 연구를 카테고리별로 요약한 것입니다.



블록 디자인

Mamba 블록의 디자인과 구조는 Mamba 모델의 전반적인 성능에 큰 영향을 미치기 때문에 이것이 주요 연구 초점이 되었습니다.



그림 5에서 볼 수 있듯이 기존 연구는 새로운 Mamba 모듈을 구축하는 다양한 방법에 따라 세 가지 범주로 나눌 수 있습니다.

  • 통합 방법: Mamba 블록을 다른 모델과 통합하여 효과와 효율성 간의 균형을 유지합니다.
  • 교체 방법: Mamba 블록을 사용하여 다른 모델 프레임워크의 기본 레이어를 교체합니다.
  • 수정 방법: 클래식 Mamba 블록 내의 구성 요소를 수정합니다.

스캔 모드

병렬 상관 스캐닝은 Mamba 모델의 핵심 구성 요소입니다. 그 목표는 선택 메커니즘으로 인해 발생하는 계산 문제를 해결하고 훈련 프로세스 속도를 향상시키며 메모리 요구 사항을 줄이는 것입니다. 이는 시변 SSM의 선형 특성을 활용하여 하드웨어 수준에서 코어 융합 및 재계산을 설계함으로써 달성됩니다. 그러나 Mamba의 단방향 시퀀스 모델링 패러다임은 이미지, 비디오 등 다양한 데이터에 대한 포괄적인 학습에 도움이 되지 않습니다.



이 문제를 완화하기 위해 일부 연구자들은 Mamba 모델의 성능을 향상하고 훈련 프로세스를 촉진하기 위한 새롭고 효율적인 스캐닝 방법을 모색했습니다. 그림 6에서 볼 수 있듯이 스캐닝 모드 개발 측면에서 기존 연구 결과는 두 가지 범주로 나눌 수 있습니다.

  • 평면화된 스캐닝 방법: 평면화된 관점에서 토큰 시퀀스를 보고 이를 기반으로 모델 입력을 처리합니다.
  • 입체 스캐닝 방법: 차원, 채널 또는 규모에 걸쳐 모델 입력을 스캐닝하며 계층적 스캐닝, 시공간 스캐닝 및 하이브리드 스캐닝의 세 가지 범주로 더 나눌 수 있습니다.

메모리 관리

RNN과 마찬가지로 상태 공간 모델 내에서 숨겨진 상태 메모리는 이전 단계의 정보를 효과적으로 저장하므로 SSM의 전체 성능에 중요한 영향을 미칩니다. Mamba는 메모리 초기화를 위해 HiPPO 기반 방법을 도입했지만 레이어 이전에 숨겨진 정보를 전송하고 무손실 메모리 압축을 달성하는 것을 포함하는 SSM 장치에서 메모리를 관리하는 것은 여전히 ​​어렵습니다.

이를 위해 몇몇 선구적인 연구에서는 메모리 초기화, 압축 및 연결을 포함한 다양한 솔루션을 제안했습니다.

Mamba를 다양한 데이터에 적용

Mamba 아키텍처는 선택적 상태 공간 모델의 확장으로 순환 모델의 기본 특성을 갖고 있으므로 텍스트, 시계열, 음성 등의 시퀀스 데이터를 처리하기 위한 일반적인 기본 모델로 매우 적합합니다.

뿐만 아니라, 최근 일부 선구적인 연구에서는 Mamba 아키텍처의 응용 시나리오를 확장하여 그림 7과 같이 시퀀스 데이터를 처리할 수 있을 뿐만 아니라 이미지, 지도와 같은 분야에서도 사용할 수 있게 되었습니다.



이러한 연구의 목표는 장거리 의존성을 확보하는 Mamba의 탁월한 능력을 최대한 활용하고 학습 및 추론 과정에서 Mamba의 효율성을 활용하는 것입니다. 표 2는 이러한 연구 결과를 간략하게 요약한 것입니다.



시퀀스 데이터

시퀀스 데이터(Sequence Data)는 특정 순서에 따라 수집되고 정리된 데이터를 말하며, 여기서 데이터 포인트의 순서가 중요합니다. 이 검토 보고서는 자연어, 비디오, 시계열, 음성 및 인간 동작 데이터를 포함한 다양한 시퀀스 데이터에 대한 Mamba의 적용을 포괄적으로 요약합니다. 자세한 내용은 원본 논문을 참조하세요.

비순차적 데이터

순차 데이터와 달리 비순차 데이터는 특정 순서를 따르지 않습니다. 데이터 포인트는 데이터의 의미에 큰 영향을 주지 않고 어떤 순서로든 구성될 수 있습니다. 이러한 고유 순서 부족은 데이터의 시간적 종속성을 캡처하도록 특별히 설계된 반복 모델(RNN, SSM 등)의 경우 어려울 수 있습니다.

놀랍게도 최근 일부 연구에서는 Mamba(대표 SSM)가 이미지, 지도, 포인트 클라우드 데이터를 포함한 비순차적 데이터를 효율적으로 처리할 수 있도록 하는 데 성공했습니다.

다중 모드 데이터

AI의 인식 및 장면 이해 능력을 향상시키기 위해 언어(순차 데이터), 이미지(비순차 데이터) 등 다양한 양식의 데이터를 통합할 수 있습니다. 이러한 통합은 매우 가치 있고 보완적인 정보를 제공할 수 있습니다.

최근 몇 년 동안 다중 모드 대형 언어 모델(MLLM)이 가장 인기 있는 연구 초점이 되었습니다. 이러한 유형의 모델은 강력한 언어 표현 및 논리적 추론 기능을 포함하여 대형 언어 모델(LLM)의 강력한 기능을 상속합니다. Transformer가 해당 분야에서 지배적인 방법이 되었지만 Mamba는 혼합 소스 데이터를 정렬하고 시퀀스 길이에 따른 선형 복잡성 확장을 달성하는 성능으로 인해 Mamba가 다중 모드 학습에서 Transformer를 대체할 가능성이 높습니다.

애플리케이션

Mamba 기반 모델의 몇 가지 주목할만한 응용 프로그램이 아래에 소개되어 있습니다. 팀은 이러한 애플리케이션을 자연어 처리, 컴퓨터 비전, 음성 분석, 약물 발견, 추천 시스템, 로봇공학 및 자율 시스템 등의 카테고리로 분류했습니다.

여기서는 너무 많이 소개하지 않겠습니다. 자세한 내용은 원본 논문을 참조하세요.

도전과 기회

Mamba는 일부 분야에서 뛰어난 성과를 거두었지만 전반적으로 Mamba 연구는 아직 초기 단계이며 앞으로 극복해야 할 몇 가지 과제가 남아 있습니다. 물론 이러한 도전은 기회이기도 합니다.

  • Mamba 기반의 기본 모델을 개발하고 개선하는 방법
  • GPU 및 TPU와 같은 하드웨어를 최대한 활용하여 모델 효율성을 향상시키기 위해 하드웨어 인식 컴퓨팅을 완전히 실현하는 방법
  • 보안과 견고성, 공정성, 설명 가능성 및 개인 정보 보호에 대한 추가 연구가 필요한 Mamba 모델의 신뢰성을 향상시키는 방법
  • 매개변수 효율적인 미세 조정, 치명적인 망각 완화, 검색 증강 생성(RAG) 등 Mamba용 Transformer 분야의 새로운 기술을 사용하는 방법.