2025-11-17T14:34:12.785982

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: बहु-चरणीय तर्क के लिए विलंबता-जागरूक परत छोड़ना

मूल जानकारी

  • पेपर ID: 2510.14211
  • शीर्षक: LiteStage: Latency-aware Layer Skipping for Multi-stage Reasoning
  • लेखक: Beomseok Kang, Jiwon Song, Jae-Joon Kim (Seoul National University)
  • वर्गीकरण: cs.CL, cs.AI
  • प्रकाशन तिथि: 16 अक्टूबर 2025
  • पेपर लिंक: https://arxiv.org/abs/2510.14211
  • कोड लिंक: https://github.com/beomseokg/LiteStage

सारांश

बहु-चरणीय तर्क जटिल समस्याओं को क्रमिक उप-चरणों में विभाजित करके छोटे भाषा मॉडल की तर्क क्षमता को बढ़ाने की एक प्रभावी रणनीति बन गई है। हालांकि, यह विलंबता में वृद्धि की कीमत पर आता है। लेखकों ने देखा है कि मौजूदा अनुकूली त्वरण तकनीकें (जैसे परत छोड़ना) इस सेटिंग में दक्षता और सटीकता को संतुलित करने में कठिनाई का सामना करती हैं, मुख्य रूप से दो महत्वपूर्ण चुनौतियों का सामना करती हैं: (1) चरणों के बीच छोड़ने की संवेदनशीलता में अंतर, (2) अनावश्यक आउटपुट टोकन की पीढ़ी। इन समस्याओं को हल करने के लिए, यह पेपर LiteStage प्रस्तावित करता है, जो बहु-चरणीय तर्क के लिए एक विलंबता-जागरूक परत छोड़ने की रूपरेखा है। LiteStage इष्टतम परत बजट आवंटित करने के लिए चरणबद्ध ऑफ़लाइन खोज और अनावश्यक डिकोडिंग को दबाने के लिए आत्मविश्वास-आधारित ऑनलाइन पीढ़ी प्रारंभिक निकास तंत्र को जोड़ता है। OBQA, CSQA और StrategyQA पर प्रयोग 1.70× तक का त्वरण प्राप्त करते हैं, 4.0% से कम सटीकता हानि के साथ, पिछली प्रशिक्षण-मुक्त परत छोड़ने की विधियों से बेहतर है।

अनुसंधान पृष्ठभूमि और प्रेरणा

समस्या परिभाषा

बहु-चरणीय तर्क जटिल समस्याओं को कई क्रमिक उप-प्रश्नों में विभाजित करके छोटे भाषा मॉडल की तर्क क्षमता को बढ़ाता है। उदाहरण के लिए, TinyThinker तीन-चरणीय तर्क अपनाता है: स्मरण (Recall), विश्लेषण (Analysis) और सारांश (Summary)। हालांकि यह दृष्टिकोण तर्क की गुणवत्ता को प्रभावी ढंग से बढ़ाता है, यह अनिवार्य रूप से अनुमान विलंबता में वृद्धि करता है।

मुख्य चुनौतियाँ

लेखकों ने गहन विश्लेषण के माध्यम से दो महत्वपूर्ण समस्याओं की खोज की है:

  1. चरणों के बीच छोड़ने की संवेदनशीलता में अंतर: विभिन्न तर्क चरणों में परत छोड़ने के प्रति संवेदनशीलता की डिग्री में महत्वपूर्ण अंतर है। प्रयोग दर्शाते हैं कि Stage 3 (सारांश चरण) परत छोड़ने के प्रति सबसे संवेदनशील है, जबकि Stage 1 (स्मरण चरण) अपेक्षाकृत मजबूत है।
  2. अनावश्यक टोकन पीढ़ी: परत छोड़ना प्रति-टोकन कम्प्यूटेशनल लागत को कम करता है, लेकिन अक्सर अधिक टोकन की पीढ़ी की ओर जाता है, जो वास्तव में अंत-से-अंत विलंबता को बढ़ाता है।

मौजूदा विधियों की सीमाएं

