ملخص تقني: kNN-MoE

بيان المشكلة

تعمل بنيات خليط الخبراء (MoE) على توسيع نطاق النماذج اللغوية الكبيرة (LLMs) بكفاءة من خلال استخدام "موجه" (router) معلمي (parametric) لتوزيع الرموز (tokens) على مجموعة فرعية متفرقة من الخبراء. ومع ذلك، فإن الموجهات القياسية هي مصنفات خفيفة الوزن يتم تدريبها مرة واحدة ثم تجميدها. ورغم فعاليتها في البيانات الموزعة داخل النطاق (in-distribution)، إلا أنها تفتقر إلى المرونة أثناء الاستدلال عند مواجهة تحولات في التوزيع (مثل المجالات المتخصصة أو المهام خارج النطاق). في مثل هذه السيناريوهات، يجب على الموجه الاستقراء من أولويات ثابتة، مما يؤدي غالباً إلى اختيار غير مثالي للخبراء وارتفاع في درجة الحيرة (perplexity). تحاول طرق التكيف عند وقت الاختبار الحالية، مثل C3PO، ضبط لوغاريتمات التوجيه (routing logits) عبر الإنترنت، لكنها تتسبب في عبء حسابي كبير (حوالي 5 أضعاف النموذج الأساسي). أما الأساليب الأخرى فتتطلب تحديثات دائمة للمعلمات عبر الضبط الدقيق الخاضع للإشراف (SFT)، وهو أمر مكلف حسابياً وأقل كفاءة في عملية الاستدلال.

المنهجية: kNN-MoE

يقترح المؤلفون kNN-MoE، وهو إطار عمل للتوجيه المعزز بالاسترجاع يعمل على تحسين قرارات تعيين الخبراء دون تعديل معلمات النموذج. تعمل الطريقة في مرحلتين متمايزتين:

1. بناء الذاكرة (خارج الخط - Offline)

يبني النظام مخزناً غير معلمي للذاكرة $\mathcal{M}$ لكل طبقة MoE، يتكون من أزواج (مفتاح-قيمة) $(k_t, v_t)$:

• المفتاح ($k_t$): مدخل الموجه (الحالة المخفية) $x_t$ من مجموعة بيانات مرجعية $D_{ref}$.
• القيمة ($v_t$): تعيين خبير "مثالي محلياً" مستمد من تحسين لوغاريتمات التوجيه لزيادة احتمالية الرمز الحقيقي $y_t$.
  • بدلاً من تخزين مخرجات الموجه المجمد، تقوم الطريقة بحل مشكلة تحسين لكل رمز في $D_{ref}$: إيجاد متجه لوغاريتمي $r^*_t$ يقلل من السلب اللوغاريتمي للاحتمالية للرمز المستهدف، بالنظر إلى معلمات الشبكة المجمدة.
  • يتم إجراء هذا التحسين عبر خطوات قليلة من الانحدار الاشتقاقي (غالباً $S=1$) بدءاً من اللوغاريتمات المعلمية الأصلية.
  • تعمل اللوغاريتمات المحسنة الناتلة $r^{(S)}_t$ كقيمة $v_t$ في الذاكرة.

2. الخلط التكيفي القائم على الثقة (الاستدلال)

أثناء الاستدلال، ولأي مدخل موجه $x$:

• الاسترجاع: يقوم النظام باسترجاع أقرب $K$ من الجيران $\mathcal{N}(x)$ من الذاكرة $\mathcal{M}$ بناءً على التشابه بين $x$ والمفاتيح المخزنة.
• الاقتراح غير المعلمي: يتم حساب اقتراح لوغاريتمي قائم على الذاكرة $r_{mem}(x)$ عن طريق تجميع اللوغاريتمات المثلى المسترجعة، موزونة بتشابهها مع المدخل.
• الخلط المدفوع بالثقة: اللوغاريتمات النهائية للتوجيه هي استكمال خطي بين لوغاريتمات الموجه المعلمي الأصلي $r(x)$ والاقتراح القائم على الذاكرة $r_{mem}(x)$:
  $$r_{final}(x) = (1 - \lambda(x))r(x) + \lambda(x)r_{mem}(x)$$
  يُعرف معامل الخلط $\lambda(x)$ بأنه متوسط التشابه للجيران المسترجعين. ويعمل هذا كدرجة ثقة:
  • تشابه عالٍ ($\lambda \approx 1$): يثق النموذج في التعيين المسترجع والمثالي محلياً.
  • تشابه منخفض ($\lambda \approx 0$): يعود النموذج إلى الموجه المجمد الأصلي، لتجنب الضجيج من الحالات غير ذات الصلة.
