SpargeAttention2: Trainable Sparse Attention via Hybrid Top-k+Top-p Masking and Distillation Fine-Tuning

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

تعاني نماذج انتشار الفيديو (Video diffusion models) من تكاليف حوسبة عالية بسبب التعقيد الحسابي بمقدار $O(N^2)$ لآلية الانتباه (attention mechanism)، خاصة بالنظر إلى أطوال التسلسلات الطويلة المتأصلة في توليد الفيديو. وبينما توجد طرق انتباه متفرقة (sparse attention) للتخفيف من ذلك، تواجه الأساليب الحالية ثلاثة تحديات حرجة:

1. فشل قواعد التغطية (Masking) القياسية عند مستويات التفرع العالية: تعتمد الطرق الحالية إما على Top-k (الاحتفاظ بأفضل $k$ من الرموز/tokens) أو Top-p (الاحتفاظ بالرموز حتى تصل الاحتمالية التراكمية إلى $p$).
   • فشل Top-k: في الصفوف ذات توزيعات الاحتمالية الموحدة، يحتفظ Top-k بعدد قليل جداً من الرموز، مما يؤدي إلى استبعاد سياق قيم.
   • فشل Top-p: في الصفوف ذات التوزيعات شديدة الانحراف (على سبيل المثال، التي تهيمن عليها "بالوعات الانتباه" أو attention sinks)، قد يستوفي Top-p العتبة باستخدام عدد قليل فقط من الرموز غير المعلوماتية، مما يسقط سياقاً حيوياً.
2. أهداف الضبط الدقيق (Fine-Tuning) غير المثالية: تقوم معظم طرق الانتباه المتفرقة القابلة للتدريب بضبط النماذج باستخدام خسارة الانتشار القياسية على بيانات من العالم الحقيقي. ومع ذلك، غالباً ما تفتقر نماذج الفيديو مفتوحة المصدر إلى الوصول إلى مجموعات بيانات التدريب الأصلية عالية الجودة. إن الضبط الدقيق على بيانات منخفضة الجودة أو متباينة التوزيع باستخدام خسارة الانتشار القياسية يسبب انحراف السلوك (behavior drift)، مما يؤدي إلى تدهور جودة التوليد حتى عند استخدام الانتباه الكامل.
3. المفاضلة بين الكفاءة والجودة: يؤدي تحقيق تفرع عالٍ جداً (على سبيل المثال، >90%) عادةً إلى تدهور كبير في الجودة البصرية والتماسك الزمني.

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

يقترح المؤلفون SpargeAttention2، وهو إطار عمل للانتباه المتفرق القابل للتدريب مصمم لتحقيق تفرع عالٍ دون المساس بجودة التوليد. ويتكون من ثلاثة مكونات أساسية:

أ. التغطية الهجينة (Hybrid Top-k + Top-p Masking)

لمعالجة أوضاع الفشل لقواعد التغطية الفردية، يقترح المؤلفون استراتيجية تغطية موحدة تجمع بين Top-k وTop-p.

• الآلية: لكل صف من مصفوفة أوزان الانتباه، تكون التغطية هي اتحاد (union) الرموز المختارة بواسطة Top-k وتلك المختارة بواسطة Top-p.
  • إذا كان التوزيع موحداً، يضمن Top-p الاحتفاظ بعدد كافٍ من الرموز لالتقاط الكتلة الاحتمالية.
  • إذا كان التوزيع منحرفاً، يضمن Top-k عدم هيمنة "بالوعة الانتباه" على الاختيار، مما يجبر النظام على تضمين رموز معلوماتية أخرى.
• النتيجة: يوفر هذا النهج الهجين أداء تغطية قوياً عبر مختلف توزيعات أوزان الانتباه، مما يحافظ على الدقة حتى عند مستويات التفرع القصوى.

ب. ضبط التقطير للسرعة (Velocity Distillation Fine-Tuning)

لحل مشكلة عدم تطابق توزيع البيانات وانحراف السلوك، يستبدل المؤلفون خسارة الانتشار المعتمدة على البيانات التقليدية بـ خسارة تقطير السرعة (Velocity Distillation Loss).

• إعداد المعلم والطالب (Teacher-Student Setup):
  • المعلم (Teacher): نموذج الانتباه الكامل المدرب مسبقاً (مجمد).
  • الطالب (Student): النموذج الذي يتم تكييفه لاستخدام SpargeAttention2.
