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

언어 모델은 급속한 발전을 이루었으며, Transformer 기반 아키텍처가 자연어 처리 분야를 선도하고 있습니다. 그러나 모델이 확장됨에 따라 긴 컨텍스트, 메모리 효율성 및 처리량 처리의 과제가 더욱 두드러졌습니다.

AI21 랩스 Transformer와 LLM(대규모 언어 모델)의 장점을 결합한 최첨단 솔루션인 Jamba를 출시했습니다. 맘바 아키텍처 하이브리드 프레임워크에서. 이 기사에서는 Jamba의 세부 사항을 탐구하고 아키텍처, 성능 및 잠재적인 응용 프로그램을 살펴봅니다.

잠바 개요

잠바 AI21 Labs에서 개발한 Transformer 계층과 Mamba 계층을 결합하여 개발한 하이브리드 대규모 언어 모델입니다. 전문가 혼합(MoE) 모듈. 이 아키텍처를 통해 Jamba는 메모리 사용, 처리량 및 성능의 균형을 유지할 수 있어 광범위한 NLP 작업에 강력한 도구가 됩니다. 이 모델은 단일 80GB GPU에 맞게 설계되어 다양한 벤치마크에서 최첨단 성능을 유지하면서도 높은 처리량과 작은 메모리 공간을 제공합니다.

잠바의 건축

Jamba의 아키텍처는 그 역량의 초석입니다. Transformer 레이어와 Mamba 레이어를 섞어서 MoE 모듈을 통합하여 계산 수요를 크게 증가시키지 않고도 모델의 용량을 향상시키는 새로운 하이브리드 디자인을 기반으로 구축되었습니다.

1. 변압기 층

Transformer 아키텍처는 병렬 처리를 효율적으로 처리하고 텍스트에서 장거리 종속성을 포착하는 능력 덕분에 현대 LLM의 표준이 되었습니다. 그러나 성능은 종종 높은 메모리 및 컴퓨팅 요구 사항으로 인해 제한을 받는데, 특히 긴 컨텍스트를 처리할 때 그렇습니다. Jamba는 Mamba 계층을 통합하여 이러한 제한 사항을 해결하는데, 다음에 살펴보겠습니다.

2. 맘바 레이어

Mamba는 전통적인 RNN 또는 Transformer보다 더 효율적으로 시퀀스에서 장거리 관계를 처리하도록 설계된 최신 상태 공간 모델(SSM)입니다. Mamba 레이어는 Transformer에서 키-값(KV) 캐시를 저장하는 것과 관련된 메모리 사용량을 줄이는 데 특히 효과적입니다. Mamba 레이어를 Transformer 레이어와 인터리빙함으로써 Jamba는 높은 성능을 유지하면서 전체 메모리 사용량을 줄이며, 특히 긴 컨텍스트 처리가 필요한 작업에서 그렇습니다.

3. 전문가 혼합(MoE) 모듈

그만큼 환경부 Jamba의 모듈은 모델 용량을 확장하는 데 유연한 접근 방식을 도입합니다. MoE를 사용하면 추론 중에 활성 매개변수를 비례적으로 늘리지 않고도 모델이 사용 가능한 매개변수의 수를 늘릴 수 있습니다. Jamba에서 MoE는 일부 MLP 계층에 적용되며, 라우터 메커니즘은 각 토큰에 대해 활성화할 최고 전문가를 선택합니다. 이 선택적 활성화를 통해 Jamba는 복잡한 작업을 처리하는 동안 높은 효율성을 유지할 수 있습니다.

아래 이미지는 Jamba의 핵심 기능인 하이브리드 Attention-Mamba 모델에서 유도 헤드의 기능을 보여줍니다. 이 예에서 Attention 헤드는 감정 분석 작업에 대한 응답으로 “긍정적” 또는 “부정적”과 같은 레이블을 예측하는 역할을 합니다. 강조된 단어는 모델의 주의가 최종 레이블을 예측하기 전의 중요한 순간에 특히 few-shot 예제의 레이블 토큰에 얼마나 강하게 집중되는지 보여줍니다. 이 주의 메커니즘은 모델이 주어진 컨텍스트와 few-shot 예제를 기반으로 적절한 레이블을 추론해야 하는 컨텍스트 내 학습을 수행하는 데 중요한 역할을 합니다.

