Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

إليك ملخص تقني مفصل لورقة البحث بعنوان: "Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection" (انتباه التوكن المتناثر: استنتاج فعال للسياقات الطويلة باستخدام اختيار التوكنات المتداخل).

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

تواجه النماذج اللغوية الكبيرة (LLMs) عائقاً جوهرياً في استنتاج السياق الطويل بسبب التعقيد التربيعي ($O(L^2)$) لآلية الانتباه خلال مرحلة الـ prefill. وبينما تعمل تحسينات الأجهزة مثل FlashAttention على تقليل عبء إدخال/إخراج الذاكرة (I/O)، إلا أنها لا تغير التعقيد الخوارزمي.

تعاني طرق التسريع الحالية من محدوديتين رئيسيتين:

• التناثر الهيكلي (Structured Sparsity): تطبق طرق مثل Minference أو FlexPrefill التناثر على مستوى الكتل (block-level). وغالباً ما يؤدي هذا إلى الاحتفاظ بتوكنات غير ذات صلة إذا تم تجميعها مع توكنات مهمة في نفس الكتلة، مما يؤدي إلى تناثر غير مثالي.
• طرد التوكنات (Token Eviction): تقوم طرق مثل FastKV أو GemFilter بطرد التوكنات بشكل دائم في الطبقات الأولى. وهذا يفترض أن أهمية التوكن ثابتة، متجاهلاً الطبيعة الديناميكية لأهمية التوكن عبر الطبقات ورؤوس الانتباه المختلفة. وبمجرد طرد التوكن، لا يمكن إعادة النظر فيه، مما يؤدي إلى تدهور الدقة.

2. المنهجية: انتباه التوكن المتناثر (Token Sparse Attention)

يقترح المؤلفون Token Sparse Attention، وهو آلية خفيفة الوزن، وديناميكية، وعكوسة لتناثر التوكنات على مستوى التوكن. بخلاف طرق الطرد، لا تقوم هذه الطريقة بإزالة التوكنات نهائياً، بل تضغط مساحة الحوسبة ثم تستعيد التسلسل.

البنية الأساسية

تعمل الطريقة على مرحلتين لكل رأس انتباه (attention head):

1. الضغط (اختيار QKV):
   • لكل رأس انتباه $h$، يتم اختيار مجموعة فرعية من مؤشرات التوكنات $S_h$ بناءً على درجات الأهمية.
   • يتم تجميع ت tensors الخاصة بالاستعلام (Query - $Q$)، والمفتاح (Key - $K$)، والقيمة (Value - $V$) لتشكيل ت tensors مضغوطة $\hat{Q}, \hat{K}, \hat{V}$ بطول تسلسل مختزل $L' \ll L$.
   • الميزة الرئيسية: يتم إجراء الاختيار بشكل مستقل لكل رأس، مما يسمح للرؤوس المختلفة بالتركيز على مجموعات فرعية مختلفة من التوكنات، مما يعالج عدم التجانس بين الرؤوس (head-wise heterogeneity).
   • تظل الت tensors المضغوطة كثيفة ومتصلة، مما يتيح استخدام نوى (kernels) محسنة مثل FlashAttention.
