ملخص تقني: A3: إطار عمل تحليلي للتقريب منخفض الرتبة لآلية الانتباه

بيان المشكلة

حققت النماذج اللغوية الكبيرة (LLMs) أداءً مذهلاً، لكنها تعاني من ضخامة عدد المعلمات، مما يؤدي إلى تكاليف حوسبة عالية، ومتطلبات ذاكرة مرتفعة، ونشر مكلف. وبينما يمثل التقريب منخفض الرتبة استراتيجية ضغط واعدة عبر تفكيك مصفوفات الأوزان ($W \approx A_r B_r$)، تواجه النهج الحالية قيوداً حرجة تتمثل في:

1. العمى المعماري: فهي عادةً ما تقلل خطأ المخرجات للطبقات الخطية الفردية دون مراعاة الخصائص الوظيفية المحددة لمكونات الـ Transformer (مثل درجات الانتباه مقابل مخرجات الانتباه).
2. العبء التشغيلي أثناء وقت التشغيل: حيث تقوم بتفكيك مصفوفة وزن واحدة كبيرة إلى مصفوفتين أصغر، مما يستبدل عملية ضرب مصفوفات عامة (GEMM) واحدة بعمليتي GEMM أصغر. وهذا يؤدي إلى أعباء تشغيلية، بما في ذلك عمليات استدعاء النواة (kernel launches) الإضافية وعمليات الذاكرة، مما ينتج غالباً أداءً أدنى من طرق التقليم (pruning) أو التكميم (quantization).

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

المنهجية: إطار عمل A3

يقترح المؤلفون A3، وهو إطار عمل للتقريب التحليلي منخفض الرتبة لما بعد التدريب، يعالج هذه القيود عبر تفكيك طبقة الـ Transformer إلى ثلاثة مكونات وظيفية: QK (الاستعلام-المفتاح)، OV (المخرج-القيمة)، و MLP (الشبكة العصبية متعددة الطبقات). بدلاً من تقليل أخطاء الأوزان أو مخرجات الطبقات العامة، يعمل A3 على تقليل الخسارة الوظيفية الخاصة بكل مكون، مستنتجاً حلولاً تحليلية مغلقة الشكل تقلل الأبعاد الخفية مباشرة.

1. مكون A3-QK (الاستعلام-المفتاح)

• الهدف: تقليل الخطأ في درجات الانتباه قبل الـ softmax ($A_i = Q_i K_i^T$).
• الطريقة: يصيغ المؤلفون المشكلة كتقليل لـ "معيار فروبينيوس" (Frobenius norm) للفرق بين درجات الانتباه الأصلية والمقربة على مجموعة معايرة.
• الحل: استنتجوا حلاً مغلق الشكل (النظرية 3.2) يتضمن مصفوفات الارتباط الذاتي لمدخلات الاستعلام والمفتاح/القيمة ($R_{X_q X_q}$ و $R_{X_{kv} X_{kv}}$). يتم الحصول على التقريب الأمثل منخفض الرتبة عبر تحليل القيم المفردة (SVD) المرجح والمبتور لمصفوفة الوزن المدمجة.
• التوسع: بالنسبة لآلية "استعلام المجموعة" (GQA)، يتم تطبيق SVD مشترك عبر الرؤوس التي تشترك في نفس المفتاح/القيمة. وبالنسبة لترميز الموضع الدوار (RoPE)، يتم استخدام تفكيك CUR للتعامل مع الدورات الزوجية، مع ترتيب الأبعاد حسب مجموع أهمية الأزواج المتجاورة.

2. مكون A3-OV (المخرج-القيمة)