मौजूदा परत छोड़ने की विधियां (जैसे SkipDecode, UnifiedSkip, AdaSkip) आमतौर पर एक समान छोड़ने की रणनीति अपनाती हैं, जो बहु-चरणीय तर्क में विभिन्न चरणों की विशेषताओं के अनुकूल नहीं हो सकती हैं, जिससे:

  • संवेदनशील चरणों में अत्यधिक संपीड़न से सटीकता में तीव्र गिरावट
  • परत छोड़ने से उत्पन्न पीढ़ी की लंबाई में वृद्धि की अनदेखी
  • विलंबता-जागरूक अनुकूलन तंत्र की कमी

मुख्य योगदान

  1. LiteStage रूपरेखा प्रस्तावित करना: बहु-चरणीय तर्क के लिए विशेष रूप से डिज़ाइन की गई पहली विलंबता-जागरूक परत छोड़ने की रूपरेखा, जो चरणों के बीच संवेदनशीलता अंतर और अनावश्यक टोकन पीढ़ी की समस्याओं को प्रभावी ढंग से हल करती है।
  2. चरणबद्ध परत बजट आवंटन रणनीति: सबसे धीमे चरण से सबसे तेज़ चरण तक एक लालची खोज एल्गोरिदम डिज़ाइन किया गया है, जो प्रत्येक तर्क चरण के लिए इष्टतम परत छोड़ने का बजट आवंटित करता है।
  3. आत्मविश्वास-संचालित पीढ़ी प्रारंभिक निकास तंत्र: ऑनलाइन आत्मविश्वास निगरानी का परिचय, कम आत्मविश्वास वाली अनावश्यक पीढ़ी को गतिशील रूप से समाप्त करना, अनुमान दक्षता को और बढ़ाना।
  4. उल्लेखनीय प्रदर्शन सुधार: तीन बेंचमार्क डेटासेट पर 1.16-1.70× त्वरण प्राप्त करना, केवल 0.4-4.0% सटीकता हानि के साथ, मौजूदा प्रशिक्षण-मुक्त विधियों को बहुत अधिक पार करना।

विधि विवरण

कार्य परिभाषा

परीक्षण डेटासेट D दिया गया है, लक्ष्य चरणबद्ध परत बजट L खोजना है जो दिए गए सटीकता सीमा ε के भीतर अनुमान विलंबता को कम करता है:

argmin_L (1/|D|) ∑_{d∈D} T(M_L(d))
subject to: A(M_L(d)) ≤ A(M(d)) - ε

जहां T और A क्रमशः अनुमान विलंबता और सटीकता का प्रतिनिधित्व करते हैं, M_L और M क्रमशः परत छोड़ने को लागू करने वाले मॉडल और पूर्ण परत मॉडल का प्रतिनिधित्व करते हैं।

मॉडल आर्किटेक्चर

LiteStage में दो पूरक घटक हैं:

1. ऑफ़लाइन कॉन्फ़िगरेशन (Offline Configuration)

चरण 1: परत महत्व अनुमान

  • उप-परत स्तर पर कोसाइन समानता को महत्व प्रॉक्सी के रूप में अपनाना
  • बहु-सिर आत्म-ध्यान (MHSA) और फीड-फॉरवर्ड नेटवर्क (FFN) की महत्ता अलग से गणना करना:
I^(j)_MHSA = (1/N) ∑_{n=0}^{N-1} cos(MHSA^(j)(x) + x, x)
I^(j)_FFN = (1/N) ∑_{n=0}^{N-1} cos(FFN^(j)(x) + x, x)

चरण 2: परत बजट खोज

  • सबसे धीमे अनुमान चरण से शुरू करके लालची खोज करना
  • सटीकता-विलंबता वक्र बनाना, सटीकता बाधा के तहत इष्टतम विलंबता कॉन्फ़िगरेशन चुनना
  • चरणबद्ध अनुकूलन, यह सुनिश्चित करना कि चरणों के बीच इंटरैक्शन सटीक रूप से परिलक्षित हो

2. ऑनलाइन समायोजन (Online Adjustment)

चरण 3: पीढ़ी प्रारंभिक निकास

  • हाल के n टोकन का आत्मविश्वास कैश बनाए रखना
  • औसत आत्मविश्वास μ_Conf की गणना करना, जब सीमा से नीचे हो तो पीढ़ी को जल्दी समाप्त करना
  • आत्मविश्वास को प्रत्येक टोकन के अधिकतम लॉगिट मान के रूप में परिभाषित करना

