YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference

إليك ملخص تقني مفصل لورقة البحث بعنوان "YOCO++: تعزيز YOCO عبر اتصالات متبقية لـ KV لرفع كفاءة استنتاج النماذج اللغوية الكبيرة (LLM)."

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

تواجه النماذج اللغوية الكبيرة (LLMs) القائمة على بنية Transformer عقبة كبيرة أثناء عملية الاستنتاج (Inference): وهي ذاكرة المفتاح-القيمة (KV cache). فمع زيادة أطوال التسلسلات، يصبح حجم الذاكرة المطلوبة لتخزين KV cache عاملاً مقيداً لحجم الدفعة (Batch size) ومعدل الإنتاجية (Throughment).

ولمعالجة ذلك، تم اقتراح طرق للضغط عبر الطبقات (Cross-layer KV compression). تعمل هذه الطرق على تقليل استهلاك الذاي عبر الاحتفاظ بـ KVs لعدد فرعي فقط من الطبقات وإعادة استخدامها في طبقات أخرى. ومن أبرز هذه الطرق هي YOCO (You Only Cache Once)، التي تشارك الـ KVs للطبقة الوسطى مع جميع طبقات النصف العلوي، مما يحقق تقليلاً بنسبة 50% في حجم KV cache.

المشكلة الجوهرية: بينما يحسن المشاركة عبر الطبقات من الكفاءة، إلا أنها تؤدي إلى تدهور غير ضئيل في الأداء. إن مشاركة الـ KVs مباشرة من طبقة إلى أخرى تضعف قدرة النموذج لأن الطبقات المتميزة تفقد قدرتها على تعلم تمثيلات فريدة. وتحاول الحلول الحالية مثل FusedKV التخفيف من ذلك باستخدام تركيبات موزونة من KVs لعدة طبقات، ولكنها تتسبب في أعباء إضافية على عمليات الإدخال/الإخراج (I/O overhead)، مما يقلل من إنتاجية الاستنتاج.

2. المنهجية: YOCO++

يقترح المؤلفون YOCO++، وهو نسخة محسنة من YOCO تعمل على تحسين أداء النموذج دون التضحية بكفاءة الاستنتاج أو إدخال أعباء إضافية على الإدخال/الإخراج.

الآلية الجوهرية: الاتصالات المتبقية الموزونة لـ KV

بدلاً من مشاركة الـ KVs مباشرة، يقدم YOCO++ اتصالاً متبقياً موزوناً (Weighted residual connection) بين الـ KVs للطبقة السفلية والـ KVs لكل طبقة لاحقة في النصف السفلي.

• البنية: يتكون النموذج من $L$ من الطبقات. النصف السفلي ($L/2$) عبارة عن "محللات ذاتية" (Self-decoders)، والنصف العلوي عبارة عن "محللات متقاطعة" (Cross-decoders).
• تركيب الـ KV: في كل طبقة من النصف السفلي $i$ (حيث $i \in \{2, ..., L/2\}$)، يتم حساب الـ KVs المخزنة ($K'_i, V'_i$) كمجموع موزون لـ KVs الطبقة الحالية ($K_i, V_i$) و KVs الطبقة الأولى ($K_1, V_1$):
  $$K'_i = \lambda(\alpha^K_{i,1}K_1 + \alpha^K_{i,2}K_i)$$
  $$V'_i = \lambda(\alpha^V_{i,1}V_1 + \alpha^V_{i,2}V_i)$$
  حيث أن $\alpha$ هي أوزان متبقية قابلة للتعلم، و $\lambda$ هو عامل قياس.
• إعادة الاستخدام في النصف العلوي: تعيد طبقات النصف العلوي استخدام الـ KVs المعدلة من الطبقة الوسطى ($K'_{L/2}, V'_{L/2}$)، والتي تحمل الآن معلومات أغنى من الـ KVs الأصلية لطبقة YOCO الوسطى.

الابتكارات التقنية الرئيسية

1. استراتيجية التخزين المؤقت (Cache Strategy): لتجنب عبء الإدخال/الإخراج الذي يسببه FusedKV (والذي يتطلب جلب KVs لطبقات متعددة أثناء فك التشفير)، يقوم YOCO++ بتخزين الـ KVs المركبة ($K', V'$) بدلاً من الأصلية. وهذا يضمن أنه خلال مرحلة فك التشفير (Decoding)، يصل النموذج فقط إلى مجموعة واحدة من KVs لكل طبقة، مما يحافظ على نفس ملف الإدخال/الإخراج الخاص بـ YOCO.
2. عامل القياس ($\lambda$): لاحظ المؤلفون أنه بدون عامل قياس، يعاني النموذج في تعلم الأوزان المتبقية (حيث تظل قريبة من قيم التهيئة). لذا، قدموا ثابتًا (Hyperparameter) $\lambda > 1$ (تم ضبطه تجريبياً عند 35) لتضخيم الإشارة المتبقية، مما يسهل التعلم الأفضل للأوزان.
3. التهيئة (Initialization): تم تهيئة الأوزان بحيث تكون $\alpha_{i,1}=0$ و $\alpha_{i,2}=1/\lambda$، مما يضمن بدء النموذج في حالة مطابقة تماماً لـ YOCO القياسي.

