혁신적인 SSM-Transformer 모델인 Jamba 소개

Name

Lynn Mikami

Published on 2024. 4. 30.

AI21 Labs은 혁신적인 Mamba 아키텍처를 기반으로 한 최초의 상업용 모델인 Jamba를 자랑스럽게 선보입니다. Mamba Structured State Space (SSM) 기술을 전통적인 Transformer 아키텍처의 요소와 원활하게 통합함으로써, Jamba는 순수 SSM 모델의 한계를 극복하고 우수한 성능과 효율성을 제공합니다.

256K 컨텍스트 창과 놀라운 처리량 향상을 갖춘 Jamba는 AI 기술의 새로운 가능성을 열어주며, 연구원, 개발자 및 기업들에게 새로운 기회를 제공합니다. Jamba는 이미 다양한 벤치마크에서 뛰어난 결과를 보여주었으며, 동일한 크기 범주의 최첨단 모델들과 일치하거나 뛰어난 결과를 보여줍니다.

TLDR: Jamba에는 안전 중재 메커니즘과 가드레일이 없으며, Apache-2.0 오픈 소스 라이선스를 사용합니다.

Jamba의 주요 기능

1. 최초의 상업용 Mamba 기반 모델: Jamba는 상업용 규모와 품질로 SSM-Transformer 하이브리드 아키텍처를 도입한 최초의 모델입니다.
2. 무시할 수 없는 처리량: Jamba는 Mixtral 8x7B와 비교하여 긴 컨텍스트에 대해 3배의 처리량을 달성하여 효율성을 높였습니다.
3. 대용량 컨텍스트 창: 256K 크기의 컨텍스트 창을 갖춘 Jamba는 다양한 컨텍스트 처리 능력을 데모크라타이즈화합니다.
4. 단일 GPU 호환성: Jamba는 단일 GPU에서 최대 140K 컨텍스트를 수용할 수 있는 해당 크기 범주의 유일한 모델로, 배포 및 실험에 더욱 접근성이 좋습니다.
5. 오픈 소스 이용 가능: Apache 2.0 라이선스에 따라 오픈 가중치로 출시된 Jamba는 AI 커뮤니티로부터 추가 최적화와 발견을 유도합니다.
6. NVIDIA API 카탈로그 통합 예정: Jamba는 NVIDIA API 카탈로그에서 NVIDIA NIM 추론 마이크로서비스로서 이용 가능해지며, 기업 애플리케이션 개발자들은 NVIDIA AI 엔터프라이즈 소프트웨어 플랫폼을 사용하여 배포할 수 있게 됩니다.

Jamba: Mamba와 Transformer 아키텍처의 최고를 결합

Jamba는 Mamba를 Transformer 아키텍처와 성공적으로 결합시키고, 하이브리드 SSM-Transformer 모델을 상업용 품질로 확장한 LLM 혁신의 중요한 단계를 나타냅니다.

기존의 Transformer 기반 LLM은 두 가지 주요 도전 과제를 가지고 있습니다:

1. 대형 메모리 사용: Transformer의 메모리 사용량은 컨텍스트 길이와 함께 증가하여, 긴 컨텍스트 창이나 여러 병렬 배치를 하드웨어 리소스 없이 실행하기 어렵게 만듭니다.
2. 긴 컨텍스트에서의 느린 추론: Transformer의 어텐션 메커니즘은 시퀀스 길이에 따라 제곱 관계로 확장되므로, 각 토큰이 전체 이전 시퀀스에 의존함에 따라 처리량이 느려집니다.

Carnegie Mellon 대학과 Princeton 대학의 연구진이 제시한 Mamba는 이러한 단점을 해결합니다. 그러나 전체 컨텍스트를 대상으로한 어텐션이 없기 때문에, Mamba는 특히 리콜 관련 작업에서 최상의 기존 모델의 출력 품질을 맞추기 어렵습니다.

Jamba vs Mamba vs Transformer

