ملخص تقني: موازنة المسار في مرحلة ما بعد التدريب لنماذج اللغة الانتشارية

المشكلة: قفل المسار في التعلم التعزيزي للنماذج الانتشارية

تقدم نماذج اللغة الانتشارية (dLLMs) بديلاً واعداً للنماذج ذات التوليد التتابعي (Autoregressive) في مجالات الاستنتاج وتوليد الكود البرمجي. ومع ذلك، فإن معظم طرق مرحلة ما بعد التدريب لنماذج dLLMs تتبنى أهداف التعلم التعزيزي (RL) الموجهة لتعظيم المكافأة (مثل PPO وGRPO) والمصممة للنماذج التتابعية. يحدد المؤلفون نمط فشل مركزي في هذا التكيف يسمى "قفل المسار" (Trajectory Locking).

في النماذج الانتشارية، يتم تعيين المكافأة فقط للإكمال النهائي المزال من الضجيج ($y$)، بينما يتم استحثاث توزيع النموذج عبر مسار إزالة الضجيج بأك entier ($\tau$). ولأن المكافآت النهائية هي "غير مبالية بالمسار" (Path-indifferent)—بمعنى أن أي مسار يصل إلى نفس الـ $y$ الصحيح يتلقى نفس المكافأة—فإن تحديثات تدرج السياسة التي تعتمد على العينات تميل إلى تعزيز مسارات إزالة الضجيج المحددة التي تم أخذ عينات منها مبكراً في التدريب. وهذا يخلق حلقة تغذية راجعة ذاتية التعزيز حيث تنهار الكتلة الاحتمالية على مجموعة ضيقة من مسارات إزالة الضجيج. وبناءً على ذلك، حتى لو تعلم النموذج إنتاج إجابة صحيحة، فإنه يفقد القدرة على توليد مجموعة متنوعة من الحلول الصالحة (الأوضاع/Modes) تحت عمليات أخذ العينات المتكررة. وقد صاغ المؤلفون ذلك عبر المبرهنة 1، حيث أثبتوا أن تغطية الأوضاع النهائية محدودة بحد أقصى بتنوع المسارات؛ وبالتالي، فإن قفل المسار يحد مباشرة من تغطية الحلول للنموذج.

المنهجية: TraFL (موازنة تدفق المسار)

لمعالجة مشكلة قفل المسار، يقترح المؤلفون TraFL، وهو هدف لما بعد التدريب يستند إلى مبدأ موازنة المسار من شبكات التدفق التوليدية (GFlowNets). بدلاً من تعظيم المكافأة مباشرة، يقوم TraFL بتدريب السياسة لمطابقة توزيع مستهدف مائل للمكافأة (Reward-tilted target distribution) مرتبط بنموذج مرجعي مجمد.

الهدف

يتم تعريف التوزيع المستهدف $p^*(\tau | x)$ كالتالي:
$$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$$
حيث أن:

• $p_{\text{ref}}$ هو نموذج مرجعي مجمد.
• $r(x, y)$ هي المكافأة النهائية.
• $\beta$ يتحكم في حدة المكافأة.
• $Z(x)$ هو حد التجزئة (Partition function) المعتمد على المطالبة (Prompt).

يهدف هدف التدريب إلى تقليل المتبقي التربيعي بين لوغاريتم احتمالية النموذج الحالي والهدف:
$$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$$
$$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$$

التنفيذ العملي للنماذج الانتشارية

بما أن نماذج dLLMs لا تكشف عن احتمالات المسار الدقيقة، فقد قدم المؤلفون مكونين رئيسيين لجعل TraFL عملياً:

1. بديل على مستوى التسلسل متوافق مع الانتشار: اشتقوا بديلاً قابلاً للحساب لـ $\log \hat{p}_\theta(y|x)$ يعتمد على حد أدنى لإعادة البناء المقنع (Masked-reconstruction lower bound). يتضمن ذلك أخذ عينة من مستوى القناع $l$، وتشويه $l$ من الرموز (Tokens)، ومتوسط لوغاريتمات الرموز المعاد بناؤها. ويُستخدم عامل $1/l$ لضمان استقلالية الدرجة (Score) عن عدد الرموز المقنعة.
2. تجزئة متعلمة تعتمد على المطالبة: على عكس الطرق التي تقدر $Z(x)$ عبر Softmax على دفعة واحدة (والتي قد تسبب التركيز/Concentration)، يتعلم TraFL دالة التجزئة $Z_\phi(x)$ عبر شبكة MLP صغيرة. تتلقى هذه القيمة إشارات التدرج في كل خطوة بغض النظر عن الإكمالات التي تم أخذ عينات منها، مما يمنع فشل التجزئة "المحلية للدفعة" الذي تسببه أساليب مطابقة التوزيع الأخرى.

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

