Jamba: Model Bahasa Hibrida AI21 Labs yang Menggabungkan Transformer dan Mamba

Model bahasa telah mengalami kemajuan pesat, dengan arsitektur berbasis Transformer memimpin dalam pengolahan bahasa alami. Namun, ketika model skala, tantangan penanganan konteks panjang, efisiensi memori, dan throughput menjadi lebih menonjol.

AI21 Labs telah memperkenalkan solusi baru dengan Jamba, model bahasa besar (LLM) yang menggabungkan kekuatan arsitektur Transformer dan Mamba dalam kerangka hibrida. Artikel ini menjelaskan arsitektur Jamba, kinerja, dan potensi aplikasinya.

Ikhtisar Jamba

Jamba adalah model bahasa besar hibrida yang dikembangkan oleh AI21 Labs, memanfaatkan kombinasi lapisan Transformer dan Mamba, terintegrasi dengan Mixture-of-Experts (MoE) modul. Arsitektur ini memungkinkan Jamba untuk menyeimbangkan penggunaan memori, throughput, dan kinerja, membuatnya menjadi alat yang kuat untuk berbagai tugas NLP. Model ini dirancang untuk muat dalam satu GPU 80GB, menawarkan throughput tinggi dan jejak memori kecil sambil mempertahankan kinerja tingkat tinggi pada berbagai benchmark.

Arsitektur Jamba

Arsitektur Jamba adalah fondasi kemampuannya. Ini dibangun pada desain hibrida novel yang menyelipkan lapisan Transformer dengan lapisan Mamba, mengintegrasikan modul MoE untuk meningkatkan kapasitas model tanpa meningkatkan biaya komputasi secara signifikan.

1. Lapisan Transformer

Arsitektur Transformer telah menjadi standar untuk model bahasa besar modern karena kemampuannya untuk menangani pemrosesan paralel secara efisien dan menangkap ketergantungan jangka panjang dalam teks. Namun, kinerjanya sering terbatas oleh persyaratan memori dan komputasi yang tinggi, terutama saat memproses konteks panjang. Jamba menangani keterbatasan ini dengan mengintegrasikan lapisan Mamba, yang akan kita jelajahi selanjutnya.

2. Lapisan Mamba

Mamba adalah model ruang negara (SSM) terbaru yang dirancang untuk menangani hubungan jarak jauh dalam urutan dengan lebih efisien daripada RNN tradisional atau bahkan Transformer. Lapisan Mamba khususnya efektif dalam mengurangi jejak memori yang terkait dengan penyimpanan cache kunci-nilai (KV) dalam Transformer. Dengan menyelipkan lapisan Mamba dengan lapisan Transformer, Jamba mengurangi penggunaan memori secara keseluruhan sambil mempertahankan kinerja tinggi, terutama dalam tugas yang memerlukan penanganan konteks panjang.

3. Mixture-of-Experts (MoE) Modul

Modul MoE dalam Jamba memperkenalkan pendekatan fleksibel untuk menskalakan kapasitas model. MoE memungkinkan model untuk meningkatkan jumlah parameter yang tersedia tanpa meningkatkan parameter aktif secara proporsional selama inferensi. Dalam Jamba, MoE diterapkan pada beberapa lapisan MLP, dengan mekanisme router yang memilih ahli teratas untuk diaktifkan untuk setiap token. Aktivasi selektif ini memungkinkan Jamba untuk mempertahankan efisiensi tinggi sambil menangani tugas yang kompleks.

Gambar di bawah ini menunjukkan fungsionalitas dari kepala induksi dalam model Perhatian-Mamba hibrida, fitur kunci dari Jamba. Dalam contoh ini, kepala perhatian bertanggung jawab untuk memprediksi label seperti “Positif” atau “Negatif” dalam tugas analisis sentimen. Kata-kata yang disoroti menunjukkan bagaimana perhatian model sangat fokus pada token label dari contoh few-shot, terutama pada saat kritis sebelum memprediksi label akhir. Mekanisme perhatian ini memainkan peran kritis dalam kemampuan model untuk melakukan pembelajaran dalam konteks, di mana model harus menyimpulkan label yang tepat berdasarkan konteks yang diberikan dan contoh few-shot.