Mixture-of-Experts(MoE)를 Attention-Mamba 하이브리드 아키텍처와 통합하여 제공되는 성능 개선 사항은 표에 강조 표시되어 있습니다. MoE를 사용함으로써 Jamba는 계산 비용을 비례적으로 증가시키지 않고도 용량을 늘립니다. 이는 HellaSwag, WinoGrande, Natural Questions(NQ)와 같은 다양한 벤치마크에서 성능이 크게 향상된 것에서 특히 분명하게 드러납니다. MoE를 사용한 모델은 더 높은 정확도(예: MoE가 없는 경우 62.5%에 비해 WinoGrande에서 66.0%)를 달성할 뿐만 아니라 다양한 도메인에서 향상된 로그 확률(예: C4에서 -0.534)도 보여줍니다.

주요 건축적 특징

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

Jamba의 성능 및 벤치마킹

Jamba는 광범위한 벤치마크에 대해 엄격하게 테스트되었으며, 전반적으로 경쟁력 있는 성과를 보였습니다. 다음 섹션에서는 Jamba가 뛰어난 주요 벤치마크 중 일부를 강조하여 일반적인 NLP 작업과 장기 컨텍스트 시나리오에서 강점을 보여줍니다.

1. 일반적인 NLP 벤치마크

Jamba는 다음을 포함한 여러 학문적 기준에 따라 평가되었습니다.

• 헬라스웨그(10발): Jamba는 상식적 추론 작업에서 87.1%의 성능 점수를 달성하여 많은 경쟁 모델을 앞지르게 되었습니다.
• 위노그란데(5샷): Jamba는 또 다른 추론 작업에서 82.5%의 점수를 받았으며, 이는 다시 한번 복잡한 언어적 추론을 처리할 수 있는 능력을 보여주었습니다.
• ARC 챌린지(25샷): Jamba는 64.4%라는 높은 점수로 강력한 성과를 보여주었으며, 이는 까다로운 객관식 문제에 대처하는 능력을 보여준다.

MMLU(5샷)와 같은 집계 벤치마크에서 Jamba는 67.4%의 점수를 달성하여 다양한 작업에서 견고함을 나타냈습니다.

2. 장기적 맥락 평가

Jamba의 두드러진 특징 중 하나는 매우 긴 컨텍스트를 처리할 수 있는 능력입니다. 이 모델은 최대 256K 토큰의 컨텍스트 길이를 지원하며, 이는 공개적으로 사용 가능한 모델 중 가장 깁니다. 이 기능은 Needle-in-a-Haystack 벤치마크를 사용하여 테스트되었으며, Jamba는 최대 256K 토큰을 포함하여 다양한 컨텍스트 길이에서 뛰어난 검색 정확도를 보였습니다.

3. 처리량 및 효율성

Jamba의 하이브리드 아키텍처는 특히 긴 시퀀스의 처리량을 크게 향상시킵니다.

다양한 모델에서 처리량(초당 토큰)을 비교하는 테스트에서 Jamba는 특히 대규모 배치 크기와 긴 컨텍스트를 포함하는 시나리오에서 경쟁사를 지속적으로 앞지르는 성과를 보였습니다. 예를 들어, 128K 토큰의 컨텍스트에서 Jamba는 비슷한 모델인 Mixtral보다 3배 높은 처리량을 달성했습니다.

Jamba 사용: Python

Jamba를 실험하고 싶어하는 개발자와 연구자를 위해 AI21 Labs는 Hugging Face와 같은 플랫폼에서 모델을 제공하여 광범위한 애플리케이션에서 사용할 수 있도록 했습니다. 다음 코드 조각은 Jamba를 사용하여 텍스트를 로드하고 생성하는 방법을 보여줍니다.

from transformers import AutoModelForCausalLM, AutoTokenizermodel = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1")tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")input_ids = tokenizer("In the recent Super Bowl LVIII,", return_tensors='pt').to(model.device)["input_ids"]outputs = model.generate(input_ids, max_new_tokens=216)print(tokenizer.batch_decode(outputs))

import torchfrom datasets import load_datasetfrom trl import SFTTrainer, SFTConfigfrom peft import LoraConfigfrom transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArgumentstokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)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")dataset = load_dataset("Abirate/english_quotes", split="train")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")trainer = SFTTrainer(model=model, tokenizer=tokenizer, args=training_args,peft_config=lora_config, train_dataset=dataset,)trainer.train()

게시물 Jamba: AI21 Labs의 새로운 하이브리드 Transformer-Mamba 언어 모델 처음 등장 유나이트.AI.