LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning
Kang, Song, Kim
Multi-stage reasoning has emerged as an effective strategy for enhancing the reasoning capability of small language models by decomposing complex problems into sequential sub-stages. However, this comes at the cost of increased latency. We observe that existing adaptive acceleration techniques, such as layer skipping, struggle to balance efficiency and accuracy in this setting due to two key challenges: (1) stage-wise variation in skip sensitivity, and (2) the generation of redundant output tokens. To address these, we propose LiteStage, a latency-aware layer skipping framework for multi-stage reasoning. LiteStage combines a stage-wise offline search that allocates optimal layer budgets with an online confidence-based generation early exit to suppress unnecessary decoding. Experiments on three benchmarks, e.g., OBQA, CSQA, and StrategyQA, show that LiteStage achieves up to 1.70x speedup with less than 4.0% accuracy loss, outperforming prior training-free layer skipping methods.
academic
LiteStage: تخطي الطبقات مع الوعي بالكمون للاستدلال متعدد المراحل
أصبح الاستدلال متعدد المراحل استراتيجية فعّالة لتعزيز قدرات الاستدلال في نماذج اللغة الصغيرة من خلال تقسيم المشاكل المعقدة إلى مراحل فرعية متسلسلة. ومع ذلك، يأتي هذا على حساب زيادة الكمون. لاحظ المؤلفون أن تقنيات التسريع التكيفية الموجودة (مثل تخطي الطبقات) تواجه صعوبة في موازنة الكفاءة والدقة في هذا السياق، وتواجه بشكل أساسي تحديين رئيسيين: (1) الاختلاف في حساسية التخطي بين المراحل، (2) توليد رموز الإخراج الزائدة. لمعالجة هذه المشاكل، يقترح هذا البحث LiteStage، وهو إطار عمل لتخطي الطبقات مع الوعي بالكمون للاستدلال متعدد المراحل. يجمع LiteStage بين البحث غير المتصل المرحلي الذي يخصص ميزانية الطبقات المثلى وآلية الخروج المبكر من التوليد القائمة على الثقة لقمع فك التشفير غير الضروري. تُظهر التجارب على ثلاث معايير (OBQA و CSQA و StrategyQA) أن LiteStage يحقق تسريعاً يصل إلى 1.70× مع فقدان دقة أقل من 4.0%، متفوقاً على طرق تخطي الطبقات السابقة الخالية من التدريب.
يعزز الاستدلال متعدد المراحل قدرات الاستدلال في نماذج اللغة الصغيرة من خلال تقسيم المشاكل المعقدة إلى عدة مشاكل فرعية متتالية. على سبيل المثال، يعتمد TinyThinker على استدلال ثلاثي المراحل: الاسترجاع (Recall)، والتحليل (Analysis)، والملخص (Summary). بينما تحسّن هذه الطريقة جودة الاستدلال بشكل فعّال، فإنها تزيد حتماً من كمون الاستدلال.
من خلال التحليل المتعمق، اكتشف المؤلفون مشكلتين رئيسيتين:
اختلاف حساسية التخطي بين المراحل: توجد اختلافات كبيرة في درجة حساسية مراحل الاستدلال المختلفة لتخطي الطبقات. تُظهر التجارب أن المرحلة 3 (مرحلة الملخص) هي الأكثر حساسية لتخطي الطبقات، بينما المرحلة 1 (مرحلة الاسترجاع) أكثر قوة نسبياً.
توليد الرموز الزائدة: بينما يقلل تخطي الطبقات من التكلفة الحسابية لكل رمز، فإنه غالباً ما يؤدي إلى توليد رموز أكثر، مما يزيد في الواقع من الكمون من طرف إلى طرف.
تعتمد طرق تخطي الطبقات الموجودة (مثل SkipDecode و UnifiedSkip و AdaSkip) عادةً على استراتيجية تخطي موحدة، وغير قادرة على التكيف مع خصائص المراحل المختلفة في الاستدلال متعدد المراحل، مما يؤدي إلى:
ضغط مفرط في المراحل الحساسة يسبب انخفاضاً حاداً في الدقة
تجاهل مشكلة زيادة طول التوليد الناجمة عن تخطي الطبقات
اقتراح إطار عمل LiteStage: أول إطار عمل متخصص لتخطي الطبقات مع الوعي بالكمون للاستدلال متعدد المراحل، يحل بفعالية مشاكل اختلاف الحساسية بين المراحل وتوليد الرموز الزائدة.
استراتيجية توزيع ميزانية الطبقات المرحلية: تصميم خوارزمية بحث جشعة من المرحلة الأبطأ إلى الأسرع، لتخصيص ميزانية تخطي الطبقات المثلى لكل مرحلة استدلال.
آلية الخروج المبكر من التوليد المدفوعة بالثقة: إدخال مراقبة الثقة عبر الإنترنت، لإنهاء التوليد الزائد منخفض الثقة بشكل ديناميكي، مما يعزز كفاءة الاستدلال بشكل أكبر.
تحسينات الأداء الكبيرة: تحقيق تسريع بمعدل 1.16-1.70× على ثلاث مجموعات بيانات معيارية، مع فقدان دقة يبلغ فقط 0.4-4.0%، متفوقاً بشكل كبير على الطرق الخالية من التدريب الموجودة.
توزيع ميزانية الطبقات غير المنتظم: توزيع ميزانيات تخطي الطبقات المختلفة بشكل تكيفي وفقاً لاختلافات حساسية كل مرحلة، مما يتجنب الضغط المفرط في المراحل الحساسة.
التحسين مع الوعي بالكمون: عدم الاعتبار فقط للدقة، بل أيضاً الكمون الفعلي للاستدلال، مما يستبعد تلقائياً التكوينات التي تتخطى طبقات أكثر لكن لها كمون أعلى.
التحكم الديناميكي في التوليد: التحكم الاستباقي في طول التوليد من خلال مراقبة الثقة، مما يخفف من مشكلة الرموز الزائدة الناجمة عن تخطي الطبقات.
اختلاف حساسية المراحل: تؤكد التجارب ذات المرحلة الواحدة أن المرحلة 3 هي الأكثر حساسية لتخطي الطبقات، وأن منحنى دقتها يحدد تقريباً الحد الأعلى للأداء الكلية.
مفارقة الكمون: تخطي طبقات أكثر لا يؤدي دائماً إلى استدلال أسرع، وبسبب زيادة طول التوليد، قد تؤدي بعض التكوينات في الواقع إلى زيادة الكمون.
أنماط الثقة: تُظهر نماذج تخطي الطبقات اتجاهاً متناقصاً بشكل رتيب في ثقة الرموز، بينما قد تستعيد النماذج ذات الطبقات الكاملة الثقة في المراحل اللاحقة.
من خلال حالة محددة من CSQA، يُظهر أن الخروج المبكر من التوليد يمكنه قطع النصوص الزائدة منخفضة الثقة بفعالية، مع الحفاظ على منطق الاستدلال الأساسي سليماً، والإجابة النهائية تبقى متسقة.
عدم الانتظام في الحساسية في الاستدلال متعدد المراحل: توجد اختلافات كبيرة في درجة حساسية مراحل الاستدلال المختلفة لضغط الطبقات، مما يتطلب استراتيجيات تحسين متمايزة.
ضرورة التحسين مع الوعي بالكمون: قد يؤدي تخطي الطبقات البسيط إلى تدهور الكمون بسبب زيادة طول التوليد، مما يتطلب الأخذ في الاعتبار الدقة والكمون معاً.
فعالية التحكم في التوليد: يمكن للخروج المبكر من التوليد القائم على الثقة أن يخفف بفعالية من مشكلة الرموز الزائدة الناجمة عن تخطي الطبقات.
نقص التحليل النظري: غياب التفسير النظري لاختلاف حساسية المراحل، يعتمد بشكل أساسي على الملاحظات التجريبية.
إعدادات المعاملات الاستكشافية: المعاملات الرئيسية مثل حد الثقة وحجم الذاكرة المؤقتة تعتمد بشكل أساسي على الإعدادات الاستكشافية، مع نقص التحليل المنهجي.
قابلية التعميم المحدودة: الأداء على بنى نماذج مختلفة متغيرة بشكل كبير، مع قدرة تعميم تحتاج إلى تحسين.
تستشهد الورقة بعدة أعمال ذات صلة مهمة، بما في ذلك:
TinyThinker (Piao and Park, 2024): العمل التمثيلي للاستدلال متعدد المراحل
AdaSkip (He et al., 2025): أحدث طريقة لتخطي الطبقات على مستوى الطبقة الفرعية
Mixture-of-Depths (Raposo et al., 2024): العمل الرائد في توزيع الحسابات الديناميكي
التقييم الشامل: تقترح هذه الورقة حلاً مبتكراً لمشكلة تحسين تخطي الطبقات في الاستدلال متعدد المراحل، مع مساهمات كبيرة في الرؤى النظرية والتأثيرات العملية. على الرغم من وجود بعض القيود، فإنها تفتح اتجاهاً بحثياً جديداً لتسريع الاستدلال الفعّال في نماذج اللغة الصغيرة، مع قيمة أكاديمية وعملية مهمة.