Cascade Token Selection for Transformer Attention Acceleration

ملخص تقني: اختيار الرموز المتتالي (Cascade Token Selection) لتسريع انتباه المحولات (Transformer Attention)

بيان المشكلة
تواجه نماذج المحولات (Transformers) عقبة حوسبية كبيرة في آلية الانتباه الذاتي (self-attention)، حيث تتناسب تكلفة حساب حاصل الضرب $QK^\top$ وما يليه من مجموع الأوزان المرجحة بـ "softmax" طردياً مع مربع طول التسلسل ($O(T^2 d_h)$). وبينما نجحت أساليب حديثة مثل FlashAttention في تحسين عرض نطاق الذاكرة (memory bandwidth)، إلا أن التعقيد الخوارزمي لا يزال يشكل عائقاً أمام التسلسلات الطويلة.

اقترحت الورقة المصاحبة، Activation Decorrelation Attention (ADA)، حلاً لتقليل هذا التعقيد من $O(T^2)$ إلى $O(r^2)$ عبر اختيار مجموعة صغيرة من الرموز "الممثلة" (representative tokens) بعدد $r \ll T$ عند كل طبقة. تحدد ADA هذه الرموز باستخدام عتبة "غرام" (Gram threshold): حيث يُعتبر الرمز ممثلاً إذا كان أقصى ارتباط جيب تمام (cosine correlation) له مع جميع الرموز السابقة أقل من $1 - \tau^2$. ومع ذلك، فإن عملية الاختيار نفسها تتطلب حساب مصفوفة "غرام" كاملة بحجم $T \times T$ عند كل طبقة، مما يكلف $O(T^2 d)$. وبالنسبة لأطوال التسلسل المتوسطة (مثل $T=2048$)، تكون تكلفة عملية الاختيار هذه مقاربة لتكلفة الانتباه التي تهدف لتقليلها، مما يلغي الفائدة المرجوة من التسريع. المشكلة الجوهرية التي تعالجها هذه الورقة هي التكلفة الحوسبية العالية لخطوة اختيار الرموز الممثلة في نموذج ADA.

المنهجية: آلية التتالي (The Cascade Mechanism)
تقدم الورقة آلية "تتالي" تستغل التماسك العمقي (depth coherence) لمجموعة الرموز الممثلة. لاحظ المؤلفون أن مجموعة الرموز المعلوماتية في الطبقة $l$ تعد مؤشراً قوياً للطبقة $l+1$. وهذا أمر مدعوم نظرياً بخاصية "اتصال ليبشيتز" (Lipschitz continuity) لكتل المحولات: فالرموز التي تُعد نسخاً متكررة (redundant) أو غير مترابطة (representative) في طبقة ما، تميل إلى الحفاظ على تلك العلاقة في الطبقة التالية، مع حدوث اضطرابات طفيفة فقط مع إضافة المعلومات السياقية.

بدلاً من إعادة حساب مصفوفة "غرام" الكاملة $T \times T$ عند كل طبقة، تقوم آلية التتالي بما يلي:

1. توريث (Inherits) المجموعة الممثلة $R^{(l)}$ من الطبقة السابقة.
2. التحقق (Validates) من المجموعة الموروثة عبر حساب حاصلي ضرب "غرام" أصغر:
   • حاصل ضرب $r \times r$ بين الرموز الموروثة الممثلة للتحقق مما إذا كان أي منها قد أصبح زائداً عن الحاجة (redundant).
   • حاصل ضرب متقاطع $(T-r) \times r$ بين الرموز غير الممثلة والرموز الموروثة الصالحة للتحقق مما إذا كان أي من الرموز غير الممثلة قد أصبح غير مترابط (أي أصبح رمزاً ممثلاً جديداً).
3. تحديث (Updates) المجموعة بعدد صغير من الإضافات والحذف بناءً على هذه التحققات.

يقلل هذا النهج من تعقيد الاختيار من $O(T^2 d)$ إلى $O(Trd)$لكل طبقة، وهو تسريع بمقدار$T/r$.

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

