Reward-Based Online LLM Routing via NeuralUCB

إليك ملخص تقني مفصل لورقة البحث بعنوان "Reward-Based Online LLM Routing via NeuralUCB" من إعداد مينج-هوا تساي وفات تران.

1. تعريف المشكلة

تعالج الورقة تحدي توجيه النماذج اللغوية الكبيرة (LLM) الواعي بالتكلفة. في عمليات النشر الواقعية، تتباين النماذج اللغوية الكبيرة بشكل كبير في كل من القدرة (الجودة) وتكلفة الاستدلال. توفر النماذج عالية المستوى أداءً فائقاً ولكنها مكلفة، بينما قد تفشل النماذج الأرخص في المهام المعقدة.

الهدف هو تصميم سياسة توجيه ديناميكية تختار النموذج الأكثر ملاءمة ($a_t$) لاستعلام مستخدم معين ($x_t$) لتعظيم المنفعة (utility) على المدى الطويل. يتم صياغة هذا كمسألة Bandit سياقي (Contextual Bandit) تتميز بالخصائص التالية:

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

2. المنهجية

يتكون الحل المقترح، التوجيه القائم على NeuralUCB، من ثلاثة مكونات أساسية: دالة المنفعة، والمتنبئ العصبي (UtilityNet)، وسياسة الاستكشاف والاستغلال.

أ. دالة المنفعة (Utility Function)

للموازنة بين الجودة والتكلفة، عرّف المؤلفون مكافأة المنفعة $r(x, a)$ كالتالي:
$$r(x, a) = q(x, a) \cdot \exp(-\lambda \tilde{c}(x, a))$$
حيث أن:

• $q(x, a)$ هي جودة النموذج.
• $\tilde{c}(x, a)$ هي التكلفة المنظمة لوغاريتمياً، والتي تقوم بخرائط التكاليف الخام إلى نطاق $[0, 1]$ لاستقرار التدريب.
• $\lambda$ هو معلم فائق (hyperparameter) يتحكم في قوة العقوبة المفروضة على التكلفة.

ب. بنية UtilityNet

UtilityNet هو شبكة عصبية مصممة للتنبؤ بمكافأة المنفعة وتوفير تمثيلات للميزات لتقدير عدم اليقين. وهي تعالج:

1. المدخلات: تضمينات النصوص ($x_{emb}$)، الميزات المساعدة ($x_{feat}$)، ومعرفات النطاق ($d$).
2. المشفرات (Encoders): طبقات MLP منفصلة للنصوص، والميزات، وتضمينات النطاق.
3. الدمج (Fusion): يتم دمج تضمينات النص والميزة والعملة لتشكيل تمثيل مشترك.
4. الرؤوس المزدوجة (Dual Heads):
   • رأس المنفعة (Utility Head): يقوم بانحدار المنفعة المتوقعة $\mu(x, a)$ باستخدام خسارة هوب (Huber loss).
   • رأس البوابة (Gating Head): يتنبأ باحتمالية $p(x)$ تشير إلى ما إذا كان الاستكشاف (إضافة مكافأة UCB) مفيداً للسياق الحالي. وهذا يسمح للنظام بالتبديل بين الاستكشاف الهجومي والاستغلال المتحفظ.

ج. سياسة NeuralUCB

تستخدم سياسة التوجيه خوارزمية Neural Upper Confidence Bound (NeuralUCB) لاختيار الإجراء:

• تمثيل الميزات: يعمل الطبقة المخفية الأخيرة من UtilityNet، وهي $h(x, a)$، كمتجه ميزات لحساب UCB.
• درجة UCB: هي $s(x, a) = \mu(x, a) + \beta \sqrt{g(x, a)^\top A^{-1} g(x, a)}$، حيث $A^{-1}$ هي مصفوفة تغاير عكسية مشتركة تقدر عدم اليقين في مساحة الميزات، و $\beta$ تتحكم في مكافأة الاستكشاف.
• قاعدة الاختيار المعتمدة على البوابة (Gated Selection Rule):
  • إذا كان احتمال البوابة $p(x) \geq \tau_g$ (العتبة)، يختار النظام الإجراء الذي يحقق أعلى درجة UCB (تمكين الاستكشاف).
  • خلاف ذلك، يعود النظام إلى إجراء آمن ($a_{safe}$) يعظم متوسط المنفعة المتوقعة $\mu(x, a)$ (الاستغلال).

