SLA2: Sparse-Linear Attention with Learnable Routing and QAT

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

تعالج الورقة البحثية عدم الكفاءة الحسابية في نماذج الانتشار (Diffusion Models)، ولا سيما في توليد الفيديو، حيث تتوسع آلية الانتباه الذاتي (Self-attention) بشكل تربيعي مع طول التسلسل ($O(N^2)$). وبينما تم اقتراح الانتباه الخطي المتناثر (Sparse-Linear Attention - SLA) كحل عبر الجمع بين الانتباه المتناثر (للأوزان ذات المقدار العالي) والانتباه الخطي (للأوزان ذات المقدار المنخفض)، حدد المؤلفون قصورين حرجين في صيغة SLA الأصلية:

1. عدم تطابق الصيغة (Formulation Mismatch): لا يتوافق مخرج SLA الأصلي رياضياً مع التفكيك النظري للانتباه إلى مكونات متناثرة ($P_1$) وخطية ($P_2$). وتحديداً، يقوم الفرع المتناثر بإعادة معايرة الاحتمالات، مما يؤدي إلى إدخال عامل قياس صفّي ($\alpha$) يخلق عدم تطابق. تحاول SLA تعويض ذلك عبر إسقاط قابل للتعلم، وهو أمر غير مثالي ويصعب تدريبه.
2. التوجيه الاستدلالي (Heuristic Routing): تستخدم SLA الأصلية تقسيم "Top-k" استدلالي يعتمد على مقدار أوزان الانتباه المستمدة من مصفوفات الاستعلام (Queries) والمفاتيح (Keys) المجمعة. هذا التقسيم الثابت وغير القابل للتعلم ليس مثالياً؛ فقد يوجه الحسابات إلى الفرع الخاطئ (على سبيل المثال، نقل الأوزان التي يمكن أن تكون منخفضة الرتبة إلى الفرع المتناثر)، مما يفشل في تحقيق أقصى قدر من التناثر دون تدهور الجودة.
3. حدود التكميم (Quantization Limits): لا تستفيد الطرق الحالية بشكل كامل من التكميم منخفض البت (Low-bit quantization) لتحقيق المزيد من التسريع دون إدخال تدهور كبير في الدقة.

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

اقترح المؤلفون SLA2، وهي بنية انتباه خطي-متناثر محسنة تقدم ثلاثة ابتكارات جوهرية:

أ. صيغة خطية-متناثرة أمينة (Faithful Sparse-Linear Formulation)

بدلاً من الإسقاط الجمعي في SLA الأصلية، تعيد SLA2 صياغة مخرج الانتباه ليتطابق مباشرة مع التفكيك النظري $P \approx P_1 + P_2$.

• النسبة القابلة للتعلم ($\alpha$): يتعلم النموذج متجهاً $\alpha$ (قيم بين 0 و 1) لوزن مساهمات الفروع المتناثرة والخطية ديناميكياً.
• المعادلة: يتم حساب المخرج كالتالي:
  $$O = \alpha \odot O_s + (1 - \alpha) \odot O_l$$
  حيث $O_s$ هو مخرج الانتباه المتناثر و $O_l$ هو مخرج الانتباه الخطي. تلغي هذه الصيغة الحاجة إلى طبقة إسقاط تعويضية وتضمن بقاء المخرج معايراً صفياً (Row-normalized)، مما يمنع انحراف المقدار.

ب. الموجه القابل للتعلم (Learnable Router - $R$)

لاستبدال تقسيم Top-k الاستدلالي، تقدم SLA2 موجهاً قابلاً للتعلم يقرر ديناميكياً أي الرموز (Tokens) يجب معالجتها بواسطة الفرع المتناثر.

• المدخلات: يأخذ الموجه مصفوفات الاستعلام ($\bar{Q}$) والمفاتيح ($\bar{K}$) المجمعة كمدخلات. ولا يستخدم القيم ($V$) لأن أوزان الانتباه تتحدد فقط بواسطة $Q$ و $K$.
• الآلية: يطبق الموجه عمليات إسقاط خطي قابلة للتعلم ($proj_q, proj_k$) على المدخلات المجمعة، ويحسب درجة انتباه مضغوطة، ويستخدم عامل SoftTop-k أثناء التدريب لإنشاء قناع (Mask) قابل للاشتقاق $M$.
• الكفاءة: من خلال تجميع الرموز المتجاورة واستخدام قناع مضغوط، يتجنب الموجه تعقيد $O(N^2)$، مما يجعله قابلاً للتوسع لتسلسلات الفيديو الطويلة.

ج. التدريب الواعي بالتكميم (Quantization-Aware Training - QAT)