Jamba의 하이브리드 아키텍처는 Transformer, Mamba 및 mixture-of-experts (MoE) 레이어로 구성되어 메모리, 처리량 및 성능을 최적화합니다. MoE 레이어를 사용함으로써 Jamba는 추론 중 52B의 파라미터 중 12B만을 사용하여, 동등한 크기의 Transformer-만 모델보다 더 효율적인 액티브 파라미터를 활용할 수 있게 됩니다.

Jamba의 하이브리드 아키텍처 확장

Jamba의 하이브리드 구조를 성공적으로 확장하기 위해, AI21 Labs은 여러 가지 핵심 아키텍처 혁신을 시행했습니다:

1. 블록과 레이어 접근 방식: Jamba의 아키텍처는 블록과 레이어 접근 방식을 특징으로 합니다. Transformer와 Mamba 아키텍처를 원활하게 통합할 수 있습니다. 각 Jamba 블록은 어텐션 레이어 또는 Mamba 레이어의 후속으로 멀티레이어 퍼셉트론(MLP)이 이어집니다. 전체적으로는 여덟 개의 레이어 중 하나가 Transformer 레이어로 구성됩니다.

2. 전문가집합(mixture-of-experts, MoE) 활용: MoE 레이어를 사용함으로써, Jamba는 추론 중 사용되는 액티브 파라미터의 수를 최소화하면서 전체적인 모델 파라미터 수를 늘릴 수 있습니다. 이로써 계산 요구사항을 증가시키지 않고 모델 용량을 높일 수 있습니다. MoE 레이어 및 전문가 수는 단일 80GB GPU에서 모델의 품질과 처리량을 극대화하기 위해 최적화되었으며, 일반적인 추론 작업에 충분한 메모리를 남겨둡니다.

Jamba의 뛰어난 성능과 효율성

Jamba vs Llama 70B vs Mixtral 8x7B

Jamba의 초기 평가는 처리량 및 효율성과 같은 주요 지표에서 뛰어난 결과를 나타내었습니다. 이런 벤치마크는 커뮤니티가 이 혁신적인 기술을 실험하고 최적화함에 따라 더욱 개선될 것으로 예상됩니다.

1. 효율성: Jamba는 긴 문맥에서 3배의 처리량을 제공하여 Mixtral 8x7B와 같은 유사한 Transformer 기반 모델보다 효율적입니다.
2. 비용 효율성: 단일 GPU에서 140K 문맥을 적용할 수 있는 능력으로 인해 Jamba는 비슷한 크기의 다른 오픈 소스 모델에 비해 더 쉽고 접근 가능한 배포 및 실험 기회를 제공합니다.

향후 개선 사항으로는 향상된 MoE 병렬화 및 더 빠른 Mamba 구현이 이미 인상적인 이득을 더욱 높여줄 것으로 예상됩니다.

Jamba로 빌드 시작하기

Jamba는 이제 Apache 2.0 라이선스로 오픈 가중치와 함께 Hugging Face에서 사용할 수 있습니다. 기본 모델로서, Jamba는 Fine-tuning, 훈련 및 사용자 정의 솔루션 개발을 위한 기초로 제공됩니다. 책임있고 안전한 사용을 위해 적절한 가드레일을 추가하는 것이 중요합니다.

Jamba의 강화 버전은 곧 AI21 플랫폼을 통해 베타 버전으로 제공될 예정입니다. 프로젝트를 공유하거나 피드백 또는 질문을 하려면 Discord에서 대화에 참여하세요.

결론

Jamba의 소개는 하이브리드 SSM-Transformer 아키텍처의 엄청난 잠재력을 보여주며 AI 기술에서 큰 발전을 나타냅니다. Mamba와 Transformer의 강점을 결합하고 효율성과 성능을 최적화함으로써, Jamba는 자신의 크기 범주에서 AI 모델에 대한 새로운 기준을 설정합니다.

인상적인 문맥 처리, 처리량 및 비용 효율성을 갖춘 Jamba는 AI 분야를 혁신시키고 연구원, 개발자 및 기업이 가능한 선을 넘어 나아갈 수 있도록 조건을 만들고 있습니다. 커뮤니티가 Jamba의 혁신을 탐구하고 발전시키는 동안, 인공지능의 미래를 형성할 새로운 파동의 AI 응용 프로그램을 기대할 수 있습니다.