2023년 12월 소개된 Mamba 아키텍처는 ‘시퀀스 모델링’의 영역에서 매우 중요한 발전 사항 중 하나라고 할 수 있는데요.

튜링 포스트 코리아의 ‘AI 101’ 두 번째 글인 이번 글에서는, Mamba 아키텍처의 출발점, Mamba 아키텍처가 해결해 주는 문제점들을 알아보고, 왜 이 Mamba 아키텍처가 트랜스포머 모델을 대신할 훌륭한 대안이 될 수도 있다고 보는지 살펴보려고 합니다. Context Length를 효율적으로 확장할 수 있고 연산 비용이 선형적으로만 증가하기 때문에, Mamba는 AI가 맞닥뜨리는 대규모 데이터 처리의 난제를 다룰 수 있는 유망한 해결책의 하나로 평가받고 있습니다.

이번 에피소드에서 다뤄보고자 하는 내용은 아래와 같습니다:

먼저 파운데이션 모델, 트랜스포머 및 관련된 몇 가지 개념의 맥락을 짚어보는 것이 좋겠습니다. 시퀀스 모델링은 순차 데이터를 대상으로 하는 연구 분야인데, 즉 ‘내재된 질서가 있는 데이터'를 대상으로 연구를 한다는 뜻이죠 - 텍스트 문장, 이미지, 음성, 오디오, 시계열 자료, 유전체 정보 등이 그 대상이 됩니다.

1990년대, 2000년대에 신경망의 부상, 그리고 컴퓨팅 성능의 향상 덕분에 시퀀스 모델링 영역에 큰 발전이 있었습니다. LSTM 및 RNN (순환 신경망) 같은 모델이 학계와 업계의 표준이 되기도 했습니다. 이 모델들은 ‘가변 길이의 시퀀스’를 처리하고 데이터의 Long-term Dependency (장기적 종속성)를 파악하는데 탁월하지만, 순차적으로 데이터를 처리한다는 특성 때문에 단일 트레이닝 케이스 안에서 연산을 병렬적으로 처리하기는 어렵습니다.

RNN의 문장 처리 과정을 묘사한 그림. Image credit: KDNuggets

2017년, 트랜스포머 아키텍처가 소개되면서 시퀀스 데이터의 처리 방식에 큰 변화가 생겼습니다. 트랜스포머는 Attention Mechanism을 사용하기 때문에 이전 모델에서 사용하던 Recurrence나 Convolution이 필요없습니다. 이런 설계 상의 변화로 모델 트레이닝 중의 병렬 처리 효율이 높아지고, 결국 트랜스포머가 대규모 데이터셋을 효율적이고 확장 가능하게 처리할 수 있도록 해 주게 됩니다.

동일 계층의 하나 또는 그 이상의 Attention Head에 대한 Attention을 보여주는 Head View. 트랜스포머 언어모델에서의 Attention을 시각화하기 위해 BertViz로 작성함.
Image credit: GitHub 

트랜스포머의 Self-Attention Mechanism은 입력 시퀀스의 한 부분과 나머지 모든 부분의 관련성을 계산합니다. 이 작업이 시퀀스의 모든 부분에 대해서 동시에 수행되는데, 따라서 본질적으로 병렬적인 작업입니다. 각각의 Attention Head는 입력 시퀀스의 모든 위치에 대해서 Attention Score를 독립적으로 계산할 수 있기 때문에 이 연산을 동시에 실행 가능한 것이죠. 이런 방법이 바로 출력값이 차례대로 계산되는 RNN과 크게 다른 점이고, 이 점 때문에 트레이닝 시간이 현저하게 빨라지고 더 큰 대규모의 데이터셋을 효율적으로 처리할 수 있는 겁니다.

트랜스포머는 이미 현대적인 파운데이션 모델의 근간이 되었습니다. 그렇지만 트랜스포머가 시퀀스 모델링 영역의 유일한 아키텍처도 아니고, 트랜스포머가 Long Context를 다루게 될 때 몇 가지 문제가 있다는 것도 밝혀졌습니다 - 자, 여기서 Mamba가 등장합니다.

트랜스포머가 많은 장점을 가지고 있지만, 그렇다고 단점이 없는 것은 아닙니다. 우선 한 가지 한계는, Self-Attention Mechanism 때문에 한정된 Context Window를 선택해야 한다는 것입니다. Selt-Attention 덕분에 시퀀스의 각 세그먼트가 다른 세그먼트과의 관련성을 고려하게 되지만 그 범위는 본질적으로 제한되어 있죠. 각 요소는 전체 시퀀스가 아닌 지정된 범위의 주변 요소와만 동시에 상호 작용할 수 있습니다. 이를 "Context Window"라고 부릅니다.

Image Credit: Stanford University course “Natural Language Processing with Deep Learning CS224N/Ling284

더불어, 이 Context Window의 크기를 늘리면 계산 요구사항이 크게 증가하는 구조인데, Window의 길이에 따라 이차 확장 (Quadratic Scaling)됩니다 - 즉, Context Length가 x인 경우 스케일링에 필요한 계산 리소스는 x² (x의 2승)이 됩니다.

엄청나게 증가하는 거죠.

