ملخص تقني: إعادة النظر في بارامترية طبقة الترانسفورمر عبر تقليل الطاقة السببي

بيان المشكلة

تُعد بنيات الترانسفورمر (Transformer)، وتحديداً الجمع بين الانتباه متعدد الرؤوس (MHA) ومركبات MLP ذات البوابات (gated MLPs)، اللبنات الأساسية لنماذج اللغات الكبيرة (LLMs) الحديثة. ومع ذلك، فإن الخيارات المحددة المتعلقة ببارامترية الطبقات — مثل ربط الأوزان (weight tying)، وهياكل الإسقاط (projection structures)، وقواعد التحديث — مدفوعة إلى حد كبير بالدراسات التجريبية بدلاً من المبادئ الحسابية الصريحة. وبينما قدمت النماذج القائمة على الطاقة (EBMs) منظوراً حيث تُعتبر الحوسبة عملية بحث عن حالات منخفضة الطاقة (مثل شبكات هوبفيلد)، إلا أن الأعمال الحالية لم تربط بشكل كامل هذه التفسيرات القائمة على الطاقة ببارامترات طبقات الترانسفورمر القياسية، ولم تستكشف بشكل منهجي كيف يمكن لهذا المنظور أن يوجه تصميم متغيرات الطبقات الأكثر كفاءة أو فعالية.

المنهجية: تقليل الطاقة السببي (CEM)

يقدم المؤلفون تقليل الطاقة السببي (Causal Energy Minimization - CEM)، وهو إطار عمل يعيد صياغة طبقات الترانسفورمر كخطوات تحسين على دالات طاقة شرطية. الفرضية الأساسية هي أن الطبقة تقوم برسم خريطة من تسلسل مدخل $h_{1:J}$ إلى مخرج $h'_{1:J}$ عبر معاملة كل موضع $i$ كمتغير تحسين $x_i$ (يتم تهيئته عند $h_i$) يتم تحديثه عبر إجراء $\mathcal{A}$ لتقليل طاقة شرطية $\epsilon(x_i | h_{1:i})$ تعتمد على التاريخ السببي.

يشتق البحث مكونات الترانسفورمر القياسية من هذا الإطار ويقترح توسعات:

1. اشتقاق الانتباه المرتبط بالأوزان (Weight-Tied Attention):
   من خلال تعريف طاقة تفاعل بناءً على صيغة (log-sum-exp) (مشابهة لشبكات هوبفيلد الحديثة)، يوضح المؤلفون أن خطوة واحدة من التدرج النزولي (gradient descent) تستعيد تحديث MHA القياسي. ومن المهم أن هذا الاشتقاق يستلزم طبيعياً قيوداً محددة لتقاسم الأوزان:

   • يتم ربط إسقاط المفتاح ($W^K$) بإسقاط القيمة ($W^V$).
   • يتم ربط إسقاط الاستعلام ($W^Q$) بإسقاط المخرج ($W^O$).
   • ترتبط معلمة درجة الحرارة $\tau$ ببعد الرأس (head dimension).

2. اشتقاق الـ MLPs ذات البوابات المرتبطة بالأوزان:
   من خلال تعريف طاقة عنصرية تخصص الطاقة بشكل مستقل لكل متجه سمات الرمز (token)، يستعيد التحديث التدرجي هيكل الـ gated MLPs. يشير هذا التفسير إلى أن مصفوفات الإسقاط "للأعلى" (up) و"للأسفل" (down) في الـ MLP يجب أن تكون مرتبطة (تحديداً، منقول مصفوفة الإسقاط للأسفل يساوي إسقاط الأعلى).

3. التوسعات المقترحة:
   بناءً على منظور تحسين الطاقة، يقترح المؤلفون ثلاث تحسينات معمارية محددة:

   • بارامترية القطر زائد الرتبة المنخفضة (D+LR Parameterization): بدلاً من الإسقاطات منخفضة الرتبة البحتة لمصفوفة التفاعل $A_k$، يقترحون $A_k = \text{diag}(d_k) + W^{Q\top}_k W^K_k$. يلتقط هذا التفاعل الذي قد تغفله المصفوفات منخفضة الرتبة مع بقائه فعالاً حسابياً. كما يقترحون مشاركة الحد القطري عبر الرؤوس لتقليل عدد البارامترات.
   • الممهدات خفيفة الوزن المتعلمة (Learned Lightweight Preconditioners): لتحسين التقارب بما يتجاوز خطوة التدرج الواحدة، يقدمون مصفوفات ممهد $P$ ذات هيكل D+LR ($P = \text{diag}(\text{softplus}(d)) + UV^\top + VU^\top$). تعمل هذه الممهدات كبدائل متعلمة لمعلومات الانحناء (curvature)، حيث تعيد قياس التدرجات دون تكلفة الطرق من الدرجة الثانية الكاملة.
   • التحديثات التكرارية داخل الطبقة (Within-Layer Recursive Updates): بدلاً من تحديث واحد، يسمح الإطار بـ $T$ من الخطوات التكرارية داخل طبقة واحدة ($x^{(t+1)}_i = x^{(t)}_i - \eta \nabla \epsilon$). يتيح هذا للطبقة تقليل دالة الطاقة الخاصة بها بشكل أعمق دون إضافة بارامترات جديدة، وهو ما يختلف عن نماذج الترانسفورمر "الحلقية" (looped) التي تعيد استخدام كتل كاملة.

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