لتحقيق تسريع إضافي، تدمج SLA2 الانتباه منخفض البت (مثل INT8/FP8) خصيصاً في الفرع المتناثر.

• الاستراتيجية: أثناء المرور الأمامي (Forward pass)، يتم تطبيق التكميم على $Q, K, P,$ و $V$. ومع ذلك، يظل المرور الخلفي (Backward pass) في صيغة FP16 باستخدام المدخلات والمخرجات الأصلية ذات الدقة الكاملة.
• الفائدة: يسمح هذا لمعلمات النموذج بالتكيف مع ضوضاء التكميم أثناء التدريب، مما يقلل من فقدان الدقة مع تمكين تسريع كبير في الاستنتاج (Inference).

د. استراتيجية التدريب ثنائية المرحلة

لضمان الاستقرار والاتساق بين التدريب والاستنتاج:

1. المرحلة 1 (التهيئة): يتم تدريب الموجه ($R$) ونسبة الخلط ($\alpha$) باستخدام خسارة متوسط مربع الخطأ (MSE) مقابل مخرجات الانتباه الكامل. يُستخدم SoftTop-k القابل للاشتقاق هنا للسماح بتدفق التدرج (Gradient flow).
2. المرحلة 2 (الضبط الدقيق): يتم ضبط نموذج الانتشار بالكامل عبر النهاية إلى النهاية باستخدام خسارة الانتشار القياسية. يتحول الموجه إلى Hard Top-k (مطابق لمنطق الاستنتاج) بينما يستمر تحسين $\alpha$.

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

1. التصحيح النظري: تحديد وإصلاح عدم التطابق الرياضي في SLA الأصلية، واقتراح صيغة متسقة مع التفكيك تلغي الحاجة إلى إسقاطات تعويض الخطأ.
2. التوجيه القابل للتعلم: استبدال تقسيم الأوزان الاستدلالي بموجه قابل للتدريب يحسن التقسيم المتناثر/الخطي بناءً على هدف صريح لتقليل الخطأ.
3. تكامل QAT: دمج التكميم منخفض البت بنجاح في الانتباه المتناثر عبر QAT، مما حقق تناثراً عالياً دون تدهور الجودة.
4. الكفاءة: أثبتت أن الجمع بين التوجيه القابل للتعلم، والصيغة الأمينة، والتكميم يحقق تسريعاً هائلاً مع الحفاظ على جودة التوليد أو تحسينها.

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

تم تقييم الطريقة على نماذج انتشار الفيديو Wan2.1 (بمعلمات 1.3B و 14B) باستخدام مجموعة بيانات فيديو خاصة.

• التناثر والسرعة:
  • حققت تناثراً في الانتباه بنسبة 97%.
  • قدمت تسريعاً بمقدار 18.6 ضعفاً في وقت تشغيل الانتباه مقارنة بالانتباه الكامل (FlashAttention).
  • في نموذج 14B، انخفض زمن انتقال توليد الفيديو من النهاية إلى النهاية بمقدار 4.35 ضعفاً.
• الجودة:
  • عند تناثر بنسبة 97%، تفوقَت SLA2 على الطرق المرجعية (بما في ذلك الانتباه الكامل) في مقاييس جودة متعددة (جودة التصوير، الجودة الجمالية، سلاسة الحركة) على نموذج 1.3B.
  • في نموذج 14B، حافظت SLA2 على جودة مماثلة للانتباه الكامل، وتفوقت بشكل ملحوظ على طرق التناثر الأخرى (SLA, VSA, VMoBA) عند مستويات تناثر أقل (90-95%).
• دراسات الاستئصال (Ablation Studies):
  • أدى إزالة QAT إلى انخفاض ملحوظ في جودة الفيديو، مما أكد ضرورته.
  • تفوق الموجه القابل للتعلم بشكل كبير على موجه Top-k الثابت المستخدم في SLA الأصلية.

5. الأهمية

تمثل SLA2 تقدماً كبيراً في توليد الفيديو عالي الكفاءة. فمن خلال الانتقال من تقسيم الانتباه الاستدلالي والثابت إلى نهج قابل للتعلم، ومبني على أسس رياضية، وواعٍ بالتكميم، فإنها تسد الفجوة بين الكفاءة النظرية والأداء العملي. وتثبت أن نماذج الانتشار يمكنها العمل بمستويات تناثر قصوى (97%) مع توفير حسابي هائل (18.6 ضعفاً) دون التضحية بالدقة العالية المطلوبة لتوليد الفيديو الاحترافي. وهذا يجعل توليد الفيديو عالي الدقة أكثر سهولة وتوافراً على الأجهزة الاستهلاكية.