तकनीकी नवाचार बिंदु

  1. गैर-समान परत बजट आवंटन: प्रत्येक चरण की संवेदनशीलता अंतर के अनुसार, विभिन्न परत छोड़ने के बजट को अनुकूल रूप से आवंटित करना, संवेदनशील चरणों में अत्यधिक संपीड़न से बचना।
  2. विलंबता-जागरूक अनुकूलन: न केवल सटीकता पर विचार करना, बल्कि वास्तविक अनुमान विलंबता को भी ध्यान में रखना, उन कॉन्फ़िगरेशन को स्वचालित रूप से बाहर करना जो अधिक परतें छोड़ते हैं लेकिन विलंबता अधिक है।
  3. गतिशील पीढ़ी नियंत्रण: आत्मविश्वास निगरानी के माध्यम से पीढ़ी की लंबाई को सक्रिय रूप से नियंत्रित करना, परत छोड़ने से उत्पन्न अनावश्यक टोकन समस्या को कम करना।

प्रयोग सेटअप

डेटासेट

TinyThinker के तीन-चरणीय तर्क प्रवाह को अपनाते हुए, तीन प्रश्नोत्तर बेंचमार्क पर मूल्यांकन:

  • OpenBookQA (OBQA): खुली किताब प्रश्नोत्तर कार्य
  • CommonSenseQA (CSQA): सामान्य ज्ञान तर्क प्रश्नोत्तर
  • StrategyQA: रणनीतिगत तर्क प्रश्नोत्तर

मूल्यांकन मेट्रिक्स

  • सटीकता: प्रश्नोत्तर सही दर
  • त्वरण अनुपात: पूर्ण परत मॉडल के सापेक्ष अनुमान गति सुधार
  • विलंबता: अंत-से-अंत अनुमान समय

तुलना विधियां

  • SkipDecode: क्रमिक गहरी परत छोड़ना
  • UnifiedSkip: आवधिक परत छोड़ना
  • AdaSkip: कोसाइन समानता आधारित उप-परत महत्व अनुमान

कार्यान्वयन विवरण

  • मुख्य रूप से TinyLlama-1.1B-Chat-v1.0 मॉडल का उपयोग
  • 10 epoch के लिए प्रशिक्षण, बैच आकार 16 (OBQA/CSQA) या 24 (StrategyQA)
  • सीखने की दर 5×10^-5
  • 10 पुनरावृत्तियों के लिए स्व-संगति प्रोटोकॉल अपनाना
  • आत्मविश्वास सीमा 0.5 पर सेट, कैश आकार n=5

प्रयोग परिणाम

मुख्य परिणाम

तीन बेंचमार्क डेटासेट पर, LiteStage आधार रेखा विधियों से काफी बेहतर है:

डेटासेटआधार सटीकताLiteStage सटीकतात्वरण अनुपात
OBQA64.0%60.0%1.32×
CSQA54.8%53.2%1.16×
StrategyQA62.4%62.0%1.70×

मुख्य निष्कर्ष

  1. चरण संवेदनशीलता अंतर: एकल-चरण छोड़ने के प्रयोगों के माध्यम से पुष्टि, Stage 3 परत छोड़ने के प्रति सबसे संवेदनशील है, इसका सटीकता वक्र समग्र प्रदर्शन की ऊपरी सीमा को लगभग निर्धारित करता है।
  2. विलंबता विरोधाभास: अधिक परत छोड़ना हमेशा तेज़ अनुमान नहीं लाता है, पीढ़ी की लंबाई में वृद्धि के कारण, कुछ कॉन्फ़िगरेशन वास्तव में विलंबता में वृद्धि का कारण बनते हैं।
  3. आत्मविश्वास पैटर्न: परत छोड़ने वाले मॉडल का टोकन आत्मविश्वास एकदिष्ट रूप से घटता है, जबकि पूर्ण परत मॉडल बाद के चरणों में आत्मविश्वास को पुनः प्राप्त कर सकता है।

विलोपन प्रयोग

गैर-समान परत बजट का प्रभाव:

  • समान परत छोड़ने की संख्या के तहत, LiteStage की सटीकता समान छोड़ने की रणनीति से काफी अधिक है
  • छोड़ी गई परतों की संख्या बढ़ने के साथ, प्रदर्शन अंतर आगे बढ़ता है

