Jamba: โมเดลภาษาไฮบริด Transformer-Mamba ของ AI21 Labs

โมเดลภาษาได้พัฒนาไปอย่างรวดเร็ว โดยมีสถาปัตยกรรม Transformer นำหน้าในด้านการประมวลผลภาษาธรรมชาติ อย่างไรก็ตาม เมื่อโมเดลมีขนาดใหญ่ขึ้น ปัญหาในการจัดการบริบทที่ยาว ความสามารถในการใช้หน่วยความจำ และการประมวลผลได้กลายเป็นเรื่องที่ชัดเจนขึ้น

AI21 Labs ได้แนะนำวิธีแก้ปัญหาใหม่ด้วย Jamba โมเดลภาษาขนาดใหญ่ (LLM) ที่รวมจุดแข็งของสถาปัตยกรรม Transformer และ Mamba ในเฟรมเวิร์กไฮบริด โมเดลนี้มีรายละเอียดเกี่ยวกับสถาปัตยกรรม ประสิทธิภาพ และการประยุกต์ใช้งาน

ภาพรวมของ Jamba

Jamba เป็นโมเดลภาษาขนาดใหญ่ไฮบริดที่พัฒนาโดย AI21 Labs โดยใช้การผสมผสานระหว่างเลเยอร์ Transformer และ Mamba รวมถึงโมดูล Mixture-of-Experts (MoE) สถาปัตยกรรมนี้ช่วยให้ Jamba สามารถสร้างสมดุลระหว่างการใช้หน่วยความจำ การประมวลผล และประสิทธิภาพ ทำให้เป็นเครื่องมือที่ทรงพลังสำหรับงานประมวลผลภาษาธรรมชาติหลายอย่าง โมเดลนี้ได้รับการออกแบบให้สามารถทำงานได้ภายใน GPU 80GB เพียงตัวเดียว โดยมีการประมวลผลที่รวดเร็วและใช้หน่วยความจำน้อย，同时ยังคงประสิทธิภาพที่ดีเยี่ยมในหลายๆ บンチมาร์ก

สถาปัตยกรรมของ Jamba

สถาปัตยกรรมของ Jamba เป็นรากฐานของความสามารถของมัน โมเดลนี้สร้างขึ้นจากดีไซน์ไฮบริดใหม่ที่สลับเลเยอร์ Transformer กับเลเยอร์ Mamba และรวมโมดูล MoE เพื่อเพิ่มความสามารถของโมเดลโดยไม่เพิ่มความต้องการการประมวลผลมากเกินไป

1. เลเยอร์ Transformer

สถาปัตยกรรม Transformer ได้กลายเป็นมาตรฐานสำหรับโมเดลภาษาขนาดใหญ่สมัยใหม่ เนื่องจากความสามารถในการประมวลผลข้อมูลแบบขนานและจับข้อมูลที่มีความสัมพันธ์ระยะไกลในข้อความ อย่างไรก็ตาม ประสิทธิภาพของมันบ่อยครั้งถูกจำกัดโดยความต้องการหน่วยความจำและคอมพิวเตอร์ที่สูง โดยเฉพาะอย่างยิ่งเมื่อประมวลผลบริบทที่ยาว Jamba ได้กล่าวถึงข้อจำกัดเหล่านี้โดยการรวมเลเยอร์ Mamba ซึ่งเราจะพูดถึงต่อไป

2. เลเยอร์ Mamba

Mamba เป็นโมเดลสถานะพื้นที่ (SSM) ที่ออกแบบมาเพื่อจัดการความสัมพันธ์ระยะไกลในลำดับข้อมูลได้อย่างมีประสิทธิภาพมากกว่า RNN หรือแม้แต่ Transformer เลเยอร์ Mamba มีประสิทธิภาพสูงในการลดการใช้หน่วยความจำที่เกี่ยวข้องกับการเก็บแคช key-value (KV) ใน Transformer โดยการสลับเลเยอร์ Mamba กับเลเยอร์ Transformer Jamba ลดการใช้หน่วยความจำโดยรวมขณะยังคงรักษาความสามารถในการประมวลผลที่สูง โดยเฉพาะอย่างยิ่งในงานที่ต้องการการจัดการบริบทที่ยาว

