ملخص تقني: EntmaxKV – فك التشفير الواعي بالدعم لآلية انتباه Entmax

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

يواجه فك التشفير في السياقات الطويلة (Long-context decoding) في النماذج اللغوية الكبيرة (LLMs) اختناقاً متزايداً ناتجاً عن حركة مرور الذاكرة (memory traffic) بدلاً من الحوسبة. أثناء التوليد التكراري (autoregressive generation)، يجب على كل رمز جديد الانتباه إلى ذاكرة مفتاح-قيمة (KV cache) تنمو خطياً مع طول السياق. إن قراءة هذه الذاكرة بالكامل من الذاكرة عالية النطاق الترددي (HBM) إلى ذاكرة SRAM الموجودة على الشريحة في كل خطوة يهيمن على زمن الاستجابة (latency)، وغالباً ما يستحوذ على أكثر من نصف وقت توليد الرمز الواحد.

تحاول طرق فك التشفير المتناثرة (sparse decoding) الحالية التخفيف من حدة ذلك عبر اختيار مجموعات فرعية من الرموز أو الكتل أو الصفحات. ومع ذلك، فإن هذه الطرق مصممة أساساً لآلية انتباه softmax. وتتمثل إحدى القيود الجوهرية لـ softmax في "ذيلها الكثيف" (dense tail): فهي تخصص احتمالية موجبة تماماً لكل رمز، بغض النظر عن مدى أهميته. وبناءً على ذلك، فإن أي اقتطاع لذاكرة KV تحت آلية softmax يؤدي حتماً إلى التخلص من كتلة احتمالية غير صفرية، مما يجعل فك التشفير المتناثر مشكلة تقريبية (approximation problem) حيث يكون الخطأ أمراً لا مفر منه ما لم تُقرأ الذاكرة بالكامل.

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

يقترح المؤلفون EntmaxKV، وهو إطار عمل لفك التشفير المتناثر المصمم خصيصاً لآلية انتباه $\alpha$-entmax. وخلافاً لـ softmax، تنتج $\alpha$-entmax (للقيمة $\alpha > 1$) توزيعات احتمالية ذات أصفار حقيقية. هذه الخاصية تحول هدف فك التشفير المتناثر من "تقريب ذيل كثيف" إلى "استعادة الدعم" (support recovery): فإذا احتوت المجموعة المختارة من الرموز على "الدعم الحقيقي" (وهو مجموعة الرموز ذات الاحتمالية غير الصفرية)، فإن مخرجات فك التشفير المتناثر ستكون متطابقة رياضياً مع مخرجات الذاكرة الكاملة.

ينفذ EntmaxKV ذلك عبر مسار فك تشفير مجزأ إلى صفحات (paged decoding pipeline) يتكون من ثلاثة مكونات أساسية:

أ. تسجيل درجات الصفحات الواعية بالاستعلام (Query-Aware Page Scoring)

يتم تقسيم ذاكرة KV إلى صفحات. وقبل تحميل مفاتيح وقيم كاملة، يقوم النظام بتسجيل درجات الصفحات باستخدام بيانات وصفية خفيفة الوزن مخزنة لكل صفحة:

• حدود الصندوق المحيط (Bounding Box Bounds): تسمح قيم الحد الأدنى والأقصى لمتجهات المفتاح ($k_{min}, k_{max}$) لكل إحداثي بوضع حد علوي حتمي لأقصى درجة انتباه ممكنة لأي رمز في الصفحة بناءً على استعلام $q$.
• الإحصاءات الغاوسية (Gaussian Statistics): لاعتماد نهج أكثر تكيفاً، يخزن النظام العزوم الأولى والثانية (المتوسط والتباين) للمفاتيح داخل الصفحة. يتيح هذا تقريب توزيع درجات الرموز في الصفحة كمتغير عشوائي غاوسي.

ب. سياسات اختيار المرشحين (Candidate Selection Policies)

تم اقتراح سياستين لاختيار الصفحات التي يجب تحميلها:

1. اختيار الصفحات الأعلى-k (Top-k Page Selection): نهج بميزانية ثابتة يختار الـ $k$ من الصفحات ذات أعلى درجات الحد العلوي المستمدة من بيانات الصندوق المحيط.
2. اختيار Entmax الواعي بالغاوس (Gaussian-Aware Entmax Selection): نهج ديناميكي يقدر عتبة $\alpha$-entmax العالمية ($\tau$) باستخدام الإحصاءات الغاوسية على مستوى الصفحة. يقوم بحل معادلة تطبيع تقريبية لإيجاد $\hat{\tau}$ ثم يختار الصفحات التي تتجاوز فيها الدرجة القصوى المقدرة هذه العتبة (مع تعديلها بهامش أمان $\Delta$). هذا يكيف ميزانية الرموز مع توزيع الدرجات الحالي للاستعلام.

