جامبا: نموذج لغة هجين جديد من مختبرات AI21

شهدت نماذج اللغة تقدمًا سريعًا، حيث تقودها هياكل Transformer في معالجة اللغة الطبيعية. ومع ذلك، مع نمو النماذج، أصبحت تحديات التعامل مع السياقات الطويلة، وكفاءة الذاكرة، والإنتاجية أكثر وضوحًا.

مختبرات AI21 قد قدمت حلًا جديدًا مع جامبا، وهو نموذج لغة كبير متقدم (LLM) يجمع بين نقاط القوة للهياكل Transformer و Mamba في إطار هجين. هذا المقال يُفصل هيكل جامبا وأدائه والتطبيقات المحتملة.

نظرة عامة على جامبا

جامبا هو نموذج لغة كبير هجين تم تطويره بواسطة مختبرات AI21، يستفيد من دمج طبقات Transformer و Mamba، متكامل مع وحدة Mixture-of-Experts (MoE). يسمح هذا الهيكل لجامبا بالتوازن بين استخدام الذاكرة والإنتاجية والأداء، مما يجعله أداة قوية لمجموعة واسعة من مهام معالجة اللغة الطبيعية. تم تصميم النموذج ليناسب داخل جهاز GPU وحيد بسعة 80GB، مع تقديم إنتاجية عالية وذاكرة صغيرة مع الحفاظ على أداء متقدم على مختلف البنود.

هيكل جامبا

هيكل جامبا هو ركن أساسي لقدراته. إنه مبني على تصميم هجين جديد يتناوب بين طبقات Transformer و Mamba، مع دمج وحدات MoE لتعزيز سعة النموذج دون زيادة الطلبات الحسابية بشكل كبير.

1. طبقات Transformer

هيكل Transformer أصبح المعيار للنماذج الحديثة بسبب قدرته على التعامل مع المعالجة الموازية بكفاءة واكتشاف الاعتماديات على المدى الطويل في النص. ومع ذلك، غالبًا ما يكون أداؤه محدودًا بسبب متطلبات الذاكرة والحوسبة العالية، خاصة عند معالجة سياقات طويلة. جامبا يعالج هذه القيود من خلال دمج طبقات Mamba، والتي سنناقشها فيما يلي.

2. طبقات Mamba

Mamba هو نموذج فضاء حالة (SSM) حديث مصمم للتعامل مع العلاقات على المدى الطويل في التسلسلات بكفاءة أكبر من النماذج التقليدية مثل RNNs أو حتى Transformers. طبقات Mamba فعالة بشكل خاص في تقليل بصمة الذاكرة المرتبطة بتخزين ذاكرة KV في Transformers. من خلال 交ل طبقات Mamba مع طبقات Transformer، يقلل جامبا من استخدام الذاكرة الإجمالي مع الحفاظ على الأداء العالي، خاصة في المهام التي تتطلب معالجة سياق طويل.

3. وحدات Mixture-of-Experts (MoE)

وحدة MoE في جامبا تقدم نهجًا مرنًا لتوسيع سعة النموذج. MoE يسمح للنموذج بزيادة عدد المعاملات المتاحة دون زيادة المعاملات النشطة بشكل متناسب خلال الاستدلال. في جامبا، يتم تطبيق MoE على بعض طبقات MLP، مع آلية الموجه التي تختار الخبراء الأفضل لتنشيط لكل رمز. يسمح هذا التنشيط الانتقائي لجامبا بالحفاظ على كفاءة عالية مع التعامل مع المهام المعقدة.

الصورة التالية توضح وظيفة رأس الاستدلال في نموذج الانتباه-مامبا الهجين، وهو ميزة رئيسية في جامبا. في هذا المثال، يكون رأس الانتباه مسؤولاً عن التنبؤ بالملabelات مثل “إيجابي” أو “سالب” في استجابة لمهام تحليل المشاعر. الكلمات المhighlight توضح كيف يركز انتباه النموذج بقوة على رموز الملABEL من الأمثلة القليلة، خاصة في اللحظة الحاسمة قبل التنبؤ بالملABEL النهائي. يلعب آلية الانتباه دورًا حاسمًا في قدرة النموذج على أداء التعلم في السياق، حيث يجب على النموذج استنتاج الملABEL المناسب بناءً على السياق والأمثلة القليلة.