3. โมดูล Mixture-of-Experts (MoE)

โมดูล MoE ใน Jamba นำเสนอวิธีการที่ยืดหยุ่นในการปรับขนาดความสามารถของโมเดล MoE ช่วยให้โมเดลสามารถเพิ่มจำนวนพารามิเตอร์ที่มีอยู่โดยไม่เพิ่มพารามิเตอร์ที่ใช้งานอยู่ในระหว่างการอนุมาน ใน Jamba MoE ถูกใช้กับบางเลเยอร์ โดยมีกลไก router ที่เลือกผู้เชี่ยวชาญที่ดีที่สุดเพื่อเปิดใช้งานสำหรับแต่ละโทเค็น การเปิดใช้งานแบบเลือกนี้ช่วยให้ Jamba รักษาความสามารถในการประมวลผลที่สูงในขณะเดียวกันก็จัดการกับงานที่ซับซ้อน

รูปด้านล่างแสดงถึงการทำงานของหัวอินดักชันในโมเดล Attention-Mamba ไฮบริด ซึ่งเป็นคุณลักษณะสำคัญของ Jamba ในตัวอย่างนี้ หัวอินดักชันรับผิดชอบในการทำนายลेबェล เช่น “Positive” หรือ “Negative” ในงานวิเคราะห์ความรู้สึก การเน้นของโมเดลมุ่งเน้นไปที่โทเค็นลेबェลจากตัวอย่าง few-shot โดยเฉพาะอย่างยิ่งก่อนทำนายลेबェลสุดท้าย กลไกการเน้นนี้มีบทบาทสำคัญในการเรียนรู้ในบริบทของโมเดล ซึ่งโมเดลต้องอนุมานลेबェลที่เหมาะสมตามบริบทและตัวอย่าง few-shot ที่ให้มา

การปรับปรุงประสิทธิภาพที่ได้รับจากการรวม MoE กับสถาปัตยกรรม Attention-Mamba ไฮบริด ถูกเน้นในตาราง โดยการใช้ MoE Jamba เพิ่มความสามารถโดยไม่เพิ่มค่าใช้จ่ายในการประมวลผลมากเกินไป ซึ่งเห็นได้ชัดในผลลัพธ์ที่ดีขึ้นอย่างมีนัยสำคัญในบンチมาร์กต่างๆ เช่น HellaSwag, WinoGrande และ Natural Questions (NQ) โมเดลที่มี MoE ไม่เพียงแต่บรรลุความแม่นยำที่สูงขึ้น (เช่น 66.0% ใน WinoGrande เมื่อเทียบกับ 62.5% โดยไม่มี MoE) แต่ยังแสดงถึงความสามารถในการจัดการข้อมูลที่ดีขึ้นในโดเมนต่างๆ (เช่น -0.534 ใน C4)

คุณลักษณะสถาปัตยกรรมหลัก

• การผสมผสานเลเยอร์: สถาปัตยกรรมของ Jamba ประกอบด้วยบล็อกที่ผสมผสานเลเยอร์ Mamba และ Transformer ในอัตราส่วนเฉพาะ (เช่น 1:7 หมายถึงเลเยอร์ Transformer หนึ่งเลเยอร์สำหรับเลเยอร์ Mamba เจ็ดเลเยอร์) อัตราส่วนนี้ถูกปรับให้เหมาะสมสำหรับการทำงานและประสิทธิภาพ
• การรวม MoE: เลเยอร์ MoE ถูกใช้ทุกๆ เลเยอร์บางส่วน โดยมีผู้เชี่ยวชาญ 16 คนและผู้เชี่ยวชาญ 2 คนที่ถูกเปิดใช้งานต่อโทเค็น การกำหนดค่านี้ช่วยให้ Jamba ขยายได้อย่างมีประสิทธิภาพในขณะเดียวกันก็จัดการกับสมดุลระหว่างการใช้หน่วยความจำและประสิทธิภาพการประมวลผล
• การปกติและเสถียรภาพ: เพื่อให้แน่ใจถึงเสถียรภาพระหว่างการฝึกอบรม Jamba รวม RMSNorm ในเลเยอร์ Mamba ซึ่งช่วยลดปัญหาด้านการกระตุ้นที่อาจเกิดขึ้นเมื่อทำงานในขนาดใหญ่

