ملخص تقني: محولات التدفق المزدوج (Dual-Flow Transformers)

بيان المشكلة

تستحوذ تكلفة خدمة النماذج اللغوية الكبيرة (LLMs) بشكل متزايد على مرحلة الاستدلال (inference) بدلاً من التدريب. ومع ذلك، فإن مرحلتي الاستدلال الأساسيتين — التعبئة المسبقة (prefill) (معالجة رموز المطالبة/البرومبت) و فك التشفير التكراري (autoregressive decode) (توليد الرموز) — تظهران اختناقات عتادية مختلفة جوهرياً:

• التعبئة المسبقة (Prefill): المعالجة المتوازية لرموز المطالبة تكون عادةً مقيدة بالحوسبة (compute-bound)، أي مقيدة بإنتاجية العمليات الحسابية.
• فك التشفير (Decode): التوليد المتسلسل يكون عادةً مقيداً بنطاق ذاكرة الوصول العشوائي (memory-bandwidth-bound)، أي مقيد بالحاجة إلى قراءة مخزن مفاتيح-قيم (KV cache) المتنامي وأوزان النموذج بشكل متكرر لكل رمز جديد.

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

المنهجية: محول التدفق المزدوج (The Dual-Flow Transformer)

يقترح المؤلفون محول التدفق المزدوج (Dual-Flow Transformer)، وهو بنية تفصل هيكلياً بين مسار معالجة المطالبة الأساسي ومسار توقع الاستمرارية المساعد.

البنية الأساسية

يتكون النموذج من مسارين لـ "بقايا الحالة" (residual state trajectories) يتشاركان معظم بارامترات النموذج ولكنهما يختلفان في إدارة الحالة ونطاق التنشيط:

1. التدفق الأساسي (Primary Flow): مسار "ترانسفورمر" سببي كامل ومعياري. يقوم بمعالجة المطالبة بأكملها، ويكتب مخزن KV الثابت، ويحدد الحالة الأساسية. وهو نظام مكتفٍ بذاته ولا يعتمد على التدفق المساعد.
2. التدفق المساعد (Auxiliary Flow): مسار بقايا يقرأ الحالة الأساسية والمخزن، ولكنه لا يكتب أبداً في مخزن KV الثابت و لا يغذي أبداً التدفق الأساسي بالعودة إليه.
   • التنشيط: يتم استبعاد التدفق المساعد أثناء المعالجة الأولية للمطالبة (prefill). ويتم تنشيطه فقط عند الموضع النهائي للمطالبة (للتنبؤ بأول رمز استمرارية) وعند كل موضع توليد لاحق.
   • مشاركة الأوزان: يتشارك كلا التدفقين جميع مصفوفات الانتباه (attention)، والـ MLP، والمخرجات (output).
   • الفصل: يستخدمان تمثيلات رموز (embeddings) منفصلة ومتجهات اقتران خفيفة الوزن ومتعلمة.

الآليات الرئيسية

• انتباه مشترك غير متماثل لـ KV (Asymmetric Shared-KV Attention): يقوم التدفق الأساسي بعملية انتباه ذاتي سببي قياسي. يقوم التدفق المساعد بتوليد استعلام جديد ($Q_2$) يتم اقترانه مع استعلام التدفق الأساسي ($Q_1$) عبر متجه مُتعلم ($a_\ell$). ومن الأهمية بمكان أن التدفق المساعد ينتبه إلى مفاتيح وقيم التدفق الأساسي ($K_1, V_1$). وهذا يسمح للتدفق المساعد بالاستفادة من الحالة المخزنة دون كتابة مخزنه الخاص، مما يتيح التنفيذ المجموعاتي حيث يتم تحميل الأوزان ومناطق KV مرة واحدة لخدمة كلا التدفقين.
• الاقتران الخفيف (Lightweight Coupling): تتدفق المعلومات أحادية الاتجاه من التدفق الأساسي إلى التدفق المساعد عبر متجهات مُتعلمة ($a_\ell, b_\ell, c_\ell$) عند مستوى الاستعلام، والـ MLP المتوسط، ومخرجات الـ MLP. وهذا يثري تمثيل التدفق المساعد دون خلق تبعيات من شأنها تغيير المسار الأساسي.
• هدف الخليط (Mixture Objective): التنبؤ النهائي هو خليط من توزيعات التدفق الأساسي والمساعد: $p_{mix} = \alpha_t p_{primary} + (1-\alpha_t) p_{aux}$. يتم ضبط وزن الخلط $\alpha_t$ ديناميكياً بناءً على الثقة (احتمالية التصادم) لتوزيع التدفق الأساسي، مما يضمن بقاء التدفق الأساسي مهيمناً (بوزن لا يقل عن 50%) أثناء الاستدلال.
• إعادة تشغيل الموجه (Router Replay - لنموذج MoE): في تكوينات "خليط الخبراء" (Mixture-of-Experts)، يعيد التدفق المساعد استخدام نفس مجموعة الخبراء المختارة من قبل التدفق الأساسي (التي حددها موجه التدفق الأساسي) ولكنه يطبقها بأوزان الخليط الخاصة به. هذا يحافظ على منطق اختيار الخبراء للتدفق الأساسي مع زيادة العمليات الحسابية التي تُجرى على الخبراء.

تكاليف الاستدلال مقابل التدريب

