ملخص تقني: تقطير القناع الظلي (Shadow Mask Distillation) لتعزيز كفاءة الذاكرة في مرحلة ما بعد التدريب لتعلم التعزيز (RL)

1. بيان المشكلة: عدم التطابق الهيكلي في السياسة خارج النطاق (Off-Policy Mismatch)

تعد أطر تعلم التعزيز (RL) مثل RLHF وRLAIF وPPO وGRPO ضرورية لإطلاق قدرات الاستنتاج في النماذج اللغوية الكبيرة (LLMs). ومع ذلك، يواجه تطبيق هذه الأساليب على المهام ذات السياق الطويل "جدار ذاكرة" شديد بسبب البصمة الضخمة لذاكرة التخزين المؤقت للمفاتيح والقيم (KV cache) المطلوبة خلال مرحلة الاستكشاف (rollout phase).

بينما يمكن لخوارزميات ضغط ذاكرة KV (مثل SnapKV وH2O) تخفيف اختناقات الذاركة من خلال جعل السياق متفرقاً (sparse) أثناء عملية الاستكشاف، إلا أنها تفرض ثنائية هيكلية حرجة في مسار تعلم التعزيز:

• مرحلة الاستكشاف (Rollout Phase): يقوم الممثل (actor) بتوليد المسارات باستخدام سياق متفرق ($\pi_{sparse}$) نتيجة لعملية طرد الـ KV.
• مرحلة المتعلم (Learner Phase): تقوم أطر العمل القياسية بتقييم هذه المسارات باستخدام السياق الكامل والكثيف ($\pi_{dense}$) لحساب التدرجات (gradients).

هذا التباين ينتهك الافتراض الجوهري لطرق تدرج السياسة (policy gradient)، حيث يعاقب المتعلم النموذج على فشله في استخدام معلومات لم يراها أصلاً أثناء عملية التوليد. تحاول الحلول الإحصائية الموجودة، مثل إعادة وزن الأهمية (Importance Reweighing) وأخذ العينات بالرفض الواعي بالتفرق (Sparsity-Aware Rejection Sampling) (مثل Sparse-RL)، معالجة هذا التباين ولكنها تعاني من:

1. تباين تدرج أسي (Exponential Gradient Variance): تنمو نسب الأهمية بشكل أسي مع طول التسلسل، مما يؤدي إلى زعزعة استقرار التدريب.
2. عدم كفاءة العينات: تؤدي عملية أخذ العينات بالرفض إلى التخلص من مسارات استكشاف مكلفة (غالباً >20%)، مما يهدر الموارد الحسابية.
3. انهيار الأداء: يؤدي الدمج الساذج لضغط KV إلى انخفاض حاد في الدقة (على سبيل المثال، انخفاض دقة GSM8K من 74.5% إلى 64.3%).

2. المنهجية: تقطير القناع الظلي (Shadow Mask Distillation - SMD)

يقترح المؤلفون تقطير القناع الظلي (SMD)، وهو إطار معماري مصمم للقضاء على انحياز السياسة خارج النطاق (off-policy bias) فيزيائياً بدلاً من إحصائياً. يعمل SMD عبر آلية ثنائية المسار:

المسار الأول: المحاذاة المعمارية داخل النطاق عبر "القناع الظلي"

بدلاً من التصحيح الإحصائي، يضمن SMD أن يختبر المتعلم نفس الاختناق المعلوماتي الذي واجهه المولد تماماً.

• تسجيل القناع الظلي (Shadow Mask Recording): أثناء الاستكشاف المتفرق، يقوم النظام بتسجيل "قناع ظلي" ثنائي ($M$) يشير إلى الرموز (tokens) التي تم الاحتفاظ بها بواسطة خوارزمية طرد الـ KV.
• حقن القناع (Mask Injection): خلال مرحلة المتعلم، يتم حقن هذا القناع في مصفوفة الانتباه السببي (causal attention matrix) للنموذج. يتم تعديل درجات الانتباه بحيث تتلقى المفاتيح المقابلة للرموز المطرودة عقوبة مسبقة لـ softmax بقيمة $-\infty$، مما يجبر أوزان الانتباه الخاصة بها على الصفر.
• النتيجة: يتم تقييد المرور الأمامي (forward pass) للمتعلم بدقة ضمن المجال المعلوماتي الذي تم اختباره أثناء التوليد، مما يضمن رياضياً أن $\pi_{sparse} \equiv \pi_{shadow}$. يحقق هذا محاذاة صارمة داخل النطاق (on-policy alignment) مع صفر من التباين خارج النطاق، بغض النظر عن طول السياق.

المسار الثاني: التقطير من الكثيف إلى المتفرق

لمنع النموذج من الإفراط في التكيف (overfitting) مع السياق المقتطع وفقدان قدرات الاستنتاج العالمي، يقدم SMD مساراً كثيفاً ثانوياً.

• مرور أمامي مزدوج (Dual Forward Pass): بالتزامن مع المرور المقنع، يتم إجراء مرور أمامي كثيف كامل باستخدام السياق الكامل.
• عقوبة تباعد KL (KL Divergence Penalty): تعمل السياسة الكثيفة كمعلم ضمني. يتم تطبيق عقوبة تباعد كولباك-لايبلر (KL) لتقطير المعرفة السياقية العالمية في السياسة المقنعة.
• الأوزان المشتركة: يتشارك كلا المسارين نفس المعلمات الأساسية ($\theta$). يتم فصل (detach) المعلم الكثيف عبر عامل إيقاف التدرج (stop-gradient). هذا الإعداد يجبر الأوزان المشتركة على تعلم تمثيلات قوية يمكنها تقريب حدود القرار الكثيف حتى عند تقييدها بسياق متفرق، مما يعمل كمنظم ضمني (implicit regularizer).