• المخرج: تمر اللوغاريتمات النهائية عبر "Softmax" من نوع Top-K لتوليد تعيين الخبراء المتفرق، مما يحافظ على قيود التفرق الأصلية للنموذج.

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

1. التوجيه المعزز بالاسترجاع: يقدم إطار عمل يعيد استخدام تعيينات الخبراء المثلى محلياً والمحسوبة مسبقاً من ذاكرة حالات سابقة مشابهة لتحسين قرارات التوجيه.
2. التكيف غير المعلمي: على عكس طرق SFT أو التكيف عبر الإنترنت، لا يقوم kNN-MoE بتحديث معلمات النموذج. بل ينقل العبء الحسابي إلى مرحلة بناء الذاكرة خارج الخط، مما يتيح استدلالاً فعالاً.
3. الرجوع المدفوع بالثقة: تستخدم الطريقة متوسط تشابه الجيران المسترجعين كمعامل خلط ديناميكي، مما يضمن المتانة من خلال العودة إلى الموجه المجمد عند عدم وجود حالات تاريخية ذات صلوة.
4. الكفاءة: تتطلب الطة خطوة واحدة فقط من الانحدار الاشتقاقي لكل رمز أثناء بناء الذاكرة وتتجنب التكاليف الحسابية الثقيلة للتحسين عبر الإنترنت أو الضبط الدقيق الكامل للنموذج.

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

قيم المؤلفون kNN-MoE على ثلاثة نماذج MoE (OLMoE، وGPT-OSS، وQwen3) عبر ستة اختبارات معيارية (GPQA، وSuperGPQA، وMMLU، وUSMLE، وMedMCQA، وMBPP).

• الأداء مقابل النماذج المرجعية: تفوق kNN-MoE باستمرار على نموذج Zero-shot baseline ونموذج "SFT (Router Only)" عبر جميع النماذج ومجموعات البيانات.
• التنافسية مع SFT: حققت الطريقة أداءً منافساً للضبط الدقيق الخاضع للإشراف الكامل (SFT)، وتفوقت في بعض الإعدادات شحيحة البيانات (مثل GPQA وUSMLE) على SFT، الذي عانى أحياناً من الإفراط في التخصيص (overfitting) على المجموعات المرجعية الصغيرة.
• الاستقرار: على عكس التلقين بـ 5-shots، الذي أدى إلى تدهور الأداء في بعض النماذج (GPT-OSS، OLMoE) بسبب مشكلات طول المدخلات، أظهر kNN-MoE تحسينات مستقرة.
• دراسات الاستئصال (Ablation Studies):
  • الجيران ($K$): خلافاً لما توضحه نتائج kNN-LM المعتادة، بلغ الأداء ذروته عند $K=1$. حيث أدى تجميع المزيد من الجيران ($K>1$) إلى إدخال الضجيج، مما يشير إلى أن إشارات التوجيه المفيدة حقاً نادرة ومحددة.
  • خطوات التحسين ($S$): لم تؤدِ زيادة خطوات الانحدار الاشتقاقي بما يتجاوز $S=1$ إلى مكاسب تذكر في الدقة مع زيادة تكاليف البناء.
  • حجم المجموعة المرجعية: تحسن الأداء بشكل رتيب مع حجم المجموعة المرجعية $|D_{ref}|$، مع ملاحظة تناقص العوائد بعد 0.5k–1k مثال.
• الكفاءة: استغرق بناء الذاكرة 0.46 ساعة لـ 1k مثال (مقابل ~1.34 ساعة لـ SFT). قللت نسخة IVFPQ (Inverted File + Product Quantization) من زمن تأخير الاستدلال الإضافي إلى ~3.67% ونمو استخدام ذاكرة الوصول العشوائي (VRAM) إلى ~9.43% مقارنة بالنموذج الأساسي.

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

يدعي البحث أن kNN-MoE يسد الفجوة بين التوجيه المعلمي الثابت والاحتياجات الديناميكية للاستدلال. تكمن أهميته الأساسية في تقديم مسار قابل للتوسع لتعزيز قدرة التكيف في نماذج MoE دون العبء الحسابي للضبط الدقيق.

يؤكد المؤلفون أن الطريقة فعالة بشكل خاص عندما يكون الموجه المعلمي غير متيقن (مرتبط بارتفاع درجة الحيرة)، حيث توفر آلية الاسترجاع إشارة تصحيحية من حالات تاريخية عالية الجودة ومشابهة. وبينما تعتمد الطريقة على توفر مجموعة مرجعية مصنفة ذات توزيع مشابه لبيانات الاختبار، فإنها توفر بديلاً عملياً وغير معلمي لتقنيات التكيف المكلفة، محققة نتائج منافسة لـ SFT مع وقت تحضير أقل بكثير ودون أي تحديث للمعلمات.