• الهدف: تقليل الخطأ في مخرج الانتباه ($O = \sum A'_i V_i W_o$).
• الطريقة: بدلاً من تقليل خطأ مخرج الانتباه العالمي (الذي يتسبب في تكاليف حوسبية عالية)، يقوم إطار العمل بتخفيف الهدف لتقليل خطأ مخرج كل رأس بشكل مستقل.
• الحل: تم اشتقاق الحل الأمثل (النظرية 3.3) بشكل مشابه لـ A3-QK ولكن مع التركيز على الارتباط الذاتي لدرجات الانتباه بعد الـ softmax مضروبة في مدخلات المفتاح/القيمة ($P_i$). وهذا يقلل من بُعد المخرج المشترك بين إسقاطات القيمة والمخرج.

3. مكون A3-MLP (الشبكة العصبية متعددة الطبقات)

• الهدف: تقليل خطأ مخرج طبقة الـ MLP.
• التحدي: وجود دوال التنشيط غير الخطية (مثل SiLU أو ReLU) يمنع تطبيق SVD المباشر.
• الطريقة: يصيغ إطار العمل المشكلة كاختيار لمجموعة فرعية من أعمدة مصفوفة أوزان الإسقاط السفلي ($W_d$) للاحتفاظ بها، مما يقلل فعلياً البُعد المتوسط ($d_{inter}$).
• الحل: باستخدام مبادئ تفكيك CUR، يختار المؤلفون أفضل-$r$ من الأعمدة بناءً على حاصل ضرب معايير $L_2$ لصفوف أوزان الإسقاط العلوي وأعمدة مصفوفة الارتباط الذاتي للتنشيطات المتوسطة. هذا يمثل استدلالاً بسيطاً وفعالاً لاختيار الأعمدة في المسائل الصعبة (NP-hard).

ميزة التصميم الرئيسية: الكفاءة التشغيلية

على عكس طرق الرتبة المنخفضة التقليدية التي تستبدل $XW$ بـ $(XA)B$ (مما يضيف عملية GEMM إضافية)، يقوم A3 بتقليل الأبعاد الخفية ($d_{qk}$, $d_{vo}$, $d_{inter}$) مباشرة. ينتج عن ذلك نفس عدد عمليات GEMM كما في النموذج الأصلي ولكن مع أحجام مصفوفات أصغر. وبالتالي، يقلل A3 من حجم النموذج، وحجم ذاكرة التخزين المؤقت لـ KV (KV cache)، وعدد العمليات الحسابية (FLOPs) دون إدخال أي عبء تشغيلي (لا توجد عمليات استدعاء نواة إضافية).

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

1. تفكيك المكونات الوظيفية: إعداد منخفض الرتبة ثلاثي الأجزاء لآلية تعدد الرؤوس (MHA) يصيغ التحسين كتقليل للخسائر الوظيفية لدرجات QK، ومخرجات OV، ومخرجات MLP، بدلاً من مخرجات الطبقات العامة.
2. حلول تحليلية مغلقة الشكل: اشتقاق الحلول المثلى للأهداف الثلاثة التي تقلل الأبعاد الخفية بشكل طبيعي، متجنبة العبء التشغيلي للمصفوفات المفككة.
3. القابلية للتكيف المعماري: توسيع إطار العمل لدعم متغيرات الـ Transformer الحديثة، بما في ذلك GQA و RoPE، متجاوزاً بذلك قيود الطرق السابقة التي اقتصرت على Vanilla MHA.
4. تطبيقات متعددة الاستخدامات: إثبات توافق A3 مع ضغط ذاكرة التخزين المؤقت لـ KV، والتكميم، والضبط الدقيق، وتخصيص الرتب المختلطة.

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

أجرى المؤلفون تجارب واسعة النطاق على نماذج لغوية متنوعة (LLaMA 1/2/3.1, MPT, Phi 3, Mistral) عبر مهام ما قبل التدريب (WikiText-2, C4, SlimPajama) والاختبارات النهائية (ARC, BoolQ, MMLU, إلخ).

• تفوق الأداء: يتفوق A3 بشكل كبير على أساليب الرتبة المنخفضة الرائدة (مثل SVD-LLM, ASVD, FWSVD).
  • مثال: في نموذج LLaMA 3.1-70B مع نسبة ضغط 10%، حقق A3 درجة إدراك (perplexity) تبلغ 4.69 على WikiText-2، متفوقاً على SVD-LLM الذي سجل 7.87 (بفارق تقليل قدره 3.18 نقطة في الإدراك).
  • في LLaMA-2-7B، يحافظ A3 باستمرار على إدراك أقل عبر نسب ضغط 10% و20% و40% مقارنة بـ SVD و FWSVD و SVD-LLM.
• مكاسب الإنتاجية (Throughput): نظراً لأن A3 يقلل أحجام المشكلات دون عمليات استدعاء نواة إضافية، فإنه يحقق تسريعاً ثابتاً في عدد الرموز في الثانية (TPS). في المقابل، غالباً ما تعاني طرق مثل SVD-LLM من عقوبات في السرعة أو تظهر مكاسب فقط عند مستويات ضغط شديدة.
• كفاءة ذاكرة التخزين المؤقت لـ KV: يوفر A3 ضغطاً أقوى لذاكرة التخบั المؤقت لـ KV مقارنة بطرق مثل Palu و CLOVER، محافظاً على إدراك أقل حتى عند نسب ضغط عالية (80% على MPT-30B).
• التوافق: يتوافق A3 بفعالية مع التكميم (مثل HQQ)، مما يخلق طيفاً مستمراً من مستويات الضغط ويحقق حدود باريتو (Pareto frontier) أفضل من التكميم وحده في مستويات ما دون 4-بت.

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

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

• دقة فائقة: الحفاظ على أداء النموذج النهائي بشكل أفضل بكثير من أساليب التقريب منخفض الرتبة الحالية.
• كفاءة الأجهزة: تقليل حجم النموذج، وبصمة الذاكرة (بما في ذلك KV cache)، وعدد العمليات الحسابية (FLOPs) دون العقوبات الزمنية المرتبطة بضرب المصفوفات المفككة.
• قابلية تطبيق واسعة: النجاح في التكيف مع البنيات الحديثة (GQA, RoPE) التي لم تستطع الطرق السابقة التعامل معها بفعالية.

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