맴바: 시퀀스 모델링을 재정의하고 트랜스포머 아키텍처를 초월하다

이 글에서는 Albert Gu와 Tri Dao가 개발한 상태 공간 모델(SSM) Mamba가 시퀀스 모델링을 어떻게 바꾸는지 살펴본다. Mamba는 언어, 유전체학, 오디오 분석의 복잡한 시퀀스를 효율적으로 처리한다. 선택적 상태 공간을 사용해 계산 비용이 시퀀스 길이에 선형적으로 증가한다.

전통적인 Transformer가 긴 시퀀스에서 겪는 계산 문제에 대응하며, 빠른 추론과 높은 처리량을 제공한다. 빠른 처리, 선택적 SSM 계층, FlashAttention에서 영감을 얻은 하드웨어 친화적 설계를 결합한다. 원문에 제시된 결과에서는 Transformer 기반 모델을 포함한 여러 기존 모델보다 좋은 성능을 보인다.

Transformer와 Mamba

GPT-4와 같은 Transformer는 자연어 처리의 기준을 높였지만 긴 시퀀스에서는 계산 효율이 떨어진다. Mamba는 다른 구조로 긴 시퀀스를 효율적으로 처리한다.

Transformer는 언어 모델의 텍스트처럼 순서가 있는 데이터를 처리한다. 순차적으로만 처리하던 이전 모델과 달리, 학습 중 시퀀스의 여러 위치를 병렬 처리하여 복잡한 관계를 포착할 수 있다. 어텐션은 예측할 때 입력의 어느 부분에 집중할지 결정한다. 입력에서 구한 쿼리, 키, 값 표현을 사용하며, 요소 사이의 비교로 상대적 중요도를 나타내는 가중치를 계산한다.

여기서 설명하는 인코더–디코더형 Transformer에서는 인코더가 입력을 처리하고 디코더가 출력을 생성한다. 이 구조가 모든 Transformer에 적용되는 것은 아니다. 인코더의 각 층은 다중 헤드 자기 어텐션과 위치별 완전 연결 피드포워드 네트워크로 구성된다. 각 하위 층의 정규화와 잔차 연결은 깊은 네트워크의 학습을 돕는다. 디코더에는 인코더 출력에 대한 어텐션이 추가된다. 인과적 마스크는 현재 위치의 예측이 이전 위치만 참조하게 하여 자기회귀 특성을 유지한다.

Mamba는 선택적 상태 공간을 사용한다. 원문 참고: 해당 비교 문장의 영어 원문은 단어 중간에서 끊겨 있다. 누락된 결론을 임의로 보충하지 않았다.

Transformer의 주요 처리 단계는 다음과 같다.

• 입력 처리: 임베딩과 위치 정보를 통해 입력을 모델이 사용할 표현으로 변환한다.
• 어텐션: 현재 요소를 이해하는 데 다른 위치가 얼마나 중요한지 점수를 구한다.
• 인코더–디코더: 이 구성에서는 여러 층이 입력 표현을 정교하게 만들고 출력을 생성한다.
• 다중 헤드 어텐션: 서로 다른 표현 공간에서 여러 부분을 동시에 참조해 다양한 맥락을 학습한다.
• 위치별 피드포워드 네트워크: 같은 변환을 각 위치에 독립적으로 적용한 뒤 잔차 연결과 층 정규화를 사용한다.
• 출력 생성: 디코더가 인코더의 맥락과 이미 생성한 내용을 바탕으로 다음 출력을 예측한다.

병렬 처리와 어텐션은 번역과 텍스트 생성에 유용하다. Mamba는 긴 시퀀스에서 커지는 비용에 선택적 상태 공간으로 대응하며, 추론 속도를 높이고 시퀀스 길이에 따라 선형적으로 확장한다.

Mamba

Mamba는 전통적인 어텐션 및 다층 퍼셉트론(MLP) 블록에서 벗어나 더 가볍고 빠르며 균일한 구조를 사용한다.

• 선택적 SSM: 관련 없는 정보를 걸러 내고 유용한 데이터에 집중하여 내용에 따른 추론을 지원한다.
• 하드웨어를 고려한 알고리즘: 특히 GPU에 최적화된 병렬 알고리즘으로 계산을 빠르게 하고 메모리 요구량을 줄인다.
• 간소화된 구조: 선택적 SSM을 통합하고 기존 어텐션 및 MLP 블록을 제거하여 확장성과 성능을 높인다.

보고된 결과는 언어, 오디오, 유전체학의 사전 학습과 분야별 작업을 포함한다. 언어 모델링 비교에서는 더 큰 Transformer와 비슷하거나 더 나은 성능을 보인다. 코드와 사전 학습 모델은 GitHub에 공개되어 있다.

일반 복사 작업은 선형 모델에도 쉽다. 선택적 복사와 귀납 헤드는 내용을 고려하는 동적 메모리가 필요하다.

구조화된 상태 공간 모델 S4는 순환 신경망(RNN), 합성곱 신경망(CNN), 고전적 상태 공간 모델의 특성을 결합한다. 잠재 상태를 통해 1차원 함수나 시퀀스를 매핑하는 연속 시스템에서 출발하며, 딥러닝에서 효율적이고 유연한 시퀀스 모델을 설계하는 방법을 제공한다.

