Reinforced Fast Weights with Next-Sequence Prediction

إليك ملخص تقني مفصل لورقة البحث بعنوان "Reinforced Fast Weights with Next-Sequence Prediction" (REFINE).

1. بيان المشكلة

محدودية التنبؤ بالرمز التالي (NTP) في نماذج الأوزان السريعة (Fast Weight Models):

• السياق: يعد نمذجة السياق الطويل أمراً ضرورياً لمهام مثل فهم المستندات والتعلم داخل السياق متعدد العينات (multi-shot in-context learning). وبينما تتفوق نماذج المحولات (Transformers) القائمة على الانتباه في هذا المجال، إلا أن تكاليف الذاكرة والحوسبة فيها تتوسع بشكل تربيعي مع طول السياق.
• البديل: توفر بنيات الأوزان السريعة (مثل DeltaNet و LaCT) حلاً عبر استبدال الانتباه العالمي بمصفوفة ذاكرة ثابتة الحجم يتم تحديثها ديناميكياً أثناء معالجة الرموز. يضمن ذلك عبئاً ثابتاً على الذاكرة بغض النظر عن طول السياق.
• المشكلة الجوهرية: رغم كفاءتها الهيكلية، يتم تدريب نماذج الأوزان السريعة حالياً باستخدام هدف التنبؤ بالرمز التالي (NTP) القياسي (تقليل خسارة الإنتروبيا المتقاطعة رمزاً تلو الآخر).
  • يعمل الـ NTP على تحسين احتمالية الرمز المنفرد الفوري، متجاهلاً التماسك الدلالي لـ التتابع الذي يلي البادئة (prefix).
  • يؤدي هذا إلى تمثيلات دون المستوى، حيث يفشل النموذج في التقاط الاعتمادات طويلة المدى بفعالية لأن إشارة التدريب لا تشجع النموذج على الحفاظ على حالة داخلية متسقة عبر خطوات متعددة.
  • كما أن توليد استمرارات كاملة متعددة الرموز لكل موضع أثناء التدريب أمر مكلف حوسبياً للغاية، مما يجعل الإشراف على مستوى التسلسل (sequence-level supervision) صعب التنفيذ بالمعايية.

2. المنهجية: REFINE

يقترح المؤلفون REFINE (الأوزان السريعة المعززة بالتنبؤ بالتسلسل التالي)، وهو إطار عمل للتعلم التعزيزي (RL) مصمم لتدريب نماذج الأوزان السريعة تحت هدف التنبؤ بالتسلسل التالي (NSP).

المفهوم الجوهري: التنبؤ بالتسلسل التالي (NSP)

بدلاً من تحسين $P(x_{t+1} | x_{\le t})$، يهدف الـ NSP إلى تحسين تسلسل متماسك من الرموز المستقبلية $P(x_{t+1:t+k} | x_{\le t})$. هذا يجبر الأوزات السريعة على تخزين معلومات تدعم التنبؤ الدقيق عبر أفق قدره $k$ من الرموز، وليس فقط الرمز التالي.

إطار عمل REFINE

يصيغ REFINE عملية الـ NSP كمسألة تعلم تعزيزي للتغلب على التكلفة الحوسبية لتوليد استمرارات كاملة لكل رمز. ويتكون من أربع خطوات رئيسية:

1. اختيار الرموز القائم على الإنتروبيا:

   • بدلاً من أخذ عينات من الرموز بشكل موحد، يحسب النموذج إنتروبيا الـ NTP ($H_t$) عند كل موضع.
   • يتم تقسيم تسلسل المدخلات إلى $c$ من الكتل المتتالية.
   • من كل كتلة، يتم أخذ عينة لموضع رمز مستهدف واحد باحتمالية تتناسب مع إنتروبيته (الإنتروبيا العالية تشير إلى مناطق عالية عدم اليقين/التحدي). يركز هذا على الأجزاء المعلوماتية من السياق.

2. توليد الاستمرارية (Rollout Generation):

   • لكل موضع عالي الإنتروبيا $t_i$ تم اختياره، يقوم النموذج بتوليد استمرارية مكونة من $k$ من الرموز $\hat{x}_{t_i+1:t_i+k}$.
   • يستخرج النموذج الحالات الخفية (hidden states) لكل من الاستمرارية المولدة والاستمرارية الحقيقية (ground-truth).

3. تعيين المكافأة:

   • أثناء التدريب وبعد التدريب: يتم حساب مكافأة تشابه جيب التمام ($R_\phi$) بين الحالات الخفية للاستمرارية المتوقعة والاستمرارية الحقيقية. يوفر هذا إشارة دلالية سلسة تكافئ الاستمرارات المعقولة حتى لو لم تتطابق مع الحقيقة حرفياً.
   • التدريب وقت الاختبار (TTT): تُستخدم مكافأة المطابقة التامة الثنائية ($R_{binary}$) لفرض الحفظ الصارم للسياق.
   • المكافأة الهجينة: بالنسبة لمرحلة ما بعد التدريب، يتم استخدام مزيج من كلا المكافأتين لتحقيق التوازن بين التعميم والحفظ.

4. التحسين باستخدام التعلم التعزيزي (GRPO):

   • يتم تحسين النموذج باستخدام تحسين السياسة النسبي للمجموعات (GRPO).
   • دالة الخسارة هي مجموع مرجح لخسارة الـ NSP (المشتقة من المكافآت) وخسارة الـ NTP القياسية لمنع النسيان الكارثي.
   • الهدف هو تعظيم المكافأة المتوقعة على مستوى التسلسل: $J(\theta) = \mathbb{E}[R]$.