Peningkatan kinerja yang ditawarkan oleh integrasi Mixture-of-Experts (MoE) dengan arsitektur Perhatian-Mamba hibrida disoroti dalam Tabel. Dengan menggunakan MoE, Jamba meningkatkan kapasitasnya tanpa meningkatkan biaya komputasi secara proporsional. Ini terutama jelas dalam peningkatan kinerja yang signifikan di berbagai benchmark seperti HellaSwag, WinoGrande, dan Pertanyaan Alami (NQ). Model dengan MoE tidak hanya mencapai akurasi yang lebih tinggi (misalnya, 66,0% pada WinoGrande dibandingkan dengan 62,5% tanpa MoE) tetapi juga menunjukkan log-probabilitas yang ditingkatkan di berbagai domain (misalnya, -0,534 pada C4).

Fitur Arsitektur Kunci

• Komposisi Lapisan: Arsitektur Jamba terdiri dari blok yang menggabungkan Mamba dan lapisan Transformer dalam rasio tertentu (misalnya, 1:7, yang berarti satu lapisan Transformer untuk setiap tujuh lapisan Mamba). Rasio ini dioptimalkan untuk kinerja dan efisiensi yang optimal.
• Integrasi MoE: Lapisan MoE diterapkan setiap beberapa lapisan, dengan 16 ahli yang tersedia dan dua ahli teratas diaktifkan per token. Konfigurasi ini memungkinkan Jamba untuk menskalakan secara efektif sambil mengelola trade-off antara penggunaan memori dan efisiensi komputasi.
• Normalisasi dan Stabilitas: Untuk memastikan stabilitas selama pelatihan, Jamba mengintegrasikan RMSNorm dalam lapisan Mamba, yang membantu mengurangi masalah seperti lonjakan aktivasi besar yang dapat terjadi pada skala.

Kinerja dan Benchmarking Jamba

Jamba telah diuji secara ketat melawan berbagai benchmark, menunjukkan kinerja yang kompetitif di seluruh papan. Bagian-bagian berikut menyoroti beberapa benchmark kunci di mana Jamba telah unggul, menunjukkan kekuatannya dalam tugas NLP umum dan skenario konteks panjang.

1. Benchmark NLP Umum

Jamba telah dievaluasi pada beberapa benchmark akademis, termasuk:

• HellaSwag (10-shot): Tugas penalaran umum di mana Jamba mencapai skor kinerja 87,1%, mengungguli banyak model yang bersaing.
• WinoGrande (5-shot): Tugas penalaran lain di mana Jamba mencapai skor 82,5%, menunjukkan kemampuannya untuk menangani penalaran linguistik yang kompleks.
• ARC-Challenge (25-shot): Jamba menunjukkan kinerja yang kuat dengan skor 64,4%, mencerminkan kemampuannya untuk mengelola pertanyaan pilihan ganda yang menantang.

Dalam benchmark agregat seperti MMLU (5-shot), Jamba mencapai skor 67,4%, menunjukkan kekuatannya di berbagai tugas.

2. Evaluasi Konteks Panjang

Salah satu fitur unggulan Jamba adalah kemampuannya untuk menangani konteks yang sangat panjang. Model ini mendukung panjang konteks hingga 256K token, yang terpanjang di antara model yang tersedia secara publik. Kemampuan ini diuji menggunakan benchmark Needle-in-a-Haystack, di mana Jamba menunjukkan akurasi pengambilan yang luar biasa di berbagai panjang konteks, termasuk hingga 256K token.

3. Throughput dan Efisiensi

Arsitektur hibrida Jamba secara signifikan meningkatkan throughput, terutama dengan urutan panjang.

Dalam tes yang membandingkan throughput (token per detik) di berbagai model, Jamba secara konsisten mengungguli rekan-rekannya, terutama dalam skenario yang melibatkan ukuran batch besar dan konteks panjang. Misalnya, dengan konteks 128K token, Jamba mencapai tiga kali throughput Mixtral, model yang setara.

Menggunakan Jamba: Python

Untuk pengembang dan peneliti yang ingin bereksperimen dengan Jamba, AI21 Labs telah menyediakan model pada platform seperti Hugging Face, membuatnya dapat diakses untuk berbagai aplikasi. Cuplikan kode berikut menunjukkan cara memuat dan menghasilkan teks menggunakan Jamba:

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

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

<p>input_ids = tokenizer(&quot;In the recent Super Bowl LVIII,&quot;, return_tensors=&#039;pt&#039;).to(model.device)[&quot;input_ids&quot;]</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(&quot;ai21labs/Jamba-v0.1&quot;)
model = AutoModelForCausalLM.from_pretrained(
&quot;ai21labs/Jamba-v0.1&quot;, device_map=&#039;auto&#039;, torch_dtype=torch.bfloat16)</p>

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

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