MaskPro: Linear-Space Probabilistic Learning for Strict (N:M)-Sparsity on LLMs

ملخص تقني: MaskPro – التعلم الاحتمالي في مساحة خطية للندرة الصارمة بنمط (N:M) للنماذج اللغوية الكبيرة (LLMs)

بيان المشكلة
أدى التوسع السريع في النماذج اللغوية الكبيرة (LLMs) إلى خلق عقبة كبيرة في كفاءة الاستدلال، مما استلزم تقنيات ضغط فعالة. توفر الندرة شبه المهيكلة، وتحديداً نمط (N:M) حيث يتم الاحتفاظ بـ $N$ من العناصر من بين كل $M$ وزناً متتالياً، حلاً صديقاً للأجهزة للتسريع وتقليل الذاكرة. ومع ذلك، فإن تحديد قناع (N:M) الأمثل يمثل مشكلة تحسين توافقية هائلة. تندرج الأساليب الحالية ضمن فئتين بخصائص متباينة:

1. الأساليب القائمة على القواعد: تستخدم عمليات بحث طبقة تلو الأخرى بطريقة جشعة باستخدام مقاييس مصممة يدوياً (مثل تقليل خطأ إعادة البناء لكل طبقة). ورغم أنها منخفضة التكلفة حسابياً، إلا أنها تعاني من انحيازات متأصلة وفجوات كبيرة في أداء المهام النهائي.
2. الأساليب القائمة على التعلم (مثل MaskLLM): تقوم بتحسين الأقنعة مباشرة عبر التدرج النازل على خسارة التوليد. وبينما هي فعالة، إلا أنها تتسبب في تكاليف تدريب باهظة. وتحديداً، يتطلب MaskLLM تخزين قيم "اللوجيت" (logits) لجميع الأقنعة الممكنة ضمن مجموعة مكونة من $M$ من الأوزان، مما يؤدي إلى تعقيد ذاكرة قدره $O(\binom{M}{N} \frac{d}{M})$. بالنسبة لقيم $M$ الكبيرة، يتوسع هذا بشكل أسي، مما يجعله غير ممكن للنماذج الكبيرة. علاوة على ذلك، تسبب المساحة التوافقية الشاسعة تبايناً عالياً في تدرجات السياسة، مما يسبب عدم استقرار التدريب.

المنهجية: MaskPro
يقترح المؤلفون MaskPro، وهو إطار عمل احتمالي جديد في مساحة خطية مصمم لتعلم ندرة (N:M) بكفاءة واستقرار.

1. الاسترخاء الاحتمالي في المساحة الخطية:
   بدلاً من تعريف توزيع فئوي لجميع $\binom{M}{N}$ من الأقنعة الممكنة لمجموعة من $M$ من الأوزان، يعيد MaskPro صياغة المشكلة. فهو يؤسس توزيعاً فئوياً لـ $M$ من العناصر الفردية داخل كل مجموعة. ومن ثم يتم إنشاء قناع (N:M) عن طريق إجراء أخذ عينات لـ $N$ من العناصر بدون إحلال من هذا التوزيع.

• الأساس النظري: يثبت المؤلفون أن أي قناع (N:M) يمكن تمثيله كمجموع احتمالي لـ $N$ من نواقل الأساس المتميزة.
• كفاءة الذاكرة: تقلل هذه إعادة الصياغة عدد المعلمات القابلة للتعلم (اللوجيت) من $O(\binom{M}{N} \frac{d}{M})$ إلى قيمة خطية $O(d)$، حيث $d$ هو إجمالي عدد الأوزان. وهذا يلغي الانفجار في الذاكرة المرتبط بالأساليب السابقة القائمة على التعلم.

2. مُقدّر تدرج السياسة المنقح (PGE):
   لتحسين "اللوجيت"، يستخدم MaskPro تدرجات السياسة. وقد حدد المؤلفون أن تدرجات السياسة القياسية باستخدام قيم الخسارة الخام تكون غير مستقرة في النماذج اللغبية الكبيرة بسبب التباين العالي الناتج عن اختلافات الدفعة الصغيرة (حيث غالباً ما يطغى التباين في الخسارة الناتج عن أخذ عينات البيانات على التباين الناتج عن أخذ عينات القناع).