S4 모델의 동역학

구조화된 SSM은 입력 시퀀스 x를 받아 학습한 매개변수 A, B, C와 시간 간격 Δ로 출력 y를 구한다. 연속 매개변수를 이산화한 뒤 SSM 연산을 적용한다. 기본적인 시불변 형태에서는 모든 시간 단계에 같은 매개변수를 사용한다.

이산화의 중요성

정해진 수식으로 연속 매개변수를 이산 매개변수로 바꾸면서 연속 시간 시스템과의 연결을 유지한다. 해상도 불변성과 적절한 정규화 같은 특성이 안정성과 성능을 뒷받침한다. 이산화는 정보 흐름을 제어하는 RNN의 게이팅 메커니즘과도 관련이 있다.

선형 시불변성(LTI)

S4의 핵심 특성은 동역학이 시간에 따라 동일하고 매개변수가 모든 단계에서 고정된다는 것이다. 이는 재귀와 합성곱 계산의 기반으로, 단순하면서도 강력한 시퀀스 모델링 틀을 제공한다.

근본적인 제약 극복

LTI는 적응적인 동역학이 필요한 데이터를 모델링하는 데 제약이 된다. 연구에서는 시간에 따라 변하는 매개변수를 도입해 이 제약을 없애고 더 다양한 시퀀스와 작업을 처리한다.

상태 공간 모델은 여러 분야에서 잠재 상태를 가진 재귀 과정을 폭넓게 가리킨다. 딥러닝의 S4 또는 구조화된 SSM은 복잡한 시퀀스 표현을 유지하면서 효율적인 계산을 위해 최적화한 특정 모델군이다. 종단 간 신경망에 독립적인 시퀀스 변환으로 통합할 수 있으며, CNN의 합성곱 층과 비슷한 기반 역할을 한다.

SSM과 선택 메커니즘을 더한 SSM

시퀀스 모델링에 선택성이 필요한 이유

구조화된 SSM

시퀀스 모델링의 핵심은 맥락을 관리 가능한 상태로 압축하는 것이다. 입력을 선택적으로 걸러 내거나 집중하면 이 압축 상태에 유용한 정보를 더 효과적으로 유지할 수 있다. 시퀀스를 따라 정보가 흐르는 방식을 적응적으로 제어하는 능력은 복잡한 언어 작업 등에 중요하다.

선택적 SSM에서는 매개변수가 입력에 의존한다. 따라서 시간에 따라 변하는 SSM이 되며, 시불변 모델처럼 합성곱으로 효율적으로 계산할 수 없으므로 선형 비용의 재귀 계산을 사용한다.

SSM에 선택성을 더한 S6에서는 B, C, Δ가 입력 x에 따라 달라진다. 선택성을 고려해 이산화하고, 스캔 연산으로 시퀀스를 처리하면서 시간 단계마다 집중 대상을 동적으로 조절한다.

Mamba의 주요 성능 결과

원문의 도표는 표시된 모든 평가에서 Mamba가 해당 범주 최고 성능이라고 제시한다.

보고된 시험에서 Mamba는 추론 속도와 정확도에서 좋은 결과를 보인다. DNA와 오디오 모델링에서 긴 맥락을 잘 활용하며, 장거리 의존성이 필요한 복잡한 작업에서 기존 모델보다 좋은 성능을 낸다. 여러 작업의 제로샷 평가도 유연성, 효율성, 확장성을 보여 준다.

Mamba 시작하기

2023년 원문에 제시된 요구 사항은 Linux, NVIDIA GPU, PyTorch 1.12+, CUDA 11.6+이다. Mamba 저장소의 패키지를 pip로 설치한다. PyTorch 버전 호환성 문제가 있으면 원문은 --no-build-isolation 옵션을 제안한다. 모델은 Pile과 SlimPajama 같은 대규모 데이터셋으로 학습되어 다양한 계산 요구를 지원한다.

인터페이스는 선택적 SSM 층, Mamba 블록, 완전한 언어 모델까지 제공된다. 핵심 Mamba 블록은 인과적 Conv1d 층을 사용하며 신경망에 통합하기 쉽다. 원문 참고: 영어 기사는 모델을 생성하고 데이터를 처리하는 Python 예제를 언급하지만 실제 코드는 포함하지 않는다.

Hugging Face의 사전 학습 모델은 1억 3천만 개에서 28억 개까지의 매개변수를 갖는다. Pile과 SlimPajama로 학습했으며 GPT-3의 차원 기준을 따른다. 다양한 계산·성능 요구를 충족하도록 설계되었고, 원문은 높은 처리량과 정확도로 언어 모델링을 비롯한 여러 응용에서 경쟁력 있는 선택이라고 설명한다.

Mamba의 영향

Mamba는 정보가 풍부한 데이터를 처리하는 Transformer의 강력한 대안이다. 현대 하드웨어에 맞춰 메모리와 병렬 처리를 최적화한다. 공개 코드와 사전 학습 모델은 AI와 딥러닝 연구자 및 개발자가 활용할 수 있는 도구를 제공한다.