• الهدف: يتم تدريب نموذج الطالب لمطابقة توقعات السرعة ($u_{sparse}$) لنموذج المعلم ($u_{full}$) بناءً على نفس مدخلات الكمون الضوضائية ($x_t$)، والخطوة الزمنية ($t$)، والمطالبات النصية ($c_{text}$).
• دالة الخسارة:
  $$L_{VD} = \mathbb{E}[\| u_{sparse}(x_t, c_{text}, t) - u_{full}(x_t, c_{text}, t) \|^2]$$
• الفائدة: يعمل هذا النهج على فصل عملية التكيف عن جودة مجموعة بيانات الضبط الدقيق. يتعلم النموذج الحفاظ على ديناميكيات التوليد الأصلية للمعلم ذو الانتباه الكامل، مما يمنع التدهور حتى عندما تكون بيانات الضبط الدقيق محدودة أو أقل جودة.

ج. تنفيذ النواة الفعال (Efficient Kernel Implementation)

يستخدم الأسلوب تنفيذ انتباه متفرق كتلوي (block-sparse attention) مبني على FlashAttention.

• يقوم بتقسيم الموترات إلى بلاطات (tiles) ويطبق التغطية الهجينة على مستوى الكتلة.
• تتخطى نواة CUDA عمليات ضرب المصفوفات وحسابات softmax بكفاءة، مما يضمن تسريعاً عملياً على وحدات معالجة الرسومات الحديثة.

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

1. التحليل النظري: حدد المؤلفون وحللوا سبب فشل Top-k وTop-p تحت ظروف توزيع معينة (موحدة مقابل منحرفة)، وأثبتوا أن الانتباه المتفرق القابل للتدريب يمكن أن يحقق تفرعاً أعلى من خلال تركيز توزيع الانتباه أثناء التدريب.
2. استراتيجية التغطية الهجينة: تقديم قاعدة تغطية اتحادية بسيطة وفعالة (Top-k+Top-p) تعمل على استقرار الانتباه المتفرق عند مستويات تفرع عالية جداً.
3. التكيف القائم على التقطير: هدف ضبط دقيق مبتكر باستخدام تقطير السرعة لمواءمة نماذج الانتباه المتفرق مع معلمي الانتباه الكامل، مما يحل مشكلة عدم تطابق توزيع البيانات الشائعة في تكييف نماذج الفيديو مفتوحة المصدر.
4. الأداء الرائد (State-of-the-Art): يحقق SpargeAttention2 تفرعاً غير مسبوق مع الحفاظ على جودة نماذج الانتباه الكامل أو تجاوزها.

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

تم تقييم الطريقة على نماذج انتشار الفيديو Wan2.1 (بمعلمات 1.3 مليار و14 مليار) عبر دقات 480p و720p.

• التفرع والسرعة:
  • حقق 95% تفرع في الانتباه.
  • تسريع بمقدار 16.2 ضعفاً في وقت تشغيل الانتباه مقارنة بالانتباه الكامل.
  • تسريع بمقدار 4.7 ضعفاً في توليد الفيديو النهائي (على سبيل المثال، تقليل وقت توليد 14B 720p من ~3043 ثانية إلى ~650 ثانية).
• جودة التوليد:
  • طابق أو تجاوز SpargeAttention2 خط الأساس للانتباه الكامل في جميع المقاييس (جودة التصوير، الجودة الجمالية، مكافأة الرؤية، ودقة VQA).
  • في المقابل، أظهرت الطرق السابقة (VSA, VMoBA, SLA) تدهوراً ملحوظاً في الجودة عند مستويات التفرع المماثلة.
• دراسات الاستئصال (Ablation Studies):
  • أدى إزالة التغطية الهجينة (باستخدام Top-k أو Top-p فقط) إلى انخفاض الجودة.
  • أدت إزالة التدريب (باستخدام انتباه متفرق مجمد) إلى تدهور حاد في الجودة.
  • استبدال تقطير السرعة بخسارة الانتشار القياسية أدى إلى انخفاض الأداء، مما يؤكد ضرورة هدف التقطير.

5. الأهمية

يمثل SpargeAttention2 قفزة نوعية في توليد الفيديو عالي الكفاءة. فمن خلال معالجة كل من القيود الخوارزمية لاستراتيجيات التغطية وتحديات التحسين للضبط الدقيق، فإنه يتيح نشر نماذج انتشار الفيديو عالية الدقة على أجهزة أكثر سهولة في الوصول إليها. إن القدرة على تحقيق تفرع بنسبة 95% دون فقدان الجودة تجعل توليد الفيديو عالي الدقة في الوقت الفعلي أو القريب من الوقت الفعلي أمراً ممكناً، مما يسد الفجوة بين الكفاءة النظرية والتطبيق العملي.