القابلية للتطبيق عبر مراحل التدريب

صُمم REFINE ليكون "محايداً للمرحلة" ويُطبق في ثلاث مراحل:

• منتصف التدريب (Mid-training): تعزيز الأوزان السريعة على مجموعات بيانات تشبه مرحلة ما قبل التدريب لتحسين التكيف مع السياق الطويل.
• ما بعد التدريب (Post-training): يتم دمجه عبر التعلم المتداخل (nested learning). يقوم النموذج بتحديث أوزانه السريعة على "المطالبة" (prompt) باستخدام REFINE قبل توليد الاستجابة النهائية، والتي يتم ضبطها لاحقاً عبر SFT القياسي.
• التدريب وقت الاختبار (TTT): تعزيز الأوزان السريلة مباشرة على مطالبة الاستدلال باستخدام إشارات ذاتية الإشراف (دون الحاجة لملصقات خارجية).

3. المساهمات الرئيسية

1. هدف جديد: تقديم هدف التنبؤ بالتسلسل التالي (NSP) خصيصاً لنماذج الأوزان السريعة لمعالجة قيود التغذية الراجعة على مستوى الرمز.
2. إطار عمل للتعلم التعزيزي: اقتراح REFINE، وهو إطار عمل فعال للتعلم التعزيزي يجمع بين اختيار الرموز القائم على الإنتروبيا (لتقليل الحوسبة) مع المكافآت على مستوى التسلسل (لتحسين التماسك الدلالي).
3. التكامل عبر دورة الحياة: إثبات أن هذا النهج فعال عبر كامل دورة حياة النموذج (منتصف التدريب، ما بعد التدريب، والتدريب وقت الاختبار)، وهي مرونة لم تُستكشف سابقاً للأوزان السريعة.

4. النتائج التجريبية

قيم المؤلفون REFINE على نموذجين: LaCT-760M و DeltaNet-1.3B.

• استرجاع السياق الطويل (RULER NIAH):

  • تفوق REFINE بشكل كبير على الضبط الدقيق الخاضع للإشراف (SFT) باستخدام NTP.
  • DeltaNet-1.3B: أظهر تحسناً بنسبة 23.5% في مهام Multi-key NIAH مقارنة بعدم وجود منتصف تدريب، و 8.8% فوق SFT mid-training.
  • LaCT-760M: أظهر تحسناً مستمراً عبر مهام Single, Multi-key, Multi-query, و Multi-value NIAH.

• الأسئلة والأجوبة للسياق الطويل (LongBench & RULER):

  • منتصف التدريب: حسن REFINE متوسط الأداء على LongBench بنسبة 16.9% (لـ LaCT) و 17.0% (لـ DeltaNet) مقارنة بالنماذج المرجعية.
  • ما بعد التدريب: أدى Nested REFINE إلى تحسين درجات SQuADQA بنسبة 17% (لـ LaCT) و 24.1% (لـ DeltaNet) فوق nested SFT.
  • التدريب وقت الاختبار (TTT): تفوق TTT الخاص بـ REFINE باستمرار على SFT TTT، مما أثبت القدرة الفائقة على التكيف مع السياقات الجديدة دون تسميات إضافية.

• دراسات الاستئصال (Ablation Studies):

  • دالة المكافأة: تفوقت مكافآت تشابه جيب التمام ($R_\phi$) على مكافآت المطابقة التامة الثنائية خلال منتصف التدريب، مما يثبت أن التشابه الدلالي أكثر فعالية للتعميم من المطابقة الصارمة.
  • اختيار الرموز: تفوق أخذ العينات الموزون بالإنتروبيا على أخذ العينات الموحد، وأخذ العينات بالإنتروبيا القصوى، والإنتروبيا الدنيا، مما يؤكد أن التركيز على المناطق غير المؤكدة يوفر إشارات تعلم أفضل.
  • طول الاستمرارية ($k$): بلغ الأداء ذروته عند $k=5$؛ حيث أدت الاستمرارات الأطول ($k=7$) إلى عوائد متناقصة بسبب تدهور إشارة المكافأة.

5. الأهمية

• سد الفجوة: نجح REFINE في مواءمة هدف التدريب لنماذج الأوزان السريعة مع قوتها الهيكلية (الذاكرة طويلة المدى)، متجاوزاً نموذج NTP قصير النظر.
• الكفاءة: باستخدام أخذ العينات القائم على الإنتروبيا والتعلم التعزيزي، يحقق REFINE إشرافاً على مستوى التسلسل دون التكلفة الباهظة لتوليد استمرارات كاملة لكل رمز.
• تعدد الاستخدامات: يثبت إطار العمل أن التعلم التعزيزي يمكن تطبيقه بفعالية على نماذج الأوزان السريعة المدربة مسبقاً حتى بدون ضبط تعليمات مسبق، مما يوفر مساراً عملياً لتوسيع قدرات السياق الطويل في البنيات الموفرة للذاكرة.
• حالة الفن (State-of-the-Art): تضع النتائج معياراً جديداً لنماذج الأوزان السريعة، حيث تظهر قدرتها على منافسة أو تجاوز النماذج القائمة على الانتباه في مهام محددة للسياق الطويل عند تدريبها بالهدف الصحيح.

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