पीढ़ी प्रारंभिक निकास का योगदान:

  • हल्के परत छोड़ने में, पीढ़ी प्रारंभिक निकास का प्रभाव बहुत कम है (-0.5% डिकोडिंग चरण)
  • भारी परत छोड़ने में, डिकोडिंग चरणों में 82.5% तक की कमी कर सकता है
  • सटीकता मूल रूप से स्थिर रहती है, कुछ मामलों में थोड़ी सुधार भी होती है

केस विश्लेषण

CSQA के एक विशिष्ट केस के माध्यम से प्रदर्शन, पीढ़ी प्रारंभिक निकास कम आत्मविश्वास वाले अनावश्यक पाठ को प्रभावी ढंग से काट सकता है, जबकि मुख्य तर्क तर्क को पूर्ण रखता है, अंतिम उत्तर सुसंगत रहता है।

संबंधित कार्य

बहु-चरणीय पीढ़ी

  • TinyThinker: स्मरण-विश्लेषण-सारांश तीन-चरणीय तर्क चक्र प्रस्तावित करना
  • DeAR: विघटन-विश्लेषण-पुनर्विचार प्रक्रिया अपनाना
  • CasCoD: विघटन विचार श्रृंखला को कैस्केड तरीके से आसवन करना
  • Self-Discover: तर्क संरचना को गतिशील रूप से संगठित करना

परत छोड़ने की तकनीक

प्रशिक्षण-आधारित विधियां:

  • LayerSkip, DeeBERT, EE-LLM: मध्य परत प्रारंभिक निकास
  • Mixture-of-Depth: मॉडल और राउटर को प्रशिक्षित करने की आवश्यकता

प्रशिक्षण-मुक्त विधियां:

  • SkipDecode: क्रमिक गहरी परत छोड़ना
  • Unified Skipping: आवधिक छोड़ना
  • ShortGPT: कोसाइन समानता आधारित
  • AdaSkip: उप-परत स्तर महत्व अनुमान

पीढ़ी प्रारंभिक निकास

मौजूदा विधियां मुख्य रूप से लंबे तर्क मॉडल पर केंद्रित हैं, मॉडल संपीड़न से उत्पन्न पीढ़ी विस्तार समस्या पर ध्यान देने में कमी है।

निष्कर्ष और चर्चा

मुख्य निष्कर्ष

  1. बहु-चरणीय तर्क में गैर-समान संवेदनशीलता: विभिन्न तर्क चरणों में परत संपीड़न के प्रति संवेदनशीलता की डिग्री में महत्वपूर्ण अंतर है, जिसके लिए विभेदित अनुकूलन रणनीति की आवश्यकता है।
  2. विलंबता-जागरूक अनुकूलन की आवश्यकता: केवल परत छोड़ना पीढ़ी की लंबाई में वृद्धि के कारण विलंबता को खराब कर सकता है, सटीकता और विलंबता दोनों पर विचार करने की आवश्यकता है।
  3. पीढ़ी नियंत्रण की प्रभावशीलता: आत्मविश्वास-आधारित पीढ़ी प्रारंभिक निकास परत छोड़ने से उत्पन्न अनावश्यक पीढ़ी समस्या को प्रभावी ढंग से कम कर सकता है।

सीमाएं

  1. ऑफ़लाइन खोज ओवरहेड: अन्य प्रशिक्षण-मुक्त विधियों की तुलना में, LiteStage के ऑफ़लाइन कॉन्फ़िगरेशन को अधिक कम्प्यूटेशनल संसाधनों की आवश्यकता है (लगभग 1-7.6 घंटे)।
  2. मॉडल आर्किटेक्चर निर्भरता: मुख्य रूप से Llama श्रृंखला मॉडल पर सत्यापित, Qwen जैसी अन्य आर्किटेक्चर पर सीमित प्रभाव।
  3. प्रयोज्यता सीमा: विशेष रूप से बहु-चरणीय तर्क परिदृश्य के लिए, एकल-चरणीय तर्क पर प्रयोज्यता पूरी तरह से सत्यापित नहीं है।

भविष्य की दिशाएं

  1. अधिक मॉडल आर्किटेक्चर तक विस्तार: विभिन्न आर्किटेक्चर की छोड़ने की संवेदनशीलता विशेषताओं का अनुसंधान
  2. गतिशील बजट आवंटन: रनटाइम पर परत बजट को स्वचालित रूप से समायोजित करने की तंत्र विकसित करना
  3. बहु-मोडल तर्क अनुकूलन: दृश्य-भाषा आदि बहु-मोडल तर्क कार्यों तक रूपरेखा का विस्तार