ประสิทธิภาพและบンチมาร์กของ Jamba

Jamba ได้ผ่านการทดสอบอย่างเข้มงวดกับบンチมาร์กหลายตัว โดยแสดงผลลัพธ์ที่แข่งขันได้ทั่วทั้งกระดาน ส่วนต่อไปนี้จะเน้นย้ำบンチมาร์กที่สำคัญที่ Jamba ได้โดดเด่น โดยแสดงถึงจุดแข็งของมันในงานประมวลผลภาษาและงานที่ต้องการบริบทยาว

1. บンチมาร์ก NLP ทั่วไป

Jamba ได้รับการประเมินในบンチมาร์กทางวิชาการหลายตัว รวมถึง:

• HellaSwag (10-shot): งานที่ต้องใช้เหตุผลเชิงสมถรรถ โดย Jamba ได้ผลลัพธ์ 87.1% ซึ่งเหนือกว่าโมเดลที่แข่งขันกันหลายตัว
• WinoGrande (5-shot): งานที่ต้องใช้เหตุผลอีกตัวหนึ่ง โดย Jamba ได้ผลลัพธ์ 82.5% ซึ่งแสดงถึงความสามารถในการจัดการเหตุผลเชิงภาษาที่ซับซ้อน
• ARC-Challenge (25-shot): Jamba ได้ผลลัพธ์ 64.4% ซึ่งสะท้อนถึงความสามารถในการจัดการคำถามแบบหลายตัวเลือกที่ท้าทาย

ในบンチมาร์กที่รวมหลายงาน เช่น MMLU (5-shot) 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:

<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()

การนำไปใช้และการผสานรวม

AI21 Labs ได้ทำให้ Jamba มีให้ใช้บนแพลตฟอร์มและตัวเลือกการนำไปใช้หลายอย่าง:

1. Cloud Platforms:
   • มีให้ใช้บนผู้ให้บริการคลาวด์หลักๆ รวมถึง Google Cloud Vertex AI, Microsoft Azure และ NVIDIA NIM (NVDA )
   • จะมาถึงเร็วๆ นี้บน Amazon Bedrock, Databricks Marketplace และ Snowflake Cortex
2. AI Development Frameworks:
   • การผสานกับเฟรมเวิร์กยอดนิยม เช่น LangChain และ LlamaIndex (กำลังจะมาถึง)
3. AI21 Studio:
   • การเข้าถึงโดยตรงผ่านแพลตฟอร์มพัฒนาของ AI21
4. Hugging Face:
   • โมเดลมีให้ดาวน์โหลดและทดลองใช้
5. On-Premises Deployment:
   • ตัวเลือกสำหรับการใช้งานในองค์กรที่ต้องการความปลอดภัยหรือการปฏิบัติตามกฎระเบียบพิเศษ
6. Custom Solutions:
   • AI21 เสนอบริการปรับแต่งและฝึกอบรมโมเดลแบบกำหนดเองสำหรับลูกค้าองค์กร

คุณลักษณะที่เป็นมิตรกับนักพัฒนา

โมเดล Jamba มีคุณลักษณะที่น่าสนใจหลายอย่างซึ่งทำให้มันเป็นตัวเลือกที่ดีสำหรับนักพัฒนา:

