ملخص تقني: StiefAttention

بيان المشكلة

تواجه عملية الاستدلال في سياقات النصوص الطويلة (Long-context inference) في النماذج اللغوية الكبيرة (LLMs) عقبات متزايدة تتمثل في متطلبات سعة الذاكرة وعرض النطاق الترددي لذاكرة المفاتيح والقيم (KV cache). فخلال عملية فك التشفير ذاتي الانحدار (autoregressive decoding)، ينمو مخزن KV بشكل خطي مع طول التسلسل وحجم الدفعة (batch size)، مما يستهلك غالبًا ذاكرة تضاهي ذاكرة أوزان النموذج نفسه (على سبيل المثال، حوالي 16 جيجابايت لنموذج Llama3-8B عند سياق طوله 32 ألفًا).

تستخدم طرق الضغط الحالية بعد التدريب (post-training compression) عادةً تفكيك المصفوفات من الرتبة المنخفضة (مثل تفكيك القيم المفردة SVD) لإسقاط المفاتيح والقيم في أبعاد أدنى. ومع ذلك، يجادل المؤلفون بأن هذه الأهداف الوسيطة (proxy objectives) — مثل تقليل خطأ إعادة البناء من نوع Frobenius لمتجهات المفاتيح والقيم الخام أو درجات الانتباه قبل الـ softmax ($QK^\top$) — تفشل في مراعاة كيفية تفاعل أخطاء إعادة البناء مع عمليات softmax للانتباه، وخلط القيم (value mixing)، وتحويلات طبقة فك التشفير اللاحقة (التطبيع، المسارات المتبقية، والـ MLPs). وبناءً على ذلك، فإن الإسقاط الذي يكون مثاليًا لإعادة بناء الموترات الوسيطة قد يتسبب في النهاية في أخطاء كبيرة في مخرجات طبقة فك التشفير النهائية، مما يؤدي إلى تدهور جودة التوليد من البداية إلى النهاية.

المنهجية: StiefAttention

يقدم البحث StiefAttention، وهو طريقة لضغط مخزن KV بعد التدريب تعمل على تقليل خطأ إعادة البناء مباشرة عند مخرجات طبقة ف decoder-layer output بدلاً من الكميات الوسيطة للانتباه.

الآلية الجوهرية

1. دالة الهدف: بدلاً من تقليل $\|K - \tilde{K}\|_F$، يعمل StiefAttention على تقليل الخطأ النسبي بين مخرج طبقة فك التشفير الأصلي $f_\ell(x)$ والمخرج $\tilde{f}_\ell(x)$ الناتج عن آلية الانتباه المضغوطة. تلتقط هذه الدالة التأثير التراكمي لعمليات softmax، وخلط القيم، واللاخطية اللاحقة.
2. الإسقاط المتعامد على منوع ستيفل (Stiefel Manifold): تتعلم الطريقة قواعد إسقاط متعامدة ($P_K, P_V$) تقع على منوع ستيفل (المصفوفات ذات الأعمدة المتعامدة).
3. التنبؤ بالقاعدة القائم على التدرج (Gradient-Based Basis Prediction):
   • بدلاً من حساب SVD لمجموعة المعايرة الكاملة (وهو أمر مستهلك للذاكرة)، يستخدم StiefAttention متنبئًا خفيف الوزن وقابلًا للتدريب (MLP) لتوليد قواعد الإسقاط من إحصائيات التنشيط البسيطة.
   • الميزات المدخلة: يأخذ المتنبئ المتوسط ($\mu$) والتباين ($\sigma^2$) لكل بُعد لمتجهات المفاتيح والقيم كمدخلات.
   • التدريب: يتم تدريب الـ MLP عبر التدرج المنحدر (gradient descent) لربط هذه الإحصائيات بالمصفوفات $A_K$ و $A_V$.
   • التعامد (Orthonormalization): يتم جعل المصفوفات الناتجة متعامدة عبر تفكيك QR ($A = QR$). ويعمل عامل التعامد $Q$ كقاعدة للإسقاط، حيث يمثل العامل المثلثي $R$ مجرد تغيير في الإحداثيات داخل الفضاء الجزئي.
4. تخصيص الرتبة (Rank Allocation): تبني الطريقة ملفات تعريف رتب الخطأ لكل طبقة. وهي تستخدم استراتيجية تخصيص ميزانية تسلسلية حيث يتم اختيار الرتب طبقة تلو الأخرى لتقليل الخطأ التراكمي لمخرجات طبقة فك التشفير تحت ميزانية إجمالية محددة لمخزن KV. وهذا يأخذ في الاعتبار انتشار الخطأ عبر الطبقات.
5. الاستدلال: في وقت الاستدلال، تُخزن المفاتيح والقيم في شكل مضغوط ($K\downarrow = K P_K$). ويتم إعادة بناء الأبعاد الكاملة فورًا ($\tilde{K} = K\downarrow P_K^\top$) قبل حساب الانتباه. ويمكن دمج قواعد الإسقاط في أوزان الانتباه الخطية، مما لا يضيف أي عبء تشغيل إضافي مقارنة بطرق الرتبة المنخفضة القياسية.

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

