Jamba: AI21 랩스의 새로운 하이브리드 Transformer-Mamba 언어 모델

언어 모델은 자연어 처리에서 급격한 발전을 이루었으며, Transformer 기반 아키텍처가 주도하고 있습니다. 그러나 모델의 규모가 커짐에 따라 긴 문맥을 처리하고, 메모리 효율성 및 처리량을 관리하는 문제가 더 심각해졌습니다.

AI21 랩스는 Jamba라는 새로운 솔루션을 도입했습니다. Jamba는 Transformer와 Mamba 아키텍처의 강점을 결합한 하이브리드 프레임워크를 사용하는 최첨단的大규모 언어 모델(Large Language Model, LLM)입니다. 이 기사에서는 Jamba의 아키텍처, 성능 및 잠재적인 응용 분야에 대해 자세히 설명합니다.

Jamba 개요

Jamba는 AI21 랩스에서 개발한 하이브리드 大규모 언어 모델로, Transformer 레이어와 Mamba 레이어를 결합하고, Mixture-of-Experts (MoE) 모듈을 통합했습니다. 이 아키텍처는 Jamba가 메모리 사용량, 처리량 및 성능을 균형 있게 관리할 수 있도록 해줍니다. 모델은 단일 80GB GPU에 맞춤으로 설계되어 높은 처리량과 작은 메모리 풋프린트를 제공하면서 다양한 벤치마크에서 최첨단 성능을 발휘합니다.

Jamba 아키텍처

Jamba의 아키텍처는 모델의 핵심입니다. 그것은 Transformer 레이어와 Mamba 레이어를 교대로 결합하는 새로운 하이브리드 디자인을 기반으로 하며, MoE 모듈을 통합하여 모델의 용량을 증가시키지 않고도 계산 요구를 크게 증가시키지 않습니다.

1. Transformer 레이어

Transformer 아키텍처는 현대적인 LLM의 표준이 되었습니다. 그것은 병렬 처리를 효율적으로 처리하고 텍스트의 긴 의존성을 캡처하는 능력으로 인해 이러한 지위를 얻었습니다. 그러나 Transformer의 성능은 종종 높은 메모리 및 컴퓨팅 요구로 제한됩니다. 특히 긴 문맥을 처리할 때 이러한 제한이 두드러집니다. Jamba는 이러한 제한을 해결하기 위해 Mamba 레이어를 통합합니다.

2. Mamba 레이어

Mamba는 시퀀스의 긴 거리 관계를 처리하는 데 더 효율적인 최근 상태 공간 모델(State Space Model, SSM)입니다. Mamba 레이어는 Transformer의 키-값 캐시를 저장하는 데 관련된 메모리 풋프린트를 줄이는 데 특히 효과적입니다. Transformer 레이어와 Mamba 레이어를 교대로 결합함으로써 Jamba는 전체 메모리 사용량을 줄이면서 높은 성능을 유지합니다. 특히 긴 문맥을 처리하는 작업에서 이러한 성능이 두드러집니다.

3. Mixture-of-Experts (MoE) 모듈

Jamba의 MoE 모듈은 모델의 용량을 확장하는 데 유연한 접근 방식을 제공합니다. MoE는 모델이 활성 매개변수의 수를 비례적으로 증가시키지 않고도 사용 가능한 매개변수의 수를 증가시킬 수 있습니다. Jamba에서 MoE는 일부 MLP 레이어에 적용되며, 라우터 메커니즘이 각 토큰에 대해 활성화할 전문가를 선택합니다. 이 선택적 활성화는 Jamba가 높은 효율성을 유지하면서 복잡한 작업을 처리할 수 있도록 합니다.

아래 이미지는 하이브리드 Attention-Mamba 모델의 유도 헤드의 기능을示す 예입니다. 이 예에서 주의 헤드는 감성 분석 작업에서 “Positive” 또는 “Negative”와 같은 레이블을 예측하는 데 책임이 있습니다. 강조된 단어는 모델의 주의가 몇샷 예에서 레이블 토큰에 강하게 집중되어 있음을 보여줍니다. 특히 레이블을 예측하기 직전에 이러한 주의 메커니즘은 모델의 문맥 학습 능력에 중요한 역할을 합니다.

MoE를 통합한 Attention-Mamba 하이브리드 아키텍처의 성능 개선은 표에서 강조됩니다. MoE를 사용하여 Jamba는 계산 비용을 비례적으로 증가시키지 않고도 용량을 증가시킵니다. 이는 다양한 벤치마크에서 두드러진 성능 향상으로 나타납니다. 예를 들어, WinoGrande에서 MoE를 사용하면 66.0%의 정확도를 달성할 수 있으며, 이는 MoE 없이 62.5%인 것보다 높습니다. 또한 다양한 도메인에서 로그 확률을 개선합니다(예: C4에서 -0.534).