2. إلغاء الضغط (استعادة المخرجات):
   • يتم توزيع مخرج الانتباه $\hat{O}$ (بشكل $L' \times d$) مرة أخرى في ت tensor مهيأ بالأصفار بطول التسلسل الأصلي $L \times d$ باستخدام المؤشرات الخاصة بكل رأس.
   • الميزة الرئيسية: يتم تصفير المواقع غير المختارة، ولكن الاتصال المتبقي (residual connection) من الطبقة السابقة يحافظ على معلومات التوكنات التي تم تخطيها. وهذا يسمح للنموذج بـ "إعادة زيارة" وإعادة تقييم أهمية التوكن في الطبقات اللاحقة.

التغطية الديناميكية للتوكنات (Dynamic Token Coverage)

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

• التسجيل (Scoring): يتم حساب خريطة انتباه وسيطة خفيفة الوزن باستخدام الاستعلامات الأخيرة مقابل جميع المفاتيح. تُشتق درجات أهمية التوكن عن طريق جمع أوزان الانتباه رأسياً.
• تخصيص الميزانية (Budget Allocation): بدلاً من الاحتفاظ بأفضل-$k$ من التوكنات، تقوم الطريقة بفرز التوكنات حسب الأهمية التصاعدية وتقليم "الذيل" (التوكنات الأقل أهمية) حتى يتجاوز الكتلة التراكمية للتوكنات المقطوعة عتبة التغطية $\tau$. هذا يعامل التقليم كنوع من التنظيم الهيكلي ضد ضجيج الانتباه.
• الاختيار على مستوى الرأس: بعد تحديد الميزانية العالمية، يختار كل رأس بشكل مستقل أفضل-$k$ من توكناته بناءً على درجاته الخاصة.

اختيار الطبقة المتناثرة (Sparse Layer Selection)

لاحظ المؤلفون أن تطبيق التناثر على جميع الطبقات يؤدي إلى تدهور الأداء. لذا قدموا مقياس انزياح التمثيل بين الطبقات ($R_\ell$) لقياس استقرار تمثيلات التوكنات بين الطبقات.

• الطبقات ذات الانزياح المنخفض (التمثيلات المستقرة) يتم اختيارها للتناثر.
• الطبقات ذات الانزياح العالي (حيث تتغير التمثيلات بشكل كبير) يتم الإبقاء عليها كثيفة للحفاظ على سلوك النموذج.
• هذه عملية اختيار تتم كخطوة معالجة مسبقة لمرة واحدة لكل نموذج.

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

1. تناثر التوكنات العكوس (Reversible Token-Level Sparsification): تصميم مبتكر "للضغط وإلغاء الضغط" يقلل من تكلفة الحوسبة دون طرد التوكنات نهائياً، مما يسمح بإعادة التقييم الديناميكي لأهمية التوكن عبر الطبقات.
2. عدم التجانس على مستوى الرأس (Head-Wise Heterogeneity): تدعم الطريقة الاختيار المستقل للتوكنات لكل رأس انتباه، مما يلتقط حقيقة أن الرؤوس المختلفة تعطي الأولوية لميزات دلالية مختلفة.
3. التوافق السلس (Seamless Compatibility): هذه الطريقة متوافقة تماماً مع نوى (kernels) الانتباه الكثيفة (مثل FlashAttention) والمتناثرة (مثل Minference و FlexPrefill)، حيث تعمل كمسرع مكمل وليس كبديل.
4. الميزانية الديناميكية (Dynamic Budgeting): نهج قائم على البيانات لتحديد مستويات التناثر بناءً على توزيعات أهمية التوكنات واستقرار الطبقات، بدلاً من النسب الثابتة.

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

أُجريت التجارب على نماذج LLaMA-3.1-8B-Instruct و Mistral-Nemo-12B-Instruct باستخدام اختبارات معيارية مثل RULER و InfiniteBench و Needle-in-a-Haystack.

• المقايضة بين الدقة والزمن (Accuracy-Latency Trade-off):
  • حقق Token Sparse Attention باستمرار تحسيناً في المقايضة بين السرعة والدقة.
  • عند سياق طوله 128K، يحقق تسريعاً في الانتباه يصل إلى $\times 3.23$ مقارنة بـ FlashAttention مع أقل من 1% تدهور في الدقة.
  • عند دمجه مع FlexPrefill، فإنه يحسن التسريع من $\times 2.44$ إلى $\times 2.76$ مع الحفاظ على نفس متوسط الدقة (87.27%).
• المقارنة مع طرق الطرد (Comparison with Eviction Methods):
  • يتفوق بشكل كبير على طرق طرد التوكنات (PyramidInfer, FastKV, GemFilter). على سبيل المثال، عند 128K، يحقق Token Sparse Attention دقة 86.84% مقابل 78.49% لـ PyramidInfer عند نفس مستويات التسريع.
• القابلية للتوسع (Scalability):
  • تزداد مكاسب التسريع مع طول السياق (على سبيل المثال، مكاسب ضئيلة عند 4K، ومكاسب كبيرة عند 128K) لأن تناثر الانتباه يزداد مع طول التسلسل.
• العبء الإضافي (Overhead):
  • العبء الإضافي لعمليات التسجيل وإدارة المؤشرات ضئيل، حيث يمثل أقل من 11% من زمن انتقال الانتباه الإجمالي حتى عند مستويات تناثر عالية.

5. الأهمية

تعالج هذه الورقة فجوة حرجة في استنتاج النماذج اللغوية الكبيرة ذات السياق الطويل من خلال الانتقال من الطرد غير العكوس للتوكنات إلى التناثر المتداخل والديناميكي.

• الرؤية النظرية: تثبت أن أهمية التوكن ديناميكية للغاية عبر الطبقات والرؤوس، مما يستلزم آليات اختيار مرنة بدلاً من التقليم الثابت.
• الأثر العملي: توفر تحسيناً "جاهزاً للتشغيل" (plug-and-play) يمكن وضعه فوق نوى الانتباه الحالية (FlashAttention, Minference, إلخ) لتحقيق تسريع هائل في سيناريوهات السياق الطويل دون الحاجة لإعادة تدريب النماذج أو التضحية بدقة كبيرة.
• الاتجاه المستقبلي: تضع نقطة تصميم جديدة للاستنتاج القابل للتوسع، مما يشير إلى أن الضغط العكوس هو استراتيجية متفوقة للتعامل مع معضلة التعقيد $O(L^2)$ مقارنة بالإزالة الدائمة للتوكنات.