1. التحسين المباشر لمخرجات طبقة فك التشريد: يقترح المؤلفون تحسين ضغط مخزن KV مباشرة مقابل خطأ مخرجات طبقة فك التشفير، متجاوزين الأهداف الوسيطة مثل إعادة بناء التنسور أو تقريب درجات الانتباه.
2. إطار عمل StiefAttention: هي طريقة ما بعد التدريب تتعلم قواعد متعامدة من إحصائيات التنشيط باستخدام متنبئ خفيف الوزن، مما يسمح بضغط فعال دون الحاجة لإعادة تدريب النموذج بالكامل.
3. تخصيص الرتبة التسلسلي: استراتيجية لتخصيص الرتب عبر الطبقات بناءً على أسطح الخطأ المحلية التي يتم تقييمها على مدخلات تتضمن بالفعل آثار ضغط الطبقات السابقة.
4. التحقق التجريبي: تقييم شامل على نماذج Llama3-8B و Qwen3-8B يظهر أن StiefAttention يتفوق على النماذج المرجعية (Baselines) المتطورة (تحديدًا EigenAttention) تحت ميزانيات مخزن KV متساوية.

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

قام المؤلفون بتقييم StiefAttention على Llama3-8B و Qwen3-8B، ومقارنته بالمرجع الأساسي FP16 غير المضغوط و EigenAttention (وهي طريقة رائدة تعتمد على SVD وتدمج الاستعلامات Queries في عملية الإسقاط).

• مكاسب الأداء: في ظل ظروف ضغط متساوية (نفس ميزانية مخزن KV)، تفوق StiefAttention بشكل كبير على EigenAttention:
  • Llama3-8B: حقق تحسنًا قدره 4.2 نقطة في C4 perplexity وزيادة قدرها 8.9 نقطة في دقة 0-shot MMLU مقارنة بـ EigenAttention.
  • Qwen3-8B: أظهر اتجاهات مماثلة، حيث حافظ StiefAttention على مستويات أقل من الـ perplexity ودقة أعلى، خاصة مع زيادة نسب الضغط (الضغط الأقوى).
• المتانة: أظهرت الطريقة تدهورًا سلسًا في الأداء مع انخفاض نسبة مخزن KV، دون حدوث انهيار مفاجئ.
  টি التوافق مع التكميم (Quantization Compatibility): وُجد أن StiefAttention أكثر متانة بشكل ملحوظ تجاه تكميم مخزن KV ذي عدد البتات المنخفض (مثل 4-bit و 8-bit) مقارنة بـ EigenAttention، حيث حافظ على perplexity أقل عند أحجام مخزن مماثلة.

تحليل الأهمية

يقدم البحث تحليلًا تشخيصيًا يشرح لماذا يحسن StiefAttention الأداء من البداية إلى النهاية:

• الاتساق الاتجاهي: بينما يعيد EigenAttention بناء مخرجات الانتباه الوسيطة بشكل أفضل (مقللًا خطأ مخرجات الانتباه بنسبة ~29.6% مقارنة بـ StiefAttention)، إلا أن StiefAttention يحقق حفظًا أفضل لـ اتجاه مخرجات طبقة فك التشفير (تشابه جيب التمام أعلى، +3.3%) وخطأ أقل في مخرجات الطبقة الإجمالية (تحسن بنسبة 5.2%).
• حساسية الطبقات المبكرة: يكشف التحليل أن فجوة تشابه جيب التمام (cosine similarity) تكون أكثر وضوحًا في الطبقات المبكرة. ويفترض المؤلفون أن الاضطرابات في الطبقات المبكرة تتضخم بواسطة الكتل اللاحقة غير الخطية (softmax، المسارات المتبقية، MLP). لذا، فإن الأساليب القائمة على SVD، التي تحسن إعادة البناء الوسيطة، غالبًا ما تفشل في الحفاظ على المعلومات الاتجاهية المحددة المطلوبة لهذه الكتل اللاحقة. ومن خلال التحسين المباشر لمخرجات فك التشفير، يقوم StiefAttention بمحاذاة قواعد الإسقاط مع "الفضاء الجزئي ذي الصلة بالمخرجات"، مما يؤدي إلى جودة توليد نهائية أفضل.

القيود والعمل المستقبلي

يقر المؤلفون بعدة قيود:

• النطاق: أجريت التجارب بشكل أساسي عند طول تسلسل 2048 بسبب قيود ذاكرة وحدة معالجة الرسومات (GPU)؛ وهناك حاجة لتقييم أوسع على سياقات أطول.
• تبسيط التحسين: تستخدم عملية اختيار الرتب مسارًا تسلسليًا جشعًا (greedy) بدلاً من التحسين العالمي الشامل.
• تمثيل الميزات: يستخدم المتنبئ حاليًا المتوسط والتباين القطري فقط؛ وترك العمل على استخدام ملخصات تنشيط أغنى (مثل مخططات التباين Covariance sketches) للعمل المستقبلي.
• تقييم النظام: لا يزال التقييم الكامل للنظام، بما في ذلك زمن استجابة الـ prefill/decode والإنتاجية (throughput) في ظروف الخدمة الواقعية، قيد الانتظار.

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