الهدف العام

يجمع هدف التدريب النهائي بين خسارة GRPO المحاذات الصارمة وعقوبة التقطير:
$$\max_{\theta} L_{SMD}(\theta) = \mathbb{E} \left[ L_{GRPO}(\pi_{\theta}(y|x, M)) - \lambda D_{KL}(\text{sg}[\pi_{\theta}(y|x)] \parallel \pi_{\theta}(y|x, M)) \right]$$

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

1. حل معماري لانحياز السياسة خارج النطاق: يحدد SMD التباين الهيكلي في RLHF الموفر للذاكرة ويقترح إطار عمل ثنائي المسار يضمن محاذاة التدرج المثالية دون هدر البيانات أو استخدام رقع إحصائية عالية التباين.
2. القضاء على طفرات VRAM: يكشف البحث أن تقسيم التنسور (tensor-slicing) الأصيل لعملية طرد KV في أطر العمل عالية المستوى يسبب "طفرات تخصيص غير مكانية" (Not-In-Place Allocation Spikes) (تجزئة لحظية للذاكرة). يتجاوز SMD ذلك باستخدام محاكاة القناع أثناء مرحلة المتعلم، مما يضمن تنفيذاً قوياً دون أخطاء الخروج من الذاكرة (OOM) الكارثية.
3. التنظيم الضمني: يوضح المؤلفون أن عمليات الاستكشاف ذات الانتباه المتفرق، عند دمجها مع SMD، تعمل كمنظم فعال. في اختبار Reddit TL;DR، يتفوق SMD على خط الأساس الكثيف في مقياس ROUGE-L، مما يشير إلى أن الاختناق المعلوماتي يمنع الإفراط في التكيف مع مكافآت RL قصيرة النظر.

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

أجريت التجارب على نموذج 4B (Qwen3-4B-Instruct-2507) عبر أربع مجموعات بيانات: TL;DR، وGovReport، وGSM8K، وHotpotQA.

• الأداء مقابل خطوط الأساس: تحت ضغط ضغط ذاكرة KV بنسبة 50%، يحقق SMD أداءً يقارب الأداء غير المنضغط للأساس الكثيف (Dense Baseline).
  • GSM8K: يحافظ SMD (73.6%) على تنافسية عالية مع الأساس الكثيف (74.5%) ويتفوق بشكل كبير على Sparse-RL (71.2%).
  • مهام السياق الطويل: في GovReport وHotpotQA، يطابق SMD أو يتجاوز قليلاً الأساس الكثيف، بينما يعاني Sparse-RL من تدهور كبير بسبب أخذ العينات بالرفض.
• كفاءة النظام: يقضي SMD تماماً على طفرات الذاكرة اللحظية المتأصلة في طرد الـ KV الفيزيائي أثناء مرحلة تحسين المتعلم. بينما يسبب Sparse-RL طفرة في VRAM قدرها +1.24 جيجابايت (150% من خط الأساس)، يحافظ SMD على نسبة مستقرة بنسبة 100% من VRAM.
• دراسات الاستئصال (Ablation Studies):
  • نسبة الضغط: يحافظ SMD على أداء قوي حتى عند ضغط بنسبة 70%، بينما ينهار Sparse-RL بسبب انفجار التباين.
  • معامل التقطير: ينتج $\lambda = 0.1$ تقارباً مثالياً؛ بينما يؤدي التقطير المفرط ($\lambda = 0.8$) إلى تدهور الأداء.
  • اختيار الرموز (Token Selection): يتميز SMD بالمتانة عبر استراتيجيات الاختيار (SnapKV، Recent، Random)، بينما تفشل الإعدادات الساذجة مع الاختيار العشوائي.

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

يزعم البحث أن تقطير القناع الظلي يحل جوهرياً معضلة "خارج النطاق" في RL طويل السياق من خلال الانتقال من الرقع الإحصائية إلى المحاذاة المعمارية.

• الضمان النظري: من خلال فرض تطابق مطلق بين المولد المتفرق والمتعلم المقنع، يحقق SMD نظرياً صفراً من التباين الإضافي خارج النطاق، وهو إنجاز يستحيل على الطرق الإحصائية تحقيقه والتي تعاني من توسع التباين الأسي.
• الأثر العملي: يُمكّن SMD من إجراء RLHF طويل السياق على الأجهزة القياسية عبر إزالة جدار الذاكرة ومنع عدم استقرار التدريب. إنه يضع معياراً قوياً وفعالاً للذاكرة يتجنب هدر البيانات في أخذ العينات بالرفض وتباين التدرج في إعادة وزن الأهمية.
• القيود: يقر المؤلفون بأن SMD يضيف عبئاً حسابياً طفيفاً (مرور أمامي مزدوج) ويعتمد على الأقنعة الثنائية، مما يجعله غير متوافق مع تقنيات الضغط المستمر مثل التكميم (quantization). ومع ذلك، يجادلون بأن هذا العبء هامشي (أقل من 10% من إجمالي وقت التدريب الزائد) بالنظر إلى هيمنة مرحلة الاستكشاف في RLHF.