주요 아키텍처 특징

• 레이어 구성: Jamba의 아키텍처는 Mamba와 Transformer 레이어를 특정 비율(예: 1:7, 즉 7개의 Mamba 레이어당 1개의 Transformer 레이어)로 결합하는 블록으로 구성됩니다. 이 비율은 최적의 성능과 효율성을 위해 조정됩니다.
• MoE 통합: MoE 레이어는 몇 개의 레이어마다 적용되며, 16개의 전문가가 사용 가능하고 각 토큰당 상위 2개의 전문가가 활성화됩니다. 이 구성은 Jamba가 효과적으로 확장하면서 메모리 사용량과 계산 효율성 간의 트레이드오프를 관리할 수 있도록 합니다.
• 정규화 및 안정성: 훈련 중 안정성을 보장하기 위해 Jamba는 Mamba 레이어에서 RMSNorm을 통합합니다. 이는 큰 활성화 스파이크와 같은 문제를 완화하는 데 도움이 됩니다.

Jamba의 성능 및 벤치마킹

Jamba는 다양한 벤치마크에서 경쟁력 있는 성능을 보여주었습니다. 다음 섹션에서는 Jamba가 일반적인 NLP 작업과 긴 문맥 시나리오에서 두드러진 성능을 보여주는 주요 벤치마크를 강조합니다.

1. 일반적인 NLP 벤치마크

Jamba는 여러 학술 벤치마크에서 평가되었습니다.

• HellaSwag (10샷): Jamba는 87.1%의 성능 점수를 달성하여 많은 경쟁 모델을 초과했습니다.
• WinoGrande (5샷): Jamba는 82.5%의 점수를 달성하여 복잡한 언어적 추론을 처리하는 능력을 보여주었습니다.
• ARC-Challenge (25샷): Jamba는 64.4%의 점수를 달성하여 어려운 다중 선택 질문을 관리하는 능력을 보여주었습니다.

집계 벤치마크인 MMLU (5샷)에서 Jamba는 67.4%의 점수를 달성하여 다양한 작업에서 강건함을 보여주었습니다.

2. 긴 문맥 평가

Jamba의 두드러진 기능 중 하나는 매우 긴 문맥을 처리하는 능력입니다. 모델은 최대 256K 토큰의 문맥 길이를 지원하며, 이는 공개적으로 사용 가능한 모델 중 가장 긴 것입니다. 이 기능은 Needle-in-a-Haystack 벤치마크를 사용하여 테스트되었으며, Jamba는 다양한 문맥 길이에서 예외적인 검색 정확도를 보여주었습니다.

3. 처리량 및 효율성

Jamba의 하이브리드 아키텍처는 특히 긴 시퀀스에서 처리량을 크게 개선합니다.

다른 모델과 비교하여 처리량(초당 토큰 수)을 테스트한 결과, Jamba는 특히 큰 배치 크기와 긴 문맥을 포함하는 시나리오에서 일관되게 최고의 성능을 보여주었습니다. 예를 들어, 128K 토큰의 문맥에서 Jamba는 비교 모델인 Mixtral의 3배의 처리량을 달성했습니다.

Jamba 사용: Python

개발자와 연구자들이 Jamba와 실험하고 싶다면, AI21 랩스는 Hugging Face와 같은 플랫폼에서 모델을 제공하여 다양한 응용 분야에 쉽게 접근할 수 있도록 했습니다. 다음 코드 조각은 Jamba를 로드하고 텍스트를 생성하는 방법을示합니다.

<p>from transformers import AutoModelForCausalLM, AutoTokenizer</p>

<p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")</p>
<p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p>

<p>input_ids = tokenizer("최근 Super Bowl LVIII에서", return_tensors='pt').to(model.device)["input_ids"]</p>

<p>outputs = model.generate(input_ids, max_new_tokens=216)</p>

print(tokenizer.batch_decode(outputs))

import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from peft import LoraConfig
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments

<p>tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")</p>
<p>model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)</p>

<p>lora_config = LoraConfig(r=8, target_modules=["embed_tokens","x_proj","in_proj","out_proj", # mamba "gate_proj","up_proj","down_proj", # mlp "q_proj","k_proj","v_proj" # attention], task_type="CAUSAL_LM", bias="none")</p>

<p>dataset = load_dataset("Abirate/english_quotes", split="train")</p>
<p>training_args = SFTConfig(output_dir="./results", num_train_epochs=2, per_device_train_batch_size=4, logging_dir='./logs', logging_steps=10, learning_rate=1e-5, dataset_text_field="quote")</p>
<p>trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args, peft_config=lora_config, train_dataset=dataset)</p>
<p>trainer.train()</p>