गहन मूल्यांकन

शक्तियां

  1. समस्या पहचान सटीक: बहु-चरणीय तर्क में मुख्य बाधाओं की सटीक पहचान, चरण संवेदनशीलता अंतर और अनावश्यक पीढ़ी समस्याओं सहित।
  2. विधि डिजाइन तर्कसंगत: ऑफ़लाइन-ऑनलाइन संयोजन रूपरेखा डिजाइन चतुर है, अनुकूलन प्रभाव सुनिश्चित करता है और रनटाइम ओवरहेड को नियंत्रित करता है।
  3. प्रयोग डिजाइन व्यापक: विस्तृत प्रेरणा प्रयोगों, विलोपन अध्ययन और केस विश्लेषण के माध्यम से, विधि की प्रभावशीलता को पूरी तरह से सत्यापित करना।
  4. व्यावहारिक मूल्य उच्च: प्रशिक्षण-मुक्त विधि के रूप में, अच्छी व्यावहारिक अनुप्रयोग संभावना है।

कमियां

  1. सैद्धांतिक विश्लेषण अपर्याप्त: चरण संवेदनशीलता अंतर के सैद्धांतिक व्याख्या की कमी, मुख्य रूप से अनुभवजन्य अवलोकन पर निर्भर।
  2. पैरामीटर सेटिंग अनुमानी: आत्मविश्वास सीमा, कैश आकार आदि महत्वपूर्ण पैरामीटर मुख्य रूप से अनुमानी सेटिंग पर आधारित, व्यवस्थित विश्लेषण की कमी।
  3. सामान्यीकरण सीमित: विभिन्न मॉडल आर्किटेक्चर पर प्रदर्शन में बड़ा अंतर, सामान्यीकरण क्षमता में सुधार की आवश्यकता।

प्रभाव

  1. शैक्षणिक योगदान: बहु-चरणीय तर्क में परत छोड़ने के अनुकूलन समस्या का पहली बार व्यवस्थित अध्ययन, संबंधित अनुसंधान के लिए नया दृष्टिकोण प्रदान करना।
  2. व्यावहारिक मूल्य: छोटे भाषा मॉडल के कुशल अनुमान के लिए व्यावहारिक समाधान प्रदान करना, किनारे तैनाती को बढ़ावा देने में सहायता।
  3. पुनरुत्पादनीयता: पूर्ण कोड कार्यान्वयन प्रदान करना, बाद के अनुसंधान और अनुप्रयोग को सुविधाजनक बनाना।

लागू परिदृश्य

LiteStage विशेष रूप से निम्नलिखित परिदृश्यों के लिए उपयुक्त है:

  • संसाधन-सीमित किनारे उपकरण तैनाती
  • बहु-चरणीय तर्क की आवश्यकता वाले जटिल कार्य
  • विलंबता-संवेदनशील वास्तविक समय अनुप्रयोग
  • छोटे भाषा मॉडल अनुमान त्वरण

संदर्भ

पेपर कई महत्वपूर्ण संबंधित कार्यों का हवाला देता है, जिनमें शामिल हैं:

  • TinyThinker (Piao and Park, 2024): बहु-चरणीय तर्क का प्रतिनिधि कार्य
  • AdaSkip (He et al., 2025): उप-परत स्तर परत छोड़ने की नवीनतम विधि
  • Mixture-of-Depths (Raposo et al., 2024): गतिशील कम्प्यूटेशन आवंटन का अग्रणी कार्य

समग्र मूल्यांकन: यह पेपर बहु-चरणीय तर्क में परत छोड़ने के अनुकूलन समस्या के लिए एक नवाचारी समाधान प्रस्तावित करता है, सैद्धांतिक अंतर्दृष्टि और व्यावहारिक प्रभाव दोनों में महत्वपूर्ण योगदान है। हालांकि कुछ सीमाएं हैं, लेकिन छोटे भाषा मॉडल के कुशल अनुमान के लिए नई अनुसंधान दिशा खोलता है, महत्वपूर्ण शैक्षणिक मूल्य और व्यावहारिक महत्व है।