ج. انتباه Entmax المتناثر (Sparse Entmax Attention)

بمجرد اختيار صفحات المرشحين، يقوم النظام بتحميل مدخلات KV المقابلة فقط. ثم يحسب أوزان $\alpha$-entmax فوق هذه المجموعة المختزلة. إذا احتوت المجموعة المختارة على الدعم الكامل لتوزيع $\alpha$-entmax، فإن المخرجات تكون دقيقة.

3. التحليل النظري

يوفر البحث إطاراً نظرياً للخطأ الناتج عن الاقتطاع:

• حد الخطأ (Error Bound): يتم تحديد الخطأ في المخرجات $\|o - \tilde{o}\|$ بـ $2B\delta$، حيث $B$ هو حد معايير متجهات القيم، و$\delta$ هي إجمالي الكتلة الاحتمالية للرموز المستبعدة.
• استعادة الدعم (Support Recovery): بالنسبة لـ softmax، تكون $\delta > 0$ دائماً طالما تم استبعاد أي رمز، مما يضمن خطأً غير صفري. أما بالنسبة لـ $\alpha$-entmax، فإذا احتوت المجموعة المختارة $I_{keep}$ على الدعم $S$ (أي $S \subseteq I_{keep}$)، فإن $\delta = 0$ ويختفي الخطأ تماماً.
• المقارنة: يوضح البحث أن $\alpha$-entmax تفقد كتلة احتمالية أقل بكثير من softmax تحت نفس المجموعة المحتفظ بها، لا سيما عندما تنجح مجموعة المرشحين في التقاط دعم entmax.

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

أُجريت التجارب باستخدام نماذج بقدر 1 مليار معلمة (NAPE) على اختبارات السياق الطويل (RULER, Passkey Retrieval, PG19) واختبارات الكفاءة.

• جودة التقريب: يتفوق EntmaxKV باستمرار في تقليل كتلة الاحتمالية المحذوفة والاحتفاظ برموز الدعم أكثر من فك التشفير المتناثر القائم على softmax عند موازنة ميزانيات ذاكرة KV. الخطأ النسبي في المخرجات أقل بكثير لـ entmax، حيث يقترب من الصفر مع استعادة الدعم.
• الأداء في المهام النهائية (Downstream Performance):
  • استرجاع المفتاح (Passkey Retrieval): يحافظ EntmaxKV على دقة استرجاع مثالية عبر سياقات طويلة (تصل إلى 100 ألف رمز) باستخدام أقل من 0.1% من ذاكرة KV، بينما تتدهور الطرق القائمة على softmax بسرعة.
  • اختبار RULER: يتفوق EntmaxKV بشكل كبير على فك التشفير المتناثر لـ softmax، خاصة عند أطوال سياق تبلغ 16 ألف رمز، محافظاً على أداء قريب من نماذج الذاكرة الكاملة.
  • النمذجة اللغوية (PG19): يحقق فك التشفير المتناثر لـ entmax مع ميزانية 1024 رمزاً تشتتاً (perplexity) أقل من فك تشفير softmax بميزانيات أكبر بكثير (مثل 8192 رمزاً).
• الكفاءة:
  • ينقل EntmaxKV الاختناق من حركة مرور ذاكرة الذاكرة الكاملة إلى اختيار الصفحات الخفيف.
  • عند طول سياق يبلغ 1 مليون رمز، يحقق المتغير الواعي بالغاوس تسريعاً قدره 3.36 ضعفاً مقارنة بـ softmax (FlashDecoding) و5.43 ضعفاً مقارنة بفك تشفير entmax الكامل.
  • يتوسع المختار الواعي بالغاوس بشكل أفضل مع طول التسلسل من متغير Top-k، متجنباً اختناق Top-k الصريح.

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

يزعم البحث أن EntmaxKV يغير جذرياً نموذج فك التشفير المتناثر عبر الاستفادة من التناثر الحقيقي لـ $\alpha$-entmax.

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

يضع المؤلفون EntmaxKV كحل أصيل لنماذج $\alpha$-entmax، متميز عن الأعمال السابقة التي تحاول فرض التناثر على توزيعات softmax الكثيفة. ويؤكدون أنه بينما جعلت الـ kernels الحديثة $\alpha$-entmax قابلة للتدريب على نطاق واسع، فإن EntmaxKV هو الأول الذي يعالج اختناق فك التشفير التكراري عبر تمكين استنتاج متناثر يتسم بالكفاءة والدقة النظرية تحت استعادة الدعم.