• بقايا الخسارة: بدلاً من استخدام الخسارة الخام $f(m \odot w, \xi)$، يقوم MaskPro بتحديث "اللوجيت" بناءً على بقايا الخسارة بالنسبة لقناع أساسي $m_0$: $f(m \odot w, \xi) - f(m_0 \odot w, \xi)$. هذا يعزل تأثير تغيير القناع.
• متتبع المتوسط المتحرك: لمزيد من استقرار التدريب ومنع الانفجارات العددية، تم إدخال متتبع متوسط متحرك $\delta$ لتنعيم خط الأساس. تصبح قاعدة التحديث:
  $$\pi_{t+1} = \pi_t - \eta (f(m_t \odot w, \xi) - f(m_0 \odot w, \xi) - \delta) \nabla \log p(m_t|\pi_t)$$
• الضمانات النظرية: تقدم الورقة براهين توضح أن هذه المُقدِّرات المنقحة تظل مُقدِّرات غير منحازة لتدرج السياسة الحقيقي. علاوة على ذلك، وتحت شرط أن تكون خسارة القناع المأخوذ عينة أقل من ضعف خسارة خط الأساس (وهو شرط يتحقق بسهولة في الممارسة العملية)، فإن تباين المُقدِّر المنقح يكون أقل بصرامة من تدرج السياسة التقليدي.

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

1. إطار العمل في المساحة الخطية: يقلل MaskPro من متطلبات الذاكرة لتعلم ندرة (N:M) من التوسع الأسي $O(\binom{M}{N} \frac{d}{M})$ إلى المساحة الخطية $O(d)$ عبر نمذجة عملية الاختيار كأخذ عينات لـ $N$ من العناصر بدون إحلال من توزيع فئوي على $M$ من العناصر.
2. تدرج السياسة المعزز: يقدم المؤلفون آلية تحديث لتقليل التباين باستخدام بقايا الخسارة ومتتبع المتوسط المتحرك، مما يعالج عدم استقرار التدريب الناتج عن التباين العالي لتدرجات السياسة في المساحات التوافقية الكبيرة.
3. التحليل النظري: تم تقديم براهين شاملة فيما يتعلق بعدم انحياز المُقدِّرات وخصائص تقليل التباين لقواعد التحديث المقترحة.

النتائج التجريبية
أُجريت تجارب مكثفة على نماذج لغوية كبيرة متعددة (Gemma-7B, Vicuna-7B, LLaMA-2-7B, DeepSeek-7B) عبر مهام متنوعة (مثل HellaSwag, RACE, PIQA) وأنماط ندرة مختلفة ((2:4), (4:8), (8:16)).

• الأداء: يتفوق MaskPro باستمرار على الأساليب القائمة على القواعد (SparseGPT, Wanda, Pruner-Zero) ويحقق أداءً مقارباً لنموذج MaskLLM كثيف الذاكرة، وغالباً مع تحسن متوسط في الدقة يتجاوز 2%.
• قابلية توسع الذاكرة: بينما يواجه MaskLLM أخطاء نفاذ الذاكرة (OOM) على وحدات معالجة رسومية 8x A100 لندرة (4:8) في نماذج 7B، يحافظ MaskPro على استخدام مستقر للذاكرة (~36-37 جيجابايت) بغض النظر عن نمط الندرة.
• كفاءة البيانات: يُظهر MaskPro متانة ملحوومة تجاه حجم البيانات. إذ يحقق أداءً مستقراً حتى عند التدريب باستخدام عينة بيانات واحدة فقط، بينما يتطلب MaskLLM مئات الآلاف من العينات للوصول إلى التقارب.
• كفاءة التدريب: يتطلب الأسلوب عدداً أقل بكثير من عينات التدريب ووقتاً أقل مقارنة بـ MaskLLM، مع الحفاظ على استخدام للذاكرة مماثل للأساليب القائمة على القواعد.

الأهمية والادعاءات
تدعي الورقة أن MaskPro يوفر مساراً قابلاً للتطبيق للنشر المحلي القابل للتوسع للنماذج اللغوية الكبيرة من خلال حل اختناقات الذاكرة والاستقرار لتعلم الندرة شبه المهيكلة. ويؤكد المؤلفون أن نهجهم يسد الفجوة بين الأداء العالي للأساليب القائمة على التعلم وكفاءة الأساليب القائمة على القواعد. ومن خلال تقليل بصمة الذاكرة إلى المساحة الخطية وتثبيت عملية التدريب عبر تدرجات السياسة المنقحة، يمكّن MaskPro من تعلم أقنعة (N:M) المثلى للنماذج واسعة النطاق دون التكاليف الحسابية الباهظة المرتبطة سابقاً بمثل هذه المهام. ويُقدم الأسلوب كإطار عمل قابل للتعميم لتعلم الأقنعة الاحتمالية الذي يتميز بالمتانة تجاه ندرة البيانات وحجم النموذج.