تحسينات الأداء التي تقدمها دمج MoE مع هيكل الانتباه-مامبا الهجين موضحة في الجدول. من خلال استخدام MoE، يزيد جامبا من سعته دون زيادة التكاليف الحسابية بشكل متناسب. هذا واضح بشكل خاص في الارتفاع الكبير في الأداء عبر مختلف البنود مثل HellaSwag و WinoGrande و Natural Questions (NQ). النموذج مع MoE لا يصل فقط إلى دقة أعلى (مثل 66.0% على WinoGrande مقارنةً بـ 62.5% بدون MoE) ولكن أيضًا يظهر تحسينات في لوغاريتمات الاحتمالات عبر مجالات مختلفة (مثل -0.534 على C4).

الميزات الهيكلية الرئيسية

• تركيب الطبقات: يتكون هيكل جامبا من كتل تجمع بين طبقات Mamba و Transformer بنسبة معينة (مثل 1:7، مما يعني طبقة Transformer واحدة لكل سبع طبقات Mamba). يتم ضبط هذه النسبة لتحقيق الأداء والكفاءة الأمثل.
• دمج MoE: يتم تطبيق طبقات MoE كل بضع طبقات، مع 16 خبيرًا متاحًا وتنشيط الخبراء الأفضل 2 لكل رمز. يسمح هذا التكوين لجامبا بالتوسع بشكل فعال مع إدارة التبادل بين استخدام الذاكرة والكفاءة الحسابية.
• التعديل والاستقرار: لضمان الاستقرار أثناء التدريب، يدمج جامبا RMSNorm في طبقات Mamba، مما يساعد في التغلب على مشاكل مثل انفجارات التنشيط الكبيرة التي يمكن أن تحدث عند النطاق.

أداء جامبا وتقييمه

تم اختبار جامبا بشكل شامل ضد مجموعة واسعة من البنود، مما يدل على أداء تنافسي عبر اللوحة. الأقسام التالية تسلط الضوء على بعض البنود الرئيسية التي تميزت فيها جامبا، مما يظهر نقاط قوتها في المهام العامة لمعالجة اللغة الطبيعية و السيناريوهات ذات السياق الطويل.

1. بنود معالجة اللغة الطبيعية الشائعة

تم تقييم جامبا على عدة بنود أكاديمية، بما في ذلك:

• HellaSwag (10-shot): مهمة تفكير بالعقل السليم حيث حقق جامبا درجة أداء 87.1%， متجاوزًا العديد من النماذج المنافسة.
• WinoGrande (5-shot): مهمة تفكير أخرى حيث سجل جامبا 82.5%， مما يظهر مرة أخرى قدرته على التعامل مع التفكير اللغوي المعقد.
• ARC-Challenge (25-shot): أظهر جامبا أداءً قويًا بنتيجة 64.4%， مما يعكس قدرته على إدارة أسئلة متعددة الخيارات الصعبة.

في بنود التجميع مثل MMLU (5-shot)، حقق جامبا درجة 67.4%， مما يدل على متانته عبر مهام متنوعة.

2. تقييمات السياق الطويل

ميزة جامبا البارزة هي قدرته على التعامل مع سياقات طويلة بشكل استثنائي. يدعم النموذج طول سياق يصل إلى 256K رمز، وهو الأطول بين النماذج المتاحة للجمهور. تم اختبار هذه القدرة باستخدام بنود Needle-in-a-Haystack، حيث أظهر جامبا دقة استرجاع استثنائية عبر أطوال سياق مختلفة، بما في ذلك 256K رمز.

3. الإنتاجية والكفاءة

يحسن هيكل جامبا الهجين الإنتاجية بشكل كبير، خاصة مع التسلسلات الطويلة.

في الاختبارات التي تقارن الإنتاجية (رموز لكل ثانية) عبر نماذج مختلفة، كان أداء جامبا دائمًا أفضل من أقرانه، خاصة في السيناريوهات التي تتضمن أحجام دفعة كبيرة وسياقات طويلة. على سبيل المثال، مع سياق 128K رمز، حقق جامبا ثلاثة أضعاف الإنتاجية من Mixtral، نموذج قابل للمقارنة.

استخدام جامبا: بايثون

لمطوري الباحثين الذين يرغبون في تجربة جامبا، قدمت مختبرات AI21 النموذج على منصات مثل Hugging Face، مما يجعله متاحًا لتطبيقات واسعة النطاق. الشفرة التالية تظهر كيفية تحميل جامبا وتوليد نص باستخدامه:

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