• الاستدلال: لا تتحمل مرحلة التعبئة المسبقة (prefill) أي تكلفة إضافية مقارنة بالترانسفورمر القياسي؛ حيث يتم تخطي التدفق المساعد تماماً. أما مرحلة فك التشفير (decode) فتتطلب حوسبة إضافية (تقريباً ضعف العمليات الحسابية لكل رمز) ولكنها لا تزيد من حجم مخزن KV أو مجموعة الأوزان الفريدة التي يتم قراءتها، حيث يتم تجميع التدفقين معاً.
• التدريب: يتم تقييم كلا التدفقين عبر التسلسل الكامل، مما يؤدي إلى حوالي ضعف عمليات (FLOPs) التدريب للنموذج الأساسي. ومع ذلك، فإن الزيادة في عدد البارامترات ضئيلة جداً (فقط التضمينات ومتجهات الاقتران المنفصلة).

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

1. البنية: تقديم تصميم "الـ KV المشترك غير المتماثل" الذي يفصل بين حوسبة نطاق المطالبة وحوسبة الاستمرارية. يتعامل التدفق الأساسي مع استمرارية الحالة، بينما يوفر التدفق المساعد مساراً "باقياً" (residual) لتعزيز التنبؤ دون تغيير الحالة الأساسية.
2. الجودة التجريبية: أظهر انخفاضاً مستمراً في خسارة التحقق (validation loss) عبر ثلاثة محاور تجريبية:
   • توسع البيانات (NanoGPT): يتفوق Dual-Flow على الترانسفورمرز القياسية عند تساوي ميزانيات الرموز، مع أرضية خسارة ملائمة أقل.
   • توسع حجم النموذج (Dense LLaMA): تستمر المكاسب مع زيادة حجم النموذج.
   • الـ MoE المتناثر (Sparse MoE): تم التحقق من ذلك في تكوينات LLaMA-style MoE مع خاصية "إعادة تشغيل الموجه".
3. تخصيص MoE النوعي للمرحلة: على عكس التصميمات السابقة ذات التدفق المتوازي التي تطبق تخصيصاً موحداً للحوسبة، يسمح Dual-Flow بتكوين "تفرع الخبراء" (expert fan-outs) للتدفق الأساسي والمساعد بشكل مستقل. وهذا يتيح مساحة للمفاضلة بين تكلفة التعبئة المسبقة، وتكلفة فك التشفير، وجودة التنبؤ، مما يسمح للأنظمة بالتحسين لوظائف عمل محددة (مثل تثبيت تكلفة التعبئة المسبقة مع زيادة حوسبة فك التشفيد).

النتائج

• خسارة التحقق: عبر تجارب NanoGPT ذات ميزانية الرموز المتساوية، حقق Dual-Flow باستمرار خسارة تحقق أقل من الترانسفورمرز القياسية. انخفضت أرضية الخسلة التقاربية الملائمة من 2.9416 (Transformer) إلى 2.9013 (Dual-Flow).
• كفاءة الحوسبة: في مقارنة مطابقة لحوسبة التدريب (حيث يؤدي نموذج Dual-Flow بخطوات بيانات أقل إلى أداء نفس عمليات الـ FLOPs المشتركة للعمود الفقري مع ترانسفورمر بخطوات مضاعفة)، حقق Dual-Flow أيضاً خسارة تحقق أقل، مما يشير إلى أن تخصيص الحوسبة عبر تدفقات متفاعلة أكثر فعالية من مجرد زيادة البيانات لتدفق واحد.
• الدراسات الاستقصائية (Ablations): أظهر تحليل المكونات أن بناء Dual البسيط (تضمينات منفصلة، KV مشترك) يتفوق بالفعل على أقرب النماذج السابقة (PHD-2). كما أدى إضافة الاقتران بين التدفقات وخلط الاحتمالات إلى تحسين الأداء بشكل أكبر.

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

يزعم البحث أن Dual-Flow يضع فرصة هيكلية لتخصيص الخبراء النوعي للمرحلة. من خلال فصل مسار المطالبة عن مسار الاستمرارية، تسمح هذه البنية لأنظمة الخدمة بمعاملة حوسبة وقت التوليد كمورد قابل للتوسع دون زيادة تكلفة معالجة المطالبة بشكل تناسبي.

على وجه التحديد، يسلط المؤلفون الض الضوء على أنه في نماذج MoE، يمكن أن يكون لتدفقات (الأساسي والمساعد) تفرعات خبراء مستقلة. وهذا يخلق مفتاح مفاضلة ثلاثي الأبعاد:

1. تكلفة التعبئة المسبقة: يحددها استخدام الخبراء في التدفق الأساسي.
2. تكلفة فك التشفير: يحددها مجموع استخدام الخبراء في التدفقين الأساسي والمساعد.
3. جودة التنبؤ: تتحسن بزيادة إجمالي الحوسبة المخصصة لمرحلة الاستمرارية.

يسمح هذا بمرونة خاصة بالنشر، حيث يمكن للنظام الحفاظ على تكاليف حوسبة المطالبة مع إنفاق ميزانية فك التشفير المتاحة على المزيد من الخبراء المساعدين، أو العكس، دون تغيير مجموعة البارامترات المشتركة أو هيكل الـ KV الأساسي الواحد. يدعم هذا العمل فرضية أن حوسبة "الرمز التالي" هي مورد حوسبي مفيد يمكن تخصيصه داخلياً داخل النموذج بدلاً من الخارج عبر مسارات استدلال أطول أو تعدد العينات.