1. تحديد قفل المسار: حددت الورقة رسمياً وأثبتت أن أهداف تعظيم المكافأة في نماذج dLLMs تؤدي إلى انهيار تنوع المسارات، مما يحد بدوره من تغطية أوضاع الحلول المتميزة.
2. هدف TraFL: اقتراح هدف موازنة المسار الذي يربط السياسة بنموذج مرجعي، مما يفصل دور المكافأة (إمالة الإكمالات النهائية) عن النموذج المرجعي (تخصيص الكتلة الاحتمالية عبر المسارات المعقولة).
3. البديل العملي والتجزئة: تطوير بديل على مستوى التسلسل متوافق مع الانتشار وحداً للتجزئة متعلمًا يتيح التدريب المستقر دون الحاجة إلى احتمالات المسار الدقيقة.
4. التفوق التجريبي: يعد TraFL الطريقة الوحيدة التي تم تقييمها والتي تحسن باستمرار عن النموذج الأساسي عبر جميع إعدادات طول الاختبار (GSM8K, MATH-500, HumanEval, MBPP) وميزانيات أخذ العينات.

النتائج

قيم المؤلفون TraFL على اختبارات الاستنتاج الرياضي وتوليد الكود البرمجي مقابل النموذج الأساسي (LLaDA-8B-Instruct) والأساليب المرجعية القوية (ESPO, JUSTGRPO).

• الأداء داخل التوزيع (In-Distribution): يحسن TraFL من قيمة Pass@k عبر جميع ميزانيات أخذ العينات ($k \in \{1, 2, 4, 8, 16\}$) ودرجات الحرارة. ومن الملاحظ أنه يحسن Pass@1 (دقة العينة الواحدة) بينما يتوسع بشكل أفضل مع زيادة ميزانيات أخذ العينات، على عكس الأساليب المرجعية التي غالباً ما تضحي بدقة العينة الواحدة مقابل مكاسب تعدد العينات.
• التعميم خارج التوزيع (Held-Out Generalization): تنتقل المكاسب إلى اختبارات خارج التوزيع. في Minerva Math، يظل TraFL أعلى من النموذج الأساسي. وفي LiveCodeBench، يعد TraFL الطريقة الأقوى عبر جميع تقسيمات الصعوبة (سهل، متوسط، صعب)، متفوقاً بشكل كبير على الأساليب المرجعية في المسائل المتوسطة والصعبة.
• تحليل التنوع: باستخدام بروتوكول "النموذج اللغوي كحكم" (LLM-as-a-judge)، أظهر المؤلفون أن مكاسب TraFL مرتبطة بتغطية أوسع لاستراتيجيات الحل المتميزة (مثل طرق إثبات مختلفة أو تفكيكات خوارزمية) بدلاً من مجرد تشديد الدرجة لوضع واحد. أما الأساليب المرجعية مثل ESPO، رغم تحسينها للدقة، فلا تظهر مكاسب كبيرة في التنوع مقارنة بالنموذج الأساسي، وهو ما يتوافق مع فرضية قفل المسار.

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

تدعي الورقة أن TraFL يوفر حلاً مبدئياً لفشل "البحث عن النمط" (Mode-seeking) في تدريب RL التقليدي للنماذج الانتشارية. من خلال الانتقال من تعظيم المكافأة إلى موازنة المسار، تحافظ الطريقة على تنوع مسارات إزالة الضجيج الضرورية لتغطية حلول صالحة متعددة.

يؤكد المؤلفون أن TraFL هو الطريقة الوحيدة التي تم تقييمها والتي تتحسن عن النموذج الأساسي في كل إعداد لطول الاختبار تم اختباره. ويشددون على أن التحسينات قوية تجاه زيادة ميزانية أخذ العينات وتعمم على مجموعات البيانات المستبعدة، مما يشير إلى أن هدف موازنة المسار يفصل بفعالية بين إشارة المكافأة وآلية تخصيص المسار. يدعو العمل إلى استخدام أهداف مطابقة التوزيع بدلاً من تعظيم المكافأة الصرف لتدريب نماذج dLLM في مرحلة ما بعد التدريب، مما يوفر إطاراً عملياً يتجنب أنماط فشل التركيز المتأصلة في أساليب تدرج السياسة القياسية.

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