• الابتكار الخوارزمي: اقتراح خوارزمية اختيار رموز متتالية تستبدل حساب مصفوفة "غرام" الكاملة بخطوات تحقق وتحديث تدريجية.
• التبرير النظري: إثبات (Proposition 3) يوضح أن التماسك العمقي للمجموعة الممثلة هو نتيجة لاتصال "ليبشيتز" لكتل المحولات والطبيعة الجمعية للروابط المتبقية (residual connections).
• الضمان التحفظي: إثبات (Proposition 4) يوضح أن طريقة التتالي هي طريقة "محافظة"؛ فهي تنتج مجموعة فائقة (superset) من المجموعة الممثلة المستقلة. إنها لا تخطئ أبداً في تحديد الرمز الممثل الحقيقي، مما يضمن عدم تدهور جودة الانتباه، وإن كان ذلك قد يزيد قليلاً من عدد الرموز الممثلة ($r$).
• التسريع المتعامد: تضع الورقة هذه الطريقة كمعادل لـ "الاستدلال في الفضاء الفرعي الموجه" (Gated Subspace Inference - GSI) في بُعد الرموز، والذي يسرع الطبقات الخطية عبر وراثة الأساس المتماسك عمقياً. معاً، يغطيان عملية المرور الأمامي (forward pass) للمحول بالكامل (محوري $d$ و $T$).

النتائج التجريبية
تم التحقق من صحة الطريقة على ثلاث عائلات من النماذج (GPT-2 124M، GPT-J 6B، OPT 6.7B) باستخدام أجهزة AMD MI300X بطول تسلسل $T=512$ وعتبة غرام $\tau=0.30$.

• التماسك العمقي: كان مؤشر "جاكارد" (Jaccard index) بين المجموعات الممثلة في الطبقات المتتالية مرتفعاً (من 0.83 إلى 0.94) في الطبقات العميقة، مما يؤكد استقرار مجموعة الرموز.
• توفير عمليات "غرام": قللت آلية التتالي من إجمالي عمليات "غرام" بنسبة:
  • 22% لنموذج GPT-2 (12 طبقة).
  • 52% لنموذج GPT-J (28 طبقة).
  • 63% لنموذج OPT (32 طبقة).
    يزداد التوفير مع زيادة عمق النموذج لأن تكلفة "غرام" الكاملة الأولية عند الطبقة 0 يتم توزيعها (amortized) على المزيد من الطبقات.
• معدل التغيير (Turnover Rate): في الطبقات العميقة، ينخفض معدل التغيير (نسبة الجزء المعدل من المجموعة) إلى 3–18%، مما يشير إلى أن المجموعة الموروثة تتطلب تعديلاً طفيفاً.
• توقعات القياس (Scaling Predictions): يتوقع المؤلفون أنه بالنسبة للتسلسلات الأطول (مثل $T=8192$)، حيث ينمو $r$ بشكل دون خطي، يمكن أن يصل تسريع الاختيار إلى $18\times$، مما يحول تكاليف الاختيار من مصطلح مهيمن إلى مصطلح مهمل.

الأهمية والادعاءات
تدعي الورقة أن آلية التتالي تكشف عن خاصية هيكلية في المحولات المدربة: وهي أن مجموعة الرموز المعلوماتية هي خاصية هيكلية متماسكة تنتقل عبر عمق الشبكة. ومن خلال استغلال هذا التماسك، تجعل هذه الطة عملية اختيار ADA عملية عملية، مما يسمح لضغط الانتباه بتحقيق تسريعه النظري دون أن يعيقها خطوة الاختيار.

يؤكد المؤلفون أن هذه الطريقة هي تسريع مكمل للأساليب الموجودة مثل FlashAttention (الذي يحسن عرض نطاق الذاكرة) والانتباه المتفرق (الذي يستخدم أنماطاً ثابتة). وخلافاً لطرق تقليم الرموز (token pruning) أو دمجها التي تغير التمثيلات بشكل دائم، يقوم "التتالي في ADA" باختيار الممثلين ديناميكياً بناءً على بنية المدخلات دون تغييرات معمارية أو إعادة تدريب. إن الجمع بين GSI و Cascade ADA يقدم تقليلاً مضاعفاً لتكلفة الاستدلال، مما يقلل فعلياً من تعقيد المرور الأمامي إلى "التعقيد الفعال" للمدخلات بدلاً من الحجم الاسمي للنموذج.