KVSculpt: KV Cache Compression as Distillation

إليك ملخص تقني مفصل لورقة البحث "KVSCULPT: ضغط ذاكرة التخزين المؤقت لـ KV كعملية تقطير".

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

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

تندرج طرق الضغط الحالية على بعد طول التسلسل عموماً ضمن فئتين:

1. الاستبعاد (Eviction): اختيار مجموعة فرعية من أزواج KV الأصلية للاحتفاظ بها (على سبيل المثال، بناءً على درجات الانتباه أو الحداثة) واستبعاد الباقي.
2. الدمج (Merging): دمج أزواج KV المتشابهة في متوسطات موزونة.

القصور: كلتا الطريقتين منفصلتان (discrete) ومرتبطتان (anchored) بمدخلات الذاكرة الأصلية. تفترض هاتان الطريقتان أن الذاكرة المضغوطة يجب أن تكون مجموعة فرعية أو مزيجاً خطياً من المتجهات الأصلية. يرى المؤلفون أن هذا قيد غير ضروري لأن أزواج KV، بعد تطبيق تضمين الموضع الدوار (RoPE)، توجد في فضاء تضمين مستمر حيث تم تشفير المعلومات الموضعية بالفعل. لا يوجد متطلب نظري بأن تكون الذاكرة المضغوتة مجموعة فرعية من الأصل؛ فأي مجموعة من المتجهات التي تعيد إنتاج سلوك الانتباه الصحيح ستكون صالحة.

2. المنهجية: KVSCULPT

يعيد KVSCULPT صياغة ضغط ذاكرة KV كمسألة تقطير (distillation). بدلاً من اختيار أو دمج الأزواج الأصلية، يقوم بتحسين مجموعة أصغر من أزواج KV غير المقيدة في فضاء التضمين المستمر للحفاظ على سلوك الانتباه للذاكرة الكاملة.

المكونات الأساسية:

• صياغة المشكلة:

  • بالنظر إلى ذاكرة كاملة مكونة من $N$ من الأزواج، الهدف هو تقطير "منطقة الضغط" (أقدم $N-m$ زوجاً) إلى $k$ من الأزواج المحسنة بحرية $(K_c, V_c)$.
  • يتم الإبقاء على "منطقة الاحتفاظ" (أحدث $m$ زوج) دون تغيير لضمان دقة السياق الحديث.
  • الهدف هو مطابقة مخرجات الانتباه للذاكرة المضغوطة مع الذاكرة الكاملة للاستعلامات المستقبلية.

• دالة الخسارة (Loss Function):
  تعمل عملية التحسين على تقليل خسارتين:

  1. متوسط مربع الخطأ للمخرجات (Output MSE): $\|\hat{Y} - Y\|_F^2$، لمطابقة متجهات مخرجات الانتباه.
  2. مطابقة LSE: $\|\hat{\ell} - \ell\|_F^2$، لمطابقة لوغاريتم مجموع التعرض (log-sum-exp) لدرجات الانتباه. وهذا يضمن أن "كتلة الانتباه" المخصصة لكتلة السياق صحيحة، وهو أمر بالغ الأهمية عند الجمع مع الرموز المستقبلية.

• استراتيجية التحسين (التناوب):

  • تحسين المفاتيح (Key Optimization - $K_c$): يتم تحسينه باستخدام L-BFGS (وهي طريقة شبه نيوتنية). هذا أمر بالغ الأهمية لأن مشهد الانتباه (الذي يتضمن softmax) سلس ولكنه غير محدب ويحتوي على وديان حادة. تفشل الطرق من الدرجة الأولى مثل Adam هنا، حيث تقع في نهايات صغرى محلية سيئة.
  • تحسين القيم (Value Optimization - $V_c$): يتم حله تحليلياً بصيغة مغلقة (closed form) عبر انحدار ريج (Ridge Regression) (المربعات الصغرى) بناءً على أوزان الانتباه الثابتة الناتجة عن المفاتيح الحالية.
  • التهيئة (Initialization): يتم تهيئة المفاتيح باستخدام أفضل $k$ من المواقع بناءً على درجات أهمية الانتباه لتوفير بداية دافئة (warm start).

• بناء استعلامات التدريب:
  بما أن استعلامات فك التشفير المستقبلية غير معروفة وقت الضغط، تستخدم الطريقة:

  1. استعلامات الاحتفاظ: استعلامات حقيقية من السياق الحديث.
  2. استعلامات مستقبلية اصطناعية: متجهات المحتوى المستخرجة من السياق (عبر تفكيك De-RoPE) وإعادة تضمينها مع ترميزات موضعية مستقبلية. يوفر أخذ العينات الموحد لهذه المتجهات عبر السياق أفضل قدرة على التعميم.

