ملخص تقني: توازن المسار-f (f-Trajectory Balance)

بيان المشكلة

يعتمد محاذاة وضبط النماذج التوليدية، بما في ذلك النماذج اللغوية الكبيرة (LLMs) وشبكات التدفق التوليدية (GFlowNets)، عادةً على تقليل التباين بين توزيع النموذج وتوزيع مستهدف. وفي الضبط الدقيق باستخدام التعلم التعزيزي (RL)، غالبًا ما يُصاغ هذا كتعظيم للمكافأة المتوقعة مع تطبيق عقوبة تباعد كولباك - ليبلر (KL divergence).

يتمثل أحد التحديات الكبيرة في هذا المجال في المقايضة بين سلوكيات البحث عن النمط (mode-seeking) وتغطية الأنماط (mode-covering). فالأهداف القياسية، مثل تباعد KL العكسي (الذي يتم تنفيذه غالبًا عبر توازن المسار أو خسارة KL المربعة)، تميل إلى البحث عن النمط، مما يتسبب في انهيار النماذج على عدد قليل من الأنماط عالية الاحتمالية. وهذا أمر غير مثالي للمهام التي تتطلب الاستكشاف، مثل اكتشاف الأدوية أو توليد المرشحين المتنوعين، حيث يُفضل سلوك "تغطية الأنماط" (أي التقاط التنوع الكامل للتوزيع المستهدف).

علاوة على ذلك، تعاني الطرق الموجودة غالبًا من مشكلة صلاحية السياسة خارج النطاق (off-policy validity). فبينما تُستخدم مقدرات مثل REINFORCE أو PPO بشكل شائع، إلا أنها تعاني من تباين عالٍ. وتقدم النهج الحديثة باستخدام خسارة KL المربعة ($KL_{sq}$) تدرجات منخفضة التباين عند العمل ضمن السياسة (on-policy) وصلاحية عند العمل خارج النطاق (off-policy)، لكنها مرتبطة جوهريًا بخصائص تباعد KL العكسي، مما يحد من القدرة على التحكم في المقايضة بين الاستكشاف والاستغلال.

المنهجية

يقترح المؤلفون توازن المسار-f، وهو عائلة من دالات خسارة البديل التي تعمم خسارة KL المربعة على عائلة كاملة من تباعدات f (f-divergences).

الأساس النظري

تتمثل المساهمة النظرية الجوهرية في إثبات وجود ارتباط واحد لواحد بين دالات الخسارة الثابتة بالترجمة (translation-invariant loss functions) على لوغاريتم الاحتمالات وتباعدات f.

1. من التباعد إلى الخسارة: لأي دالة محدبة $f$ تُعرف تباعد f، يستنتج المؤلفون دالة خسارة $L_f$ معرفة على الفرق بين لوغاريتم الاحتمالات $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$:
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   تتمتع هذه الخسارة بصفتين حرجتين:

   • تكافؤ التدرج ضمن السياسة (On-Policy Gradient Equivalence): عندما تُسحب العينات من سياسة النموذج ($\mu = p_\theta$)، فإن التدرجات المتوقعة ذاتية الاشتقاق لـ $L_f$ تطابق تدرجات تباعد f المقابل.
   • الصلاحية خارج النطاق (Off-Policy Validity): عند التقييم على بيانات خارج النطاق (مسحوبة من توزيع سلوكي $\mu$)، تظل الخسارة هدفًا صالحًا مع نفس المقلل العالمي ($p_\theta = p^*$)، بشرط أن يكون دعم الهدف متضمنًا في دعم سياسة السلوك.

2. من الخسارة إلى التباعد: على العكس من ذلك، فإن أي خسارة $\ell$ شديدة التحدب وثابتة بالترجمة على لوغاريتم الاحتمالات تقابل تقليل تباعد f محدد. وقد أثبتت الورقة أن هذه الخرائط هي عكس بعضها البعض.

التنفيذات العملية

لمعالجة التحديات العملية في التدريب، يقدم المؤلفون عدة توسعات:

• خسارة DevGrad (التطبيع على مستوى الدفعة): للتعامل مع دوال التجزئة غير القابلة للحساب (الشائعة في GFlowNets والأهداف غير المعيرة)، يعمم المؤلفون مقدر VarGrad. فبدلاً من استخدام تباين الدفعة، يستخدمون الانحراف المعمم على مستوى الدفعة لفرق لوغاريتم الاحتمالات. يتضمن ذلك تقدير ثابت التجزئة $\log Z$ عن طريق تقليل الخسارة عبر الدفعة، مما يؤدي فعليًا إلى تمركز معاملات دالة الدرجة وتقليل التباين.
• الخسارة المبردة (Tempered Loss): لمنع تجاوز السعة الرقمية عند التعامل مع معاملات درجة حرارة $\beta$ صغيرة (حيث التوزيع المستهدف هو $p^* \propto \pi_{ref} \exp(r/\beta)$)، يقترح المؤلفون خسارة مبردة. تقوم هذه الخسارة بتوسيع الخسارة بمقدار $1/\beta$ وتطبيقها على التوزيع المبرد، مما يضمن بقاء التدرجات محدودة ومستقرة.
• توازن المسار-f لشبكات GFlowNets: يتم تطبيق الإطار على GFlowNets من خلال فرض قيود توازن المسار باستخدام خسارة $L_f$. وهذا يعمم توازن المسار القياسي (الذي يقابل KL العكسي) على تباعدات f التعسفية.

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

