Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction

بناءً على ورقة البحث "اجعل كل رمز ذا قيمة: نحو تحسين أداء السياق الطويل من خلال استبعاد ذاكرة التخزين المؤقت للمفاتيح والقيم (KV Cache Eviction)"، إليك ملخص تقني مفصل.

بيان المشكلة

تعد ذاكرة التخزين المؤقت للمفاتيح والقيم (KV cache) ضرورية لعملية فك التشفير ذاتي الانحدار (autoregressive decoding) الفعالة في نماذج المحولات (Transformers) للغات الكبيرة (LLMs) والنماذج اللغوية البصرية (VLMs). ومع ذلك، مع توسع نوافذ السياق لتصل إلى عشرات الآلاف من الرموز (بما في ذلك الرموز النصية والبصرية)، تصبح ذاكرة الـ KV عائقاً رئيسياً بسبب النمو الخطي في الذاكرة والتكاليف التربيعية لحساب الانتباه (attention computation).

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

المنهجية: الاستبعاد القائم على الاحتفاظ العالمي

يقترح المؤلفون طريقة DBTrimKV (Dynamic Budget TrimKV)، وهي طريقة تتعلم التنبؤ بالمنفعة المستقبلية للرموز المخزنة وتخصص ميزانية ذاكرة عالمية واحدة عبر جميع الطبقات، والرؤوس، والأنماط (modalities).

1. الرؤية النظرية: تخفيف الانتباه (Attention Dilution)
تقدم الورقة تحليلاً نظرياً يوضح أن انتباه الذاكرة الكاملة يعاني من "تخفيف الانتباه". عندما يوجد عدد كبير من رموز المشتتات في الذاكرة، تجبر عملية تسوية (softmax normalization) الرموز المفيدة على التنافس معها، مما يتسبب في تلاشي كتلة الانتباه المخصصة للأدلة ذات الصلة. يثبت المؤلفون أن الاحتفاظ التفضيلي بالرموز المفيدة مع استبعاد المشتتات يمكن أن يقلل من هذا التخفيف ويحسن جودة الانتباه نظرياً.

2. الاحتفاظ الهندسي كمنفعة مستقبلية
بدلاً من استخدام درجات الانتباه المعتمدة على الاستعلام (والتي تعكس الارتباط قصير المدى)، تتعلم الطريقة درجة احتفاظ هندسية ($\beta_i$) لكل رمز.

• المفهوم: يتم نمذجة منفعة الرمز كحدث استمرارية. إذا كان من المرجح أن يظل الرمز مفيداً لخطوات فك التشفير المستقبلية، فإنه يحصل على درجة احتفاظ عالية.
• الصياغة: يتم تعريف وزن الاحتفاظ للرمز $i$ عند الخطوة $t$ بـ $r_{t,i} = \beta_i^{t-i}$.
• التبرير: يبرر المؤلفون هذا الاضمحلال الهندسي نظرياً كبديل مستقل عن الاستعلام (query-agnostic) للمنفعة المستقبلية، مدعوماً بملاحظات تجريبية تشير إلى أن بقاء انتباه الرموز ينخفض بسرعة بمرور الوقت في السياقات الطويلة.

3. البوابات العالمية المرتبطة بالأوزان (Weight-Tied Global Gates)
لتمكين سياسة استبعاد موحدة، قدم المؤلفون بوابات احتفاظ مرتبطة بالأوزان:

• البنية: تتنبأ بوابة خفيفة الوزن بدرجة احتفاظ عددية (scalar) $\beta$ لكل مدخل KV. والأهم من ذلك، أن طبقة الإسقاط النهائية لهذه البوابات هي مشتركة (tied) عبر جميع طبقات المحول ورؤوس الانتباه.
• المعايرة العالمية: يفرض هذا التشارك في الأوزان وضع درجات الاحتفاظ من مختلف الطبقات والرؤوس على مقي scale مشترك. وبناءً على ذلك، يمكن مقارنة الرموز من رأس كثيف النصوص مباشرة مع رموز من رأس بصري أو طبقة عميقة.
• الميزانية العالمية: بدلاً من تخصيص ميزانيات ثابتة لكل طبقة أو رأس، يحافظ النظام على ميزانية KV عالمية واحدة ($M_{global}$). عند أي خطوة استبعاد، يقوم النظام بترتيب جميع مدخلات KV المخزنة عالمياً حسب المنفعة المستقبلية المتوقعة ويحتفظ بأفضل $M_{global}$ من الرموز.

4. التنفيذ
تستخدم الطة تخطيط PagedAttention لدعم تسلسلات KV متغيرة الطول لكل رأس، مما يسمح بأحجام ذاكرة مؤقتة ديناميكية وخاصة بكل رأس دون متطلبات ذاكرة فيزيائية متصلة.

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

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

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

تم تقييم الطريقة على معايير متنوعة تشمل الاستنتاج طويل الأمد (MathVision, AIME)، والحوار متعدد الدورات (MMDU)، والإجابة على الأسئلة البصرية (VQAText, MME).

• الأداء مقابل الذاكرة الكاملة: في العديد من السيناريوهات، وخاصة تحت ميزانيات الذاكرة الضيقة، طابقت أو تجاوزت DBTrimKV أداء الاستدلال بالذاكرة الكاملة (Vanilla). على سبيل المثال، في معيار الحوار متعدد الدورات MMDU، تفوقت DBTrimKV بميزانية 512 رمزاً على خط الأساس للذاكرة الكاملة بنسبة 14%.
• المقارنة مع النماذج المرجعية: تفوقت DBTrimKV باستمرار على طرق استبعاد الـ KV الحالية (مثل SnapKV و AdaKV و R-KV) وعلى نماذج تقليم الرموز البصرية (مثل VisionZip و FastV).
• الكفاءة: قللت الطريقة بشكل كبير من استخدام ذاكرة الـ KV ووقت فك التشفير مقارنة بالاستدلال بالذاكرة الكاملة، وتوسعت بكفاءة مع طول السياق وطول التوليد.
• دراسات الاستبعاد (Ablation Studies): أدى إزالة الإسقاط المرتبط بالأوزان (مما فرض درجات محلية) إلى انخفاض كبير في الأداء، مما يؤكد ضرورة المعايرة العالمية.

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

تدعي الورقة أن استبعاد الـ KV المتعلم والمعاير عالمياً ليس مجرد تقنية ضغط، بل هو آلية لتحسين الاستنتاج في السياق الطويل.

من خلال التعامل مع الاستبعاد كوسيلة لكبح المشتتات بدلاً من مجرد توفير الذاكرة، يظهر المؤلفون أن "نسيان" الرموز غير ذات الصلة يمكن أن يشحذ الانتباه ويعزز دقة النموذج. تشير النتائج إلى أن إدارة الذاكرة الديناميكية والمتعلمة يمكن أن تتفوق على الاستدلال بالذاكرة الكاملة الثابتة في إعدادات السياق الطويل، مما يوفر مساراً لأنظمة ذكاء اصطنا_ي أكثر كفاءة وقدرة على التعامل مع الآفاق الطويلة. ويشير المؤلفون إلى القيود المتعلقة بحجم النماذج المختبرة (حتى 8 مليار معلمة) وإمكانية العمل المستقبلي في التدريب المشترك (end-to-end) لبوابات الاحتفاظ والنموذج الأساسي.