Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference

إليك ملخص تقني مفصل لورقة البحث بعنوان "الاختيار التكيفي للطبقات لتقليم الرموز على مستوى الطبقة في استنتاج النماذج اللغوية الكبيرة (LLM)."

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

تواجه النماذج اللغوية الكبيرة (LLMs) اختناقات كبيرة في الذاكرة أثناء عملية الاستنتاج بسبب ذاكرة المفتاح-القيمة (KV cache)، التي تخزن متجهات الرموز السابقة لتمكين التوليد الفعال. وبينما توجد العديد من تقنيات تقليل ذاكرة (KV cache)، برزت طرق التقليم الرمزي على مستوى الطبقة (مثل FastKV وGemFilter) كحل شائع. تقوم هذه الطرق باختيار مجموعة فرعية من الرموز المهمة في طبقات محددة للاحتفاظ بها في ذاكرة (KV cache)، وتقليم البقية.

ومع ذلك، تعاني الطرق الحالية من قصور حرج وهو: عدم المرونة. فهي تعتمد على طبقات اختيار محددة مسبقاً وثابتة يتم تحديدها بشكل مستقل عن المهمة المحددة.

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

2. المنهجية: طبقة الاختيار التكيفية (ASL)

يقترح المؤلفون ASL، وهي طريقة تعتمد على إدراك المهام ولا تتطلب تدريباً، وتحدد ديناميكياً الطبقة المثلى لاختيار الرموز خلال مرحلة التعبئة المسبقة (prefilling stage).

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

تستغل ASL ملاحظة مفادها أن درجات الانتباه تتطور عبر الطبقات:

1. الطبقات المبكرة: تكون درجات الانتباه موزعة بشكل موحد عبر رموز السياق.
2. الطبقات العميقة: تتركز درجات الانتباه في مجموعة فرعية مستقرة وصغيرة من الرموز.

تقوم ASL بمراقبة تباين رتب الرموز (variance of token ranks) المرتبة حسب درجات الانتباه للكشف عن وقت حدوث هذا التركيز.

خطوات الخوارزمية:

1. نافذة المراقبة: بدءاً من الحد الأدنى للطبقة ($L_{min}$)، تراقب ASL الطبقات اللاحقة ($L_{obs}$).
2. تجميع الدرجات: تقوم بتجميع درجات الانتباه عبر رؤوس الانتباه وتطبق تجميعاً متوسطاً أحادي البعد (1D average pooling) لتنعيم الضجيج والتقاط مناطق السياق المتصلة.
3. حساب الرتبة: بالنسبة لآخر ($L_{obs}$) من الطبقات، تحدد أفضل ($k$) من الرموز وتحسب تباين رتبها.
4. التباين النسبي: لجعل المقياس غير مرتبط بمهمة محددة، يتم تطبيع التباين الحالي من خلال التباين الأولي عند ($L_{min}$):
   $$\text{relative\_variance}(L) = \frac{\text{vartop}(L)}{\text{vartop}(L_{min})}$$
5. القرار التكيفي: يتم وضع عتبة ($\tau$) يحددها المستخدم. مع معالجة النموذج للطبقات، إذا انخفض التباين النسبي عن $\tau$، فهذا يشير إلى أن رتب الرموز قد استقرت. يتم تعيين الطبقة الحالية فوراً كـ طبقة الاختيار.
6. التنفيذ:
   • الاختيار لمرة واحدة: بمجرد اختيار الطبقة، يتم تقليم الرموز، ويتم تمرير الرموز المختارة فقط إلى جميع الطبقات اللاحقة.
   • التكامل: يمكن دمج ASL مع التقنيات الموجودة (مثل SnapKV للضغط قبل الاختيار، أو GemFilter لاستراتيجيات المرور المزدوج) لتلبية قيود ميزانية (KV) المحددة من قبل المستخدم.

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