1. การเรียกฟังก์ชัน: การผสานรวมเครื่องมือและ API ภายนอกเข้ากับเวิร์กโฟลว์ AI ได้อย่างง่ายดาย
2. การผลิต JSON ที่มีโครงสร้าง: การสร้างโครงสร้างข้อมูลที่สะอาดและสามารถแปลได้โดยตรงจากข้อความธรรมชาติ
3. การย่อยเอกสาร: การประมวลผลและทำความเข้าใจโครงสร้างเอกสารที่ซับซ้อนได้อย่างมีประสิทธิภาพ
4. การเพิ่มประสิทธิภาพ RAG: คุณลักษณะที่ถูกสร้างขึ้นเพื่อเพิ่มประสิทธิภาพของการสร้างข้อความที่ช่วยโดยการค้นหา

คุณลักษณะเหล่านี้ เมื่อรวมกับความสามารถในการจัดการบริบทยาวและการประมวลผลที่รวดเร็วของโมเดล ทำให้ Jamba เป็นเครื่องมือที่มีประโยชน์สำหรับสถานการณ์การพัฒนาที่หลากหลาย

การพิจารณาทางจริยธรรมและการใช้ AI อย่างรับผิดชอบ

ในขณะที่ความสามารถของ Jamba นั้นน่าประทับใจ แต่ก็จำเป็นต้องใช้มันด้วยมุมมองของ AI ที่รับผิดชอบ AI21 Labs เน้นย้ำจุดสำคัญหลายประการ:

1. ธรรมชาติของโมเดลฐาน: โมเดล Jamba 1.5 เป็นโมเดลฐานที่ไม่มีการปรับแต่งหรือการฝึกอบรมเฉพาะ
2. การขาดกลไกการป้องกัน: โมเดลไม่มีกลไกการป้องกันที่ถูกสร้างขึ้นมา
3. การนำไปใช้อย่างระมัดระวัง: ควรปรับและเพิ่มการป้องกันเพิ่มเติมก่อนใช้ Jamba ในสภาพแวดล้อมการผลิตหรือกับผู้ใช้สุดท้าย
4. การรักษาความเป็นส่วนตัวของข้อมูล: เมื่อใช้การนำไปใช้บนคลาวด์ ควรระวังในการจัดการและปฏิบัติตามข้อกำหนดเกี่ยวกับข้อมูล
5. การรู้จักและจัดการอคติ: เช่นเดียวกับโมเดลภาษาขนาดใหญ่อื่นๆ Jamba อาจสะท้อนอคติที่มีอยู่ในข้อมูลการฝึกอบรม ผู้ใช้ควรตระหนักถึงปัญหานี้และใช้มาตรการเพื่อลดอคติ

โดยการคำนึงถึงปัจจัยเหล่านี้ นักพัฒนาและองค์กรสามารถใช้ Jamba ได้อย่างรับผิดชอบและจริยธรรม

บทใหม่ในพัฒนาการของ AI?

การแนะนำโมเดล Jamba โดย AI21 Labs ถือเป็นก้าวสำคัญในพัฒนาการของโมเดลภาษาขนาดใหญ่ โดยการรวมจุดแข็งของ Transformer และโมเดลสถานะพื้นที่ การบูรณาการเทคนิค Mixture-of-Experts และการขยายขอบเขตของความยาวบริบทและการประมวลผล Jamba เปิดโอกาสใหม่ๆ สำหรับการใช้งาน AI ในหลายอุตสาหกรรม

เมื่อชุมชน AI ต่อเนื่องในการสำรวจและสร้างบนสถาปัตยกรรมนี้ เราสามารถคาดหวังการปรับปรุงที่เพิ่มขึ้นในด้านประสิทธิภาพ การจัดการบริบทยาว และการนำไปใช้ AI ในทางปฏิบัติ โมเดล Jamba ไม่เพียงแต่แสดงถึงโมเดลใหม่ๆ เท่านั้น แต่ยังเป็นตัวแทนของการเปลี่ยนแปลงที่อาจเกิดขึ้นในด้านการออกแบบและการใช้งานระบบ AI ขนาดใหญ่