[AI News #02] AI21 Labs, 세계 최초의 Mamba 기반 상용 수준 모델 ‘Jamba’ 출시
AI21 Labs가 세계 최초로 Mamba 기반의 상용 수준 모델 ‘Jamba’를 출시했다. Jamba의 기반이 되는 Mamba는 새로운 구조의 SSM(Structured State Space Model, 구조화 상태 공간 모델)아키텍처이다. 이번에 새로 출시된 Jamba는 전통적인 트랜스포머 아키텍처와 Mamba를 결합하여 각 아키텍처의 강점을 살렸기에, 기존 모델들과 비교하여 뛰어난 성능을 보인다. Jamba는 256K 토큰의 기존 모델들 대비 긴 컨텍스트 윈도우(Context Window)를 가지면서도, 약 3배 높은 처리량을 제공한다. 이처럼 Jamba는 긴 문맥에서의 처리량 향상을 보여주며, 비용 효율성과 접근성을 우선시하는 모델이다.
Transformer vs Mamba: SSM
현재 딥러닝은 사실상 ”트랜스포머 전성시대”라 할 수 있다. 텍스트를 생성하는 LLM(Large Language Model, 대규모 언어 모델)도 트랜스포머 아키텍처로 이루어져 있으며, 그림을 생성하는 디퓨전 모델 또한 트랜스포머를 차용한다. 그러나, 이렇게 광범위한 딥러닝 분야에서 사용되는 트랜스포머 아키텍처는 훈련 비용이 높다는 한계를 갖는다. 이러한 상황 속에서, ‘Mamba’ 아키텍처의 등장은 새로운 패러다임의 서막으로 여겨진다.
전통적인 트랜스포머 아키텍처의 경우, 시퀀스가 길어지면 효율성이 떨어진다는 한계를 갖고 있다. 특히, 트랜스포머는 훈련에서 2차(quadratic)적 시간 복잡성을 갖기에 텍스트 시퀀스가 커질수록 트랜스포머 계산 요구 사항이 상당이 커졌다. 지난 몇 년 동안 이를 타개하기 위해 여러 아키텍처들이 제안되었지만, 성능 저하로 인해 트랜스포머를 대체할 수 없었다.
카네기 멜론 대학교와 프린스턴 대학교의 연구팀에 의해 개발된 ’Mamba‘는 트랜스포머와는 차별화된 구조를 갖고 있다. Mamba는 선택적 상태 공간 모델(Selective SSM) 기반으로, 더 빠른 추론 속도와 훈련 속도, 그리고 언어 처리, 오디오 분석 등 복잡한 시퀀스에 대한 뛰어난 성능을 보인다. 더 낮은 메모리 사용량으로 더 빠른 계산을 보여주는 혁신적인 아키텍처인 것이다. 이러한 이유로 Mamba는 전문가들 사이에서 주목받기 시작되었다.
트랜스포머의 핵심에 어텐션 모델이 있는 것처럼, 맘바의 핵심에는 선택적 SSM(Selective State Space Model)이 있다. SSM은 1960년대 상태 공간 모델에서 영감을 받은 다소 새로운 언어 모델링 아키텍처다. 맘바(Mamba)는 선택적 집중을 통해 입력에 따라 특정 정보에 집중하거나 무시할 수 있다. 입력에 기반하여 SSM(Selective State Space Model) 가중치를 매개변수화함으로써, 모델이 관련 없는 정보를 걸러내고 관련 있는 데이터를 무기한 유지할 수 있게 하는 것이다. 간단히 말해, 이 모델은 컨텍스트 역할을 하는 ‘상태’ 또는 메모리를 유지하고, 다음 출력은 현재 입력과 그 시점까지 내 현재 상태의 함수가 된다.
선택적 SSM은 중요한 문맥만 메모리에 유지하기 때문에, 훈련과 추론을 위한 복잡성이 선형적이고 일정하다. 즉 시퀀스가 2배가 되면 트랜스포머의 경우 2차적 시간 복잡성을 가지기 때문에 훈련 비용이 4배나 되는 반면, 맘바의 경우 선형적 시간 복잡성을 갖기에 훈련 비용이 2배가 된다. 또한, 트랜스포머는 추론에서 선형적 복잡성을 갖지만, 맘바의 경우 일정한 복잡성을 갖는다. 즉, 시퀀스 길이에 상관없이 맘바의 경우 추론 비용이 증가하지 않는다는 것을 의미한다. 이는 상태 크기가 고정되어 있고 모델이 핵심 정보만 저장하고 나머지는 잊어버리기 때문에 가능하다.
이렇게 선택적 SSM을 사용한 선형 시간 시퀀스 모델링 아키텍처인 맘바는 특히 긴 시퀀스에서 전통적 트랜스포머 아키텍처가 직면한 비용 문제를 극복할 수 있다. SSM의 선택적 접근방식을 통해 추론 속도가 빨라지고, 시퀀스 길이에 따른 선형 확장이 가능해 처리량이 크게 향상되기 때문이다.
Jamba 모델
Jamba는 SSM과 트랜스포머 아키텍처의 장점을 결합한 세계 최초의 Mamba 기반 상용 수준 모델이다. 이 하이브리드 구조는 트랜스포머의 강력한 언어 이해 능력과 SSM의 효율적인 메모리 관리 및 처리 속도를 채택하여 활용한다. Jamba는 MoE 레이어를 통해 추론 시 사용 가능한 52B 파라미터 중 단 12B만을 활용하기에, 모델의 활성화 파라미터가 더 효율적으로 사용된다. 이 결과로, 동일한 크기의 트랜스포머 전용 모델보다 더 나은 성능을 보이는 것이다.
또한, Jamba는 256K의 대규모 컨텍스트 윈도우를 제공함으로써, 매우 긴 문서나 대화를 효율적으로 처리할 수 있다. 단일 GPU에서 최대 140K 컨텍스트까지 지원하는 동급 유일 모델로, 비용 효율적인 퍼포먼스를 보여준다. 기존의 언어 모델들이 가지고 있던 메모리 사용량 증가와 처리 속도 저하 문제를 크게 개선한다는 측면에서 기대를 한 몸에 받고 있는 모델이다. Jamba는 Apache 2.0 라이선스 하에 공개되었고, 허깅 페이스를 통해 사용이 가능하다. 현재는 상업적 사용에 필요한 안전장치 없이 연구 모델로 공개되었으나 AI21 Labs는 향후 몇 주 안에 더 안전한 버전을 출시할 계획이라 밝혔다.
Reference
- Mamba: Linear-Time Sequence Modeling with Selective State Spaces | OpenReview
- AI21 Labs Unveils Jamba: The first Production-Grade Mamba-Based AI Model | Maginative
- Mamba:트랜스포머를 대체할 차세대 아키텍처의 등장 | 모두의 연구소
대학생 인공지능 연합 동아리 Prometheus
2024년 4월 12일
Editor : 김민지