د. بروتوكول التدريب

تم تقييم النموذج في بيئة محاكاة عبر الإنترنت باستخدام مجموعة بيانات RouterBench:

• تم تقسيم البيانات إلى 20 شريحة متتالية.
• لكل شريحة، تختار السياسة الإجراءات، وتحدث ذاكرة التخزين المؤقت لإعادة التشغيل (replay buffer) ومصفوفة التغاير المشتركة $A^{-1}$، ثم تعيد تدريب UtilityNet.
• هذا يحاكي بيئة التعلم عبر الإنترنت دون التكلفة الباهظة لطلبات API الحية.

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

1. استخدام NeuralUCB لتوجيه LLM: أول تطبيق لـ NeuralUCB على مشكلة توجيه LLM الواعي بالتكلفة، مما يعالج بفعالية الخرائط غير الخطية بين الاستعلام والنموذج تحت التغذية الراجعة الجزئية.
2. آلية الاستكشاف المعتمدة على البوابة: تقديم فرع بوابة يقرر ديناميكياً ما إذا كان ينبغي تطبيق مكافأة استكشاف UCB، مما يمنع الاستكشاف غير الضروري في الاستعلامات التي يكون فيها النموذج الأمثل واضحاً بالفعل.
3. إطار تقييم شامل: تنفيذ خط معال إعادة التشغيل (offline replay pipeline) قوي على RouterBench يحاكي ديناميكيات التعلم عبر الإنترنت، مما يسمح بتتبع المكافآت التراكمية والمقايضات بين التكلفة والجودة بمرور الوقت.

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

تم تقييم الطريقة على RouterBench (36,497 عينة، 86 مجالاً، 11 نموذجاً لغوياً مرشحاً) مقابل نماذج مرجعية تشمل Random و Min-Cost و RouteLLM-BERT.

• الأداء مقابل النماذج المرجعية:
  • المكافأة: حققت NeuralUCB باستمرار أعلى متوسط مكافأة (حوالي 0.59–0.61)، متفوقة بشكل كبير على RouteLLM-BERT (0.35)، و Min-Cost (0.52)، و Random (~0.32).
  • المكافأة التراكمية: اتسعت الفجوة بين NeuralUCB والنماذج المرجعية بمرور الوقت، مما يدل على التعلم والتكيف الفعالين.
• المقايضة بين التكلفة والجودة:
  • حققت سياسة NeuralUCB جودة تنافسية مع استخدام ~33% فقط من تكلفة الاستدلال الخاصة بالمرجع "max-quality" (الذي يختار دائماً أفضل نموذج بغض النظر عن التكلفة).
• دراسات الاستئصال (Ablation Studies):
  • مشفرات النصوص: أظهرت المشفرات الأصغر مثل all-MiniLM-L6-v2 و all-mpnet-base-v2 أداءً يضاهي أو يتفوق على المشفرات الأكبر والأكثر تعقيداً (مثل multilingual-e5-large-instruct)، مما يشير إلى تناقص العوائد من حجم المشفر في سياق التوجيه هذا.

5. الأهمية والعمل المستقبلي

• الأثر العملي: تثبت الدراسة أنه من الممكن توجيه استعلامات LLM ديناميكياً لتحقيق توازن متفوق بين التكلفة والأداء، مما قد يوفر موارد حوسبية كبيرة في أنظمة الإنتاج.
• القيود: يلاحظ المؤلفون أنه بالنسبة للعديد من العينات، يكون نموذج "max-quality" هو الخيار الأمثل للمكافأة بالفعل، مما يحد من المكاسب المحتملة من التوجيه.
• الاتجاهات المستقبلية: تقترح الورقة التركيز على جمع بيانات تدريب أكثر تنوعاً وتحسين نموذج المنفعة لتحسين التمييز بين الإجراءات (action discrimination)، خاصة في السيناريوهات التي تكون فيها قيود التكلفة صارمة والنموذج الأمثل ليس واضحاً.

في الختام، تقدم هذه الورقة نهجاً قوياً ومبنياً على أسس نظرية لتوجيه LLM، يستفيد من الـ Bandits العصبية السياقية للتنقل في المشهد المعقد لقدرات النماذج وتكاليفها، مما يوفر مساراً عملياً لنشر أكثر كفاءة للنماذج اللغوية الكبيرة.