3. تحليل الكفاءة

حللت الورقة بدقة كفاءة YOCO++ مقارنة بـ YOCO و FusedKV:

• المعلمات (Parameters): زيادة طفيفة جداً (فقط 4 أوزان قابلة للتعلم لكل طبقة في النصف السفلي).
• الحوسبة: عبء ضئيل؛ فالتركيب الموزون عملية بسيطة جداً مقارنة بآلية الانتباه (Attention mechanism).
• الذاكرة: مطابقة لـ YOCO. كلاهما يخزن $L/2$ من الطبقات.
• عبء الإدخال/الإخراج (I/O Overhead):
  • التدريب/التحضير (Prefill): إدخال/إخراج أعلى قليلاً بسبب الوصول إلى $K_1, V_1$ و $K_i, V_i$، ولكن هذه المرحلة تعتمد على الحوسبة (Compute-bound)، لذا لا يتأثر زمن الاستجابة (Latency).
  • فك التشفين (Decoding): لا يوجد عبء إضافي. بما أن الـ KVs المركبة يتم حسابها مسبقاً وتخزينها، فإن فك التشفين يقرأ فقط مجموعة واحدة من KVs لكل طبقة، مما يطابق إنتاجية YOCO.

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

أُجريت التجارب على نموذج بـ 1.1 مليار معلمة (إعداد TinyLlama) باستخدام مجموعة بيانات SlimPajama.

• كفاءة الاستنتاج (الشكل 3):

  • زمن استجابة التحضير (Prefill Latency): يطابق YOCO++ نموذج YOCO و FusedKV، محققاً زمن استجابة أقل بنسبة 50% من Transformer القياسي بفضل الخروج المبكر.
  • إنتاجية فك التشفين (Decoding Throughput): يحافظ YOCO++ على إنتاجية عالية مقاربة لـ YOCO ويتفوق بشكل كبير على Transformer القياسي. وعلى عكس FusedKV، الذي تنخفض إنتاجيته إلى مستويات Transformer بسبب اختناقات الإدخال/الإخراج، يظل YOCO++ فعالاً.

• أداء النموذج (الجدول 1 و 2):

  • دقة الصفر (Zero-Shot Accuracy): يحقق YOCO++ أداءً هو الأفضل حالياً (SOTA) بين طرق الضغط عبر الطبقات عند معدل ضغط 50%.
  • المقارنة: يتفوق على Transformer القياسي، و YOCO، و FusedKV، و FusedKV-Lite في اختبارات الاستدلال المنطقي (HellaSwag, OBQA, WinoGrande, ARC, PIQA).
  • خسارة التدريب (Training Loss): يحقق YOCO++ أدنى خسارة تحقق (2.4194)، متفوقاً حتى على Transformer القياسي (2.4266).

• دراسات الاستئصال (Ablation Studies):

  • عامل القياس: إزالة $\lambda$ تؤدي إلى انخفاض الأداء إلى مستويات YOCO، مما يؤكد ضرورته للتعلم.
  • المتبقيات المفتاحية (Key Residuals): إزالة اتصال المفتاح المتبقي ($K$) تؤدي إلى تدهور الأداء، مما يثبت أن الاتصالات المتبقية لكل من المفاتيح (Keys) والقيم (Values) مفيدة.

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

1. بنية YOCO++: طريقة مبتكرة تعزز ضغط KV عبر الطبقات من خلال إدخال اتصالات متبقية موزونة بين الطبقة السفلية وطبقات النصف السفلي.
2. الحفاظ على الكفاءة: نجح في إزالة عبء الإدخال/الإخراج المرتبط بطرق "الدمج" السابقة (مثل FusedKV) عن طريق تخزين التمثيل المدمج، مما يحافظ على إنتاجية YOCO العالية.
3. عامل القياس: تقديم عامل قياس $\lambda$ لحل مشكلة صعوبة تحسين الأوزان المتبقية في الشبكات العميقة.
4. أداء SOTA: أثبت أنه من الممكن ضغط KV caches بنسبة 50% مع التفوق على Transformer القياسي غير المضغوط في كل من خسارة التدريب ودقة المهام النهائية.

6. الأهمية

تعالج هذه الورقة المقايضة الحرجة بين كفاءة استنتاج LLM وقدرة النموذج. ومن خلال إظهار أن الضغط عبر الطبقات يمكن أن لا يكتفي بمطابقة Transformer القياسي بل يتفوق عليه دون تكبد عقوبات في الإدخال/الإخراج، يوفر YOCO++ حلاً عملياً للغاية لنشر النماذج الكبيرة على الأجهزة محدودة الموارد. إنه يشير إلى أن "الاتصالات المتبقية" يمكن تطبيقها بفعالية في مجال KV cache للحفاظ على تدفق المعلومات عبر الطبقات، مما يفتح اتجاهاً جديداً لتصميم نماذج لغوية كبيرة فعالة.