• تخصيص الميزانية التكيفي:
  تختلف صعوبة الضغط بشكل كبير بين الطبقات وحتى بين رؤوس KV داخل نفس الطبقة.

  • التشغيل التجريبي (Pilot Run): عملية ضغط قصيرة ورخيصة (60 خطوة) بتخصيص موحد لتوليد إشارة متوسط مربع الخطأ (MSE) لكل مكون.
  • إعادة التوزيع: يتم إعادة توزيع الميزانية الإجمالية بما يتناسب مع الجذر التربيعي لـ MSE التجريبي ($\text{MSE}^{0.5}$). هذا يخفف من القيم المتطرفة (مثل طبقة 0 التي تمتلك خطأً أعلى بـ 100 ضعف) مع تخصيص المزيد من الأزواج للمكونات الصعبة.
  • الدقة: يتم التخصيص على مستويي الطبقة ورأس KV.

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

1. تحول النموذج المعرفي: يعيد صياغة ضغط ذاكرة KV من الاختيار/الدمج المنفصل إلى التقطير المستمر، مما يلغي البحث التوافقي عبر المواقع.
2. المُحسِّن المبتكر: يقترح مُحسّناً تناوبياً (L-BFGS + Least Squares) يحقق تباعداً (KL divergence) أقل بكثير من النماذج المرجعية القائمة على الاستبعاد.
3. التخصيص التكيفي: يقدم استراتيجية تخصيص ميزانية تكيفية قائمة على التشغيل التجريبي توزع الموارد ديناميكياً بناءً على صعوبة كل طبقة ورأس، مما يحقق مكاسب إضافية دون تكلفة استدلال.
4. التحليل التجريبي: يثبت أن صعوبة الضغط غير منتظمة للغاية (تتفاوت بما يصل إلى 467 ضعفاً بين الرؤوس في نفس الطبقة) وأن أخطاء الطبقات تتراكم بشكل كبير عبر المحول (transformer).

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

• النموذج والإعداد: تم التقييم على Qwen2.5-1.5B-Instruct مع سياقات بطول 2048 رمزاً.
• الأداء:
  • يقلل KVSCULPT تباعد KL بمقدار 3.5 إلى 4.1 ضعف مقارنة بأقوى نموذج مرجعي (Select+Fit) عبر نسب الضغط ($r \in \{0.3, 0.5, 0.7\}$).
  • في التسلسلات "السهلة"، يحقق ضغطاً يقترب من عدم الفقد (KL < 0.01)، بينما تظل النماذج المرجعية أسوأ بمرات.
  • التخصيص التكيفي يوفر تقليلاً إضافياً بمقدار 1.3 ضعف في تباعد KL.
• فعالية المُحسِّن:
  • استبدال L-BFGS بـ Adam يؤدي إلى تدهور الأداء بمقدار 17 إلى 95 ضعفاً في MSE لكل طبقة، مما يثبت ضرورة التحسين من الدرجة الثانية لمشهد الانتباه غير المحدب.
  • تصل عملية تشغيل واحدة لـ L-BFGS إلى ما ضمن 2-8% من البحث المثالي (Oracle search) (100 إعادة تشغيل عشوائية)، مما يشير إلى أن مشكلة كل طبقة قد حُلَّت فعلياً.
• تخصيص الميزانية:
  • الإشارات الثابتة (مثل معايير القيمة) تفشل في التعميم عبر التسلسلات.
  • تخصيص الميزة القائم على Pilot-MSE يقلل متوسط KL بنسبة 25% وهو فعال بشكل خاص في التسلسلات "الصعبة".
  • التخصيص لكل رأس يحسن النتائج بشكل أكبر، خاصة في الطبقات ذات التباين العالي بين الرؤوس.

5. الأهمية والتوجهات المستقبلية

• الرؤية النظرية: تثبت الورقة أن قيد "الاختيار المنفصل" هو العائق الرئيسي في طرق الضغط الحالية. السماي بحركة المفاتيح بحرية في فضاء التضمين يسمح للنموذج بـ "تلخيص" عدة رموز أصلية في رمز افتراضي واحد يلتقط تأثير انتباههم المشترك.
• تحديد الاختناق: بينما يعد التحسين لكل طبقة شبه مثالي، فإن الخطأ المتبقي في التسلسلات الصعبة ينبع من انتشار الخطأ عبر الطبقات. الاضطرابات الصغيرة في الطبقات الأولى تتراكم بمقدار 2 إلى 3 مراتب عشرية بحلول الطبقة النهائية.
• النشر: صُمم KVSCULPT لـ الإعدادات غير المتزامنة/المتراكمة (مثل ضغط مستند طويل أو نظام تعليمات قبل عدة استعلامات). تكلفة الضغط (~170 ثانية على A100 لـ 2048 رمزاً) مقبولة لهذه السيناريوهات ولكنها ليست لفك التشفير الفردي عبر الإنترنت (online single-pass decoding).
• العمل المستقبلي: يقترح المؤلفون أن التحسين المدرج (cascade-aware optimization) (حيث يأخذ هدف الطبقة $i$ في الاعتبار الأخطاء من الطبقة $i-1$) هو الخطوة المنطقية التالية لمعالجة مسألة تراكم الخطأ.

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