• التوحيد النظري: يثبت البحث أن طبقات الترانسفورمر القياسية مع مخططات مشاركة الأوزان المحددة (ربط Q/O و K/V في الانتباه؛ وربط up/down في MLPs) مكافئة رياضياً لخطوات تدرج واحدة على دالات طاقة تفاعل وعنصرية محددة.
• توسيع مجال التصميم: يحدد منظور CEM مجال تصميم جديد لطبقات الترانسفورمر، بما في ذلك تفاعلات D+LR، والممهدات خفيفة الوزن، والتحديثات التكرارية، والتي تستند إلى نظرية التحسين بدلاً من التجربة والخطأ.
• التحقق التجريبي: يقيم المؤلفون طبقات مشتقة من CEM على نطاق متوسط (مئات الملايين من البارامترات) في مهام نمذجة اللغة، مما يثبت أن هذه البارامترات المقيدة يمكن أن تضاهي أو تقترب من نماذج الترانسفورمر القياسية.

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

أُجريت التجارب على مجموعة بيانات SlimPajama باستخدام أحجام نماذج تتراوح من ~86 مليون إلى ~162 مليون بارامتر، تم تدريبها باستخدام ميزانيات توكن محسنة حسابياً (Chinchilla scaling).

• كفاءة البارامترات: حققت نماذج CEM ذات الأوزان المرتبطة (باستخدام حوالي نصف بارامترات MHA وثلثي بارامترات MLP لنماذج Llama) مستويات تنافسية في اختبار الـ perplexity. بالنسبة لـ CEM attention، كان تدهور الأداء ضئيلاً رغم الخفض الكبير في البارامترات.
• تأثير التوسعات:
  • التكرار (Recursion): أدى زيادة عدد الخطوات التكرارية من $T=1$ إلى $T=2$ إلى تحسين الأداء باستمرار لكل من CEM attention و MLPs. والجدير بالذكر أن متغيرات CEM attention التكرارية تفوقت على نماذج Llama الأساسية مع استخدام بارامترات أقل.
  • الممهدات (Preconditioners): قدمت الممهدات المتعلمة مكاسب هامشية فقط مقارنة بالتكرار، مما يشير إلى أن الهيكل التكراري هو المحرك الرئيسي للتحسن.
  • الحدود القطرية (Diagonal Terms): كان تضمين حد قطري في مصفوفة التفاعل (تحديداً قطر مشترك عبر الرؤوس) ضرورياً للأداء الجيد، حيث تفوق على المتغيرات ذات الرتبة المنخفضة البحتة أو تلك التي لا تحتوي على قطر.
• الاستقرار: تم تدريب نماذج الترانسفورمر المكونة بالكامل من طبقات مشتقة من CEM (مع $T=2$ وممهدات D+LR) بشكل مستقر من البداية للنهاية، محققة perplexity مقاربة لنماذج Llama مع بارامترات أقل.
• البيانات الاصطناعية: أكدت التجارب على بيانات Gaussian process الاصطناعية أن التكرار الأعمق ($T=4, T=8$) يحقق تحسينات إضافية في RMSE، رغم أن المكاسب في إعدادات الترانسفورمر الكاملة اقتصرت على $T=2$ تحت إعدادات التدريب القياسية.

الأهمية والادعاءات

يدعي البحث أن CEM يوفر عدسة نظرية مفيدة لفهم بارامترية طبقة الترانسفورمر، حيث يربط بفعالية بين بنيات الترانسفورمر والنماذج القائمة على الطاقة. تكمن الأهمية في:

1. التصميم المنهجي: ينقل تصميم الطبقات من التجربة والخطأ التجريبي إلى إطار عمل قائم على تقليل الطاقة وديناميكيات التحسين.
2. الكفاءة: يثبت أن مشاركة الأوزان والتحديثات التكرارية، المستمدة من هذا الإطار، تسمح بنماذج كفؤة في البارامترات تحافظ على أداء تنافسي.
3. الاتجاهات المستقبلية: يقترح المؤلفون أن هذا المنظور يحفز استكشاف المزيد في:
   • النواة المخصصة (Custom Kernels): تحسين وقت التشغيل للهياكل المحددة لمشاركة الأوزان والتحديثات التكرارية (مثل دمج العمليات/fusing).
   • توسيع الحوسبة وقت الاختبار (Test-Time Scaling): استخدام التكرار داخل الطبقة كمحور إضافي لتوسيع الحوسبة أثناء الاستدلال.
   • التصميم المشترك مع الأجهزة (Hardware Co-Design): تصميم بارامترات الطبقات خصيصاً لبنيات الأجهزة الجديدة.

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