1. عائلة الخسارة المعممة: اشتقاق عائلة عامة من خسائر البديل الثابتة بالترجمة ($L_f$) التي تطابق تدرجاتها ضمن السياسة تدرجات تباعد f المقابلة، مع الاحتفاظ بالصلاحية خارج النطاق.
2. التكافؤ النظري: إثبات العلاقة العكسية بين خسائر لوغاريتم الاحتمالات الثابتة بالترجمة وتباعدات f، مما يسمح بالتصميم المنهجي لدالات خسارة جديدة ذات خصائص مرغوبة (مثل تغطية الأنماط).
3. التوسع إلى الأهداف غير المعيرة: تعميم تقنية VarGrad إلى تباعدات f التعسفية عبر خسارة DevGrad، مما يعالج مشكلة دالة التجزئة غير القابلة للحساب في الضبط الدقيق لـ RL.
4. التحكم في الاستكشاف: إثبات أن اختيار تباعد f (المحدد بواسطة $\alpha$ في $\alpha$-divergences) يتحكم مباشرة في سلوك البحث عن النمط مقابل تغطية الأنماط. فقيم $\alpha$ الأقل (مثل KL الأمامي، وHellinger) تعزز تغطية الأنماط، بينما قيم $\alpha$ الأعلى (مثل Pearson $\chi^2$) تعزز البحث عن النمط.

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

يقيم المؤلفون عائلة الخسارة عبر أربع مهام متميزة:

1. الشبكة الاصطناعية (GFlowNets): في مهمة شبكة ثنائية الأبعاد مصممة لاختبار تغطية الأنماط، نجحت الخسارات ذات قيم $\alpha$ المنخفضة (Forward KL, Hellinger) في اكتشاف جميع الأنماط الأربعة للتوزيع المستهدف بشكل أسرع وأكثر موثوقية من توازن المسار القياسي (Reverse KL). وعلى العكس من ذلك، أظهرت Pearson $\chi^2$ ($\alpha=2$) سلوكًا قويًا في البحث عن النمط، حيث علقت في أول نمط وجدته.
2. اكتشاف الجزيئات (SynFlowNet): طُبقت على إطار عمل SynFlowNet لتوليد جزيئات يمكن الوصول إليها اصطناعيًا. وجد المؤلفون أن ضبط $\alpha$ سمح بالتحكم في تنوع الجزيئات المولدة. وقد حقق جدول التبريد (Annealing schedule) (بالبدء بـ $\alpha$ منخفض للاستكشاف وزيادته إلى $\alpha=1$ للاستغلال) أفضل مقايضة، حيث ولد جزيئات ذات مكافأة أعلى مع تنوع أكبر مقارنة بتوازن المسار القياسي.
3. نماذج الانتشار (Diffusion Models): في أخذ العينات الشرطي لأرقام MNIST (الفردية مقابل الزوجية)، فشل توازن المسار القياسي في أخذ عينات بالتساوي عبر الأنماط المستهدفة (مما أدى لزيادة فرط في أخذ عينات أرقام محددة مثل 0 و6). حققت خسائر توازن المسار-f، وخاصة تلك ذات $\alpha$ المنخفض، أخذ عينات أكثر انتظامًا عبر الأنماط المستهدفة.
4. ضبط LLM غير المتزامن: طُبقت الخسائر لضبط النماذج اللغوية الكبيرة (عائلات Qwen وOLMo) على المسائل الرياضية (GSM8k, MATH) مع تأخير غير متزامن قدره 50 خطوة. أظهرت خسائر تباعد f استقرارًا في التدريب ومقايضة ثابتة بين الإنتروبيا والمكافأة، بينما أظهرت تطبيقات PPO القياسية عدم استقرار تحت نفس الظروف غير المتزامنة.

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

يدعي البحث أن فعالية خسارة KL المربعة ليست ظاهرة معزولة بل هي حالة خاصة من عائلة أوسع من الخسائر. ومن خلال توسيع خاصية "مطابقة التدرج ضمن السياسة" و"الصلاحية خارج النطاق" لتشمل عائلة كاملة من تباعدات f، يوفر المؤلفون إطارًا موحدًا لضبط النماذج التوليدية.

تكمن الأهمية الأساسية في القدرة على التحكم الصريح في المقايضة بين الاستكشاف والاستغلال ببساطة عن طريق اختيار دالة خسارة بديلة مختلفة (أو ضبط معامل $\alpha$)، دون تغيير خوارزمية التدريب الأساسية. وهذا يسمح للممارسين بالاختيار بين سلوكيات البحث عن النمط (لتعظيم المكافأة) وتغطية الأنماط (للتنوع والاستكشاف) مع الحفاظ على الاستقرار والتوافق مع العمل خارج النطاق المطلوب للتدريب الحديث واسع النطاق وغير المتزامن لنماذج LLM وGFlowNets. ويشير العمل إلى أن هذا الإطار مفيد بشكل خاص لتطبيقات مثل اكتشاف الأدوية والوكلاء الاستكشافيين حيث يكون تغطية مساحة المرشحين ذوي المكافأة العالية أمرًا بالغ الأهمية.