최근, 시퀀스 모델링의 영역에서 SSM (State-Space Sequence Model: 상태 공간 시퀀스모델)이라고 부르는 모델이 트랜스포머를 대체할 수 있는 잠재적 모델로 부상하고 있는데, Mamba도 이 SSM의 한 변종입니다. 이 SSM은 트랜스포머가 내재적으로 가지고 있는 몇 가지 근본적 문제점을 해결해 줍니다.

여기서 잠깐! 물론, Mamba와 같은 SSM 류의 모델들이 트랜스포머가 가진 특정한 제약 조건들을 해결해 주는 것은 사실이지만, 이 외에도 트랜스포머의 효율성을 향상시키기 위해 시도되고 있는 다른 많은 접근 방법도 있다는 것은 아셔야 합니다. “Efficient Transformers: A Survey” 논문에서는 위와 같은 여러가지 접근 방법에 대해 광범위한 논의와 조사를 한 바 있는데, 아래의 이미지에 그 방법들과 모델들이 나타나 있습니다. 예를 들어서, 지난 번 AI 101에서는 아래의 벤 다이어그램 오른쪽 아래에 위치한 ‘Sparse’ 카테고리로 분류된 MoE (Mixture-of-Experts) 모델에 대해 말씀드린 바가 있습니다.

전반적으로, SSM (상태 공간 모델)을 사용하는 접근 방식은 ‘미분 방정식 (Differential Equation)을 사용하여 설명하는 동적 시스템으로 귀결됩니다. 여기서, ‘동적 시스템’이란 인공적, 물리적, 생물학적 시스템을 포함해서 어떤 시스템이든 간에 시간에 따라 변화하는 시스템을 말합니다. 지구 궤도를 도는 우주왕복선, 경쟁하는 종들이 있는 생태계, 단순한 유기체의 신경계 또는 팽창하는 우주 등이 모두 이런 동적 시스템이라고 생각하면 되겠습니다.

‘SSM (상태 공간 모델)’은 많은 다른 과학 분야에서 광범위하게 사용되어 왔지만, 상대적으로 머신 러닝 학계에서는 간과되고 있었는데요. ‘특수 상태 행렬 (Special State Matrices)’의 소개 이전에는 간단한 작업도 쉽게 할 수가 없었다고 해요 - 이런 행렬을 사용하면, SSM은 탁월한 성능을 발휘하게 됩니다.

구조화된 SSM은 스탠포드 대학교의 세 명의 연구자가 2021년에 제안했습니다: Albert Gu, Karan Goel, 그리고 Christopher Re가 그 세 사람인데요, 이 세 사람의 이전 연구에서 확장된 모델이 SSM입니다. 이 모델은 고전적인 ‘상태 공간 원리 (State Space Principle)’를 바탕으로 RNN (순환 신경망)과 CNN (컨볼루션 신경망)의 요소를 혼합한 것입니다 - 그렇지만 불연속적이고 밀도가 높은 데이터 - 텍스트 같은 데이터겠죠 - 에서는 성능이 떨어지는 경향이 있습니다.

초기에 등장한 SSM 컨셉을 기반으로, 카네기 멜론 대학과 프린스턴 대학의 연구원들이 트랜스포머와 기존 SSM의 단점을 모두 극복한 혁신적인 ‘선택적 상태 공간 모델 (Selective SSM)’인 Mamba 모델을 만들었습니다.

Selective SSM의 다이어그램. Image Credit: 논문

Mamba 아키텍처는 Attention이나 MLP (Multilayer Perceptron; 다층 퍼셉트론)과 같은 기존 구성 요소를 제거하고 ‘선택적 상태 공간 모델 (Selective SSM)’을 신경망 프레임웍에 통합합니다.

Mamba 아키텍처의 주요 특징:

왜 Mamba가 큰 관심을 받을까요? 짧게 대답하자면, ‘더 긴 Context Window (Larger Context Window)’ 때문입니다. 생성AI에 관심있는 분이라면, 최근 많은 LLM 사업자들이 LLM의 Context Length의 확장 경쟁을 벌이고 있는 것을 아실 텐데요 - 대표적 사례로 ‘무손실 Long Context’를 자사 기술 개발의 핵심 전략으로 채택한 중국의 생성 AI 유니콘 Moonshot AI가 있죠 - 이 스타트업에 대한 튜링 포스트 코리아의 기사도 있으니 참고하세요. Mamba 아키텍처는 Context Length를 효율적으로 확장할 수 있는 훌륭한 방법입니다.

Mamba는 그 단순화된 구조에도 불구하고 최첨단의 성능을 구현, 언어 모델링과 같은 작업에서 크기가 두 배나 큰 트랜스포머보다 우수하거나 비슷한 기능을 보여줍니다.

Mamba의 주요 이점은:

Mamba 아키텍처의 등장은 트랜스포머 모델의 등장 이후 아주 흥미로운 다음 단계의 발전이라고 할 수 있을 것 같습니다 - 파운데이션 모델의 영역에서 보다 나은 수준의 효율성과 확장성을 달성할 수 있는 가능성이 높은 접근 방식이죠. 그리고 이 Mamba 아키텍처의 발견 또는 개발, 그리고 이에 대한 높은 관심은, 지속적으로 이전의 머신 러닝 연구 중 어떤 혁신적인 접근이 있었는지 돌아보고, 학제 간 통섭적인 사고가 중요하다는 것의 반증이기도 합니다.

읽어주셔서 감사합니다. 친구와 동료 분들에게도 뉴스레터 추천해 주세요!