1. تحديد عدم المرونة: توضح الورقة أن طبقات الاختيار الثابتة في أحدث الأساليب (SOTA) تؤدي إلى تباين كبير في الأداء عبر المهام، وخاصه الفشل في مهام الاسترجاع عالية الصعوبة.
2. اقتراح ASL: تقديم طريقة تكيفية جديدة لا تتطلب تدريباً وتحدد طبقة الاختيار بناءً على تباين رتب الرموز، مما يسمح للنموذج بالتكيف مع صعوبة المهمة دون ضبط يدوي.
3. القدرة على التكامل: صُممت ASL لتعمل بسلاسة مع تقنيات تقليل ذاكرة (KV cache) الحالية (مثل SnapKV وGemFilter) لتحسين كل من مرحلتي التعبئة المسبقة وفك التشفيد مع الالتزام بميزانيات الذاكرة المحددة من قبل المستخدم.
4. التقييم الشامل: تجارب واسعة النطاق عبر ثلاثة معايير رئيسية (InfiniteBench وRULER وNIAH) مع أطوال سياق تصل إلى 256 ألف، مما يظهر تفوقاً في المقايضة بين الدقة والكفاءة مقارنة بالأساليب الحالية.

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

قيم المؤلفون ASL على نماذج Llama-3.1-8B-UltraLong و Qwen2.5-7B.

• الدقة:

  • المهام الصعبة: تتفوق ASL بشكل كبير على طرق الطبقة الثابتة (FastKV وGemFilter) في المهام الصعبة مثل استرجاع KV و Multi-key NIAH. على سبيل المثال، في مهمة استرجاع KV في InfiniteBench مع ميزانية KV تبلغ 2048، حققت ASL دقة تقارب دقة الـ KV الكاملة، بينما انخفضت دقة FastKV بشكل كبير.
  • الأداء العام: في InfiniteBench، حققت ASL (نسخة المرور الثاني) أعلى متوسط درجة لنموذج Llama-3.1-8B-UL. وفي RULER (سياق 128k)، تفوقت ASL باستمرار على FastKV، خاصة مع زيادة طول السياق.
  • NIAH: حققت ASL درجات كاملة عبر جميع أطوال السياق (1k–256k) في اختبار "إبرة في كومة قش"، مطابقةً أداء الـ KV الكامل، بينما واجهت الطرق الأخرى صعوبات عند أعماق معينة.

• الكفاءة (TTFT & TPOT):

  • الوقت حتى أول رمز (TTFT): تعتبر ASL أبطأ قليلاً من FastKV لأنها غالباً ما تختار الرموز عند طبقات أعمق (أي $L_{select}$ أعلى) لضمان الدقة في المهام الصعبة. ومع ذلك، تظل أسرع بكثير من الـ KV الكامل.
  • الوقت لكل رمز مخرج (TPOT): تحقق ASL سرعات فك تشفير مقاربة لطرق التقليم الأخرى، حيث تعتمد مرحلة فك التشفير على مجموعة الرموز المختزلة بغض النظر عن وقت الاختيار.
  • الإنتاجية (Throughput): بينما تضحي ASL ببعض الإنتاجية مقابل الدقة، فإن الفجوة تتقلص في السيناريوهات الواقعية ذات تسلسلات المخرجات الطويلة حيث يهيمن وقت فك التشفير.

• استخدام الذاكرة: تفرض ASL عبئاً ضئيلاً على الذاكرة (تخزين درجات مجمعة ومخازن رتب لعدد قليل من الطبقات فقط) وتحقق تقليلاً في الذاكرة مقارباً لطرق التقليم الأخرى (على سبيل المثال، تقليل استخدام 18 جيجابايت إلى حوالي 0.3 جيجابايت لنموذج Llama-3.1-8B).

5. الأهمية

تعالج هذه الورقة مشكلة الصلابة الجوهرية في تحسين استنتاج النماذج اللغوية الكبيرة حالياً. من خلال الانتقال من الاستدلالات الثابتة والمحددة مسبقاً إلى اتخاذ القرار الديناميكي القائم على البيانات، تسمح ASL للنماذج اللغوية بـ:

• الموازنة بين الدقة والسرعة: المقايضة تلقائياً بين سرعة الاستنتاج والدقة بناءً على صعوبة المهمة المدخلة.
• التعامل مع السياقات الطويلة بمتانة: الحفاظ على أداء عالٍ في مهام الاسترجاع المعقدة (مثل العثور على مفاتيح محددة في مستند ضخم) حيث تفشل الطرق السابقة.
• التعميم دون تدريب: العمل كوحدة جاهزة للتشغيل (plug-and-play) دون الحاجة إلى ضبط دقيق (fine-tuning)، مما يجعلها قابلة للتطبيق الفوري في عمليات نشر النماذج اللغوية الكبيرة الحالية.

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