YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference

यहाँ शोध पत्र "YOCO++: कुशल LLM इन्फरेंस के लिए KV रेसिडुअल कनेक्शन के साथ YOCO को उन्नत करना" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

ट्रांसफॉर्मर आर्किटेक्चर पर आधारित लार्ज लैंग्वेज मॉडल्स (LLMs) इन्फरेंस के दौरान एक महत्वपूर्ण बाधा का सामना करते हैं: की-वैल्यू (KV) कैश। जैसे-जैसे सीक्वेंस की लंबाई बढ़ती है, KV कैश को स्टोर करने के लिए आवश्यक मेमोरी, बैच साइज और थ्रूपुट के लिए एक सीमित कारक बन जाती है।

इसे संबोधित करने के लिए, क्रॉस-लेयर KV कम्प्रेशन विधियों का प्रस्ताव दिया गया है। ये विधियाँ केवल कुछ लेयर्स के लिए KVs को बनाए रखकर और दूसरों में उन्हें पुन: उपयोग करके मेमोरी खपत को कम करती हैं। एक प्रमुख विधि, YOCO (You Only Cache Once), मध्य परत (middle layer) के KVs को सभी ऊपरी-आधी (top-half) परतों के साथ साझा करती है, जिससे KV कैश के आकार में 50% की कमी आती है।

मुख्य समस्या: जबकि क्रॉस-लेयर शेयरिंग दक्षता में सुधार करती है, यह गैर-नगण्य प्रदर्शन गिरावट (non-negligible performance degradation) भी लाती है। एक परत से दूसरी परत में सीधे KVs साझा करने से मॉडल की क्षमता कमजोर हो जाती है क्योंकि अलग-अलग परतें अद्वितीय रिप्रजेंटेशन (unique representations) सीखने की अपनी क्षमता खो देती हैं। मौजूदा समाधान जैसे FusedKV, कई परतों के KVs के भारित संयोजन (weighted combinations) का उपयोग करके इसे कम करने का प्रयास करते हैं, लेकिन वे महत्वपूर्ण कैश I/O ओवरहेड पेश करते हैं, जो इन्फरेंस थ्रूपुट को कम कर देता है।

2. कार्यप्रणाली: YOCO++ (Methodology: YOCO++)

लेखक YOCO++ का प्रस्ताव करते हैं, जो YOCO का एक उन्नत संस्करण है जो अतिरिक्त I/O ओवरहेड या इन्फरेंस दक्षता से समझौता किए बिना मॉडल के प्रदर्शन में सुधार करता है।

मुख्य तंत्र: भारित KV रेसिडुअल कनेक्शन (Weighted KV Residual Connections)

KVs को सीधे साझा करने के बजाय, YOCO++ बॉटम लेयर के KVs और प्रत्येक क्रमिक बॉटम-हाफ लेयर के KVs के बीच एक भारित रेसिडुअल कनेक्शन (weighted residual connection) पेश करता है।

• आर्किटेक्चर: मॉडल में $L$ परतें होती हैं। निचली $L/2$ परतें "सेल्फ-डिकोडर" (self-decoders) हैं, और ऊपरी $L/2$ "क्रॉस-डिकोडर" (cross-decoders) हैं।
• KV संयोजन: प्रत्येक बॉटम-हाफ लेयर $i$ (जहाँ $i \in \{2, ..., L/2\}$) में, कैश किए गए KVs ($K'_i, V'_i$) वर्तमान परत के KVs ($K_i, V_i$) और बॉटम लेयर के KVs ($K_1, V_1$) के भारित योग के रूप में गणना किए जाते हैं:
  $$K'_i = \lambda(\alpha^K_{i,1}K_1 + \alpha^K_{i,2}K_i)$$
  $$V'_i = \lambda(\alpha^V_{i,1}V_1 + \alpha^V_{i,2}V_i)$$
  जहाँ $\alpha$ सीखने योग्य रेसिडुअल वेट्स (learnable residual weights) हैं, और $\lambda$ एक स्केलिंग फैक्टर है।
• टॉप-हाफ पुन: उपयोग: ऊपरी-आधी परतें मध्य परत ($K'_{L/2}, V'_{L/2}$) से संशोधित KVs का पुन: उपयोग करती हैं, जिनमें अब मूल YOCO मध्य-लेयर KVs की तुलना में अधिक समृद्ध जानकारी है।

प्रमुख तकनीकी नवाचार (Key Technical Innovations)

1. कैश रणनीति: FusedKV में देखे गए I/O ओवरहेड (जिसके लिए डिकोडिंग के दौरान कई परतों के KVs को फेच करने की आवश्यकता होती है) से बचने के लिए, YOCO++ मूल KVs के बजाय संयोजित KVs ($K', V'$) को कैश करता है। यह सुनिश्चित करता है कि डिकोडिंग चरण के दौरान, मॉडल प्रति परत केवल एक ही KV सेट तक पहुँचता है, जिससे YOCO के समान ही I/O प्रोफाइल बना रहता है।
2. स्केलिंग फैक्टर ($\lambda$): लेखकों ने देखा कि स्केलिंग फैक्टर के बिना, मॉडल रेसिडुअल वेट्स को प्रभावी ढंग से सीखने में संघर्ष करता है (वे इनिशियलाइजेशन के करीब रहते हैं)। वे रेसिडुअल सिग्नल को बढ़ाने के लिए एक फिक्स्ड हाइपरपैरामीटर $\lambda > 1$ (प्रायोगिक रूप से 35 निर्धारित) पेश करते हैं, जिससे वेट्स को बेहतर ढंग से सीखने में मदद मिलती है।
3. इनिशियलाइजेशन: वेट्स को इस तरह इनिशियलाइज किया गया है कि $\alpha_{i,1}=0$ और $\alpha_{i,2}=1/\lambda$ हो, जिससे यह सुनिश्चित होता है कि मॉडल मानक YOCO के समान स्थिति से शुरू होता है।

3. दक्षता विश्लेषण (Efficiency Analysis)

शोध पत्र YOCO और FusedKV की तुलना में YOCO++ की दक्षता का कड़ाई से विश्लेषण करता है:

• पैरामीटर्स: नगण्य वृद्धि (केवल 4 सीखने योग्य वेट्स प्रति बॉटम-हाफ लेयर)।
• कंप्यूटेशन: न्यूनतम ओवरहेड; भारित संयोजन अटेंशन मैकेनिज्म की तुलना में बहुत सरल है।
• मेमोरी: YOCO के समान। दोनों $L/2$ परतों को कैश करते हैं।
• I/O ओवरहेड:
  • ट्रेनिंग/प्रीफिल: $K_1, V_1$ और $K_i, V_i$ तक पहुँचने के कारण थोड़ा अधिक I/O, लेकिन यह चरण कंप्यूट-बाउंड है, इसलिए लेटेंसी अप्रभावित रहती है।
  • डिकोडिंग: कोई अतिरिक्त ओवरहेड नहीं। चूंकि संयोजित KVs को प्री-कंप्यूट और कैश किया गया है, इसलिए डिकोडर प्रति लेयर केवल एक ही KV सेट पढ़ता है, जो YOCO के थ्रूपुट से मेल खाता है।

4. प्रयोगात्मक परिणाम (Experimental Results)

SlimPajama डेटासेट का उपयोग करके 1.1B पैरामीटर वाले मॉडल (TinyLlama कॉन्फ़िगरेशन) पर प्रयोग किए गए।

• इन्फरेंस दक्षता (चित्र 3):

  • प्रीफिल लेटेंसी (Prefill Latency): YOCO++, अर्ली एग्जिट (early exit) के माध्यम से मानक ट्रांसफॉर्मर की तुलना में 50% कम लेटेंसी प्राप्त करते हुए, YOCO और FusedKV के बराबर है।
  • डिकोडिंग थ्रूपुट (Decoding Throughput): YOCO++ उच्च थ्रूपुट बनाए रखता है जो YOCO के तुलनीय है और मानक ट्रांसफॉर्मर से काफी बेहतर है। FusedKV के विपरीत, जिसका थ्रूपुट I/O बाधाओं के कारण ट्रांसफॉर्मर के स्तर तक गिर जाता है, YOCO++ कुशल बना रहता है।

• मॉडल प्रदर्शन (तालिका 1 और 2):

  • जीरो-शॉट सटीकता (Zero-Shot Accuracy): YOCO++ 50% कम्प्रेशन रेट पर क्रॉस-लेयर कम्प्रेशन विधियों के बीच स्टेट-ऑफ-द-आर्ट (SOTA) प्रदर्शन प्राप्त करता है।
  • तुलना: यह कॉमनसेन्स रीजनिंग बेंचमार्क (HellaSwag, OBQA, WinoGrande, ARC, PIQA) पर मानक ट्रांसफॉर्मर, YOCO, FusedKV और FusedKV-Lite से बेहतर प्रदर्शन करता है।
  • ट्रेनिंग लॉस: YOCO++ सबसे कम वैलिडेशन लॉस (2.4194) प्राप्त करता है, जो मानक ट्रांसफॉर्मर (2.4266) से भी बेहतर है।

• एब्लेशन स्टडीज (Ablation Studies):

  • स्केलिंग फैक्टर: $\lambda$ को हटाने से प्रदर्शन गिरकर YOCO के स्तर पर आ जाता है, जो इसके महत्व की पुष्टि करता है।
  • की रेसिडुअल्स (Key Residuals): की रेसिडुअल कनेक्शन ($K$) को हटाने से प्रदर्शन में गिरावट आती है, जो यह सिद्ध करता है कि कीज़ (Keys) और वैल्यूज़ (Values) दोनों के लिए रेसिडुअल कनेक्शन फायदेमंद हैं।

5. मुख्य योगदान (Key Contributions)

1. YOCO++ आर्किटेक्चर: एक नवीन विधि जो बॉटम लेयर और बॉटम-हाफ लेयर्स के बीच भारित रेसिडुअल कनेक्शन पेश करके क्रॉस-लेयर KV कम्प्रेशन को बेहतर बनाती है।
2. दक्षता संरक्षण: संयोजित रिप्रजेंटेशन को कैश करके, यह पिछले "फ्यूजन" तरीकों (जैसे FusedKV) से जुड़े I/O ओवरहेड को सफलतापूर्वक समाप्त करता है, जिससे YOCO का उच्च थ्रूपुट बना रहता है।
3. स्केलिंग फैक्टर: डीप नेटवर्क में रेसिडुअल वेट्स को सीखने की अनुकूलन कठिनाई (optimization difficulty) को हल करने के लिए एक स्केलिंग फैक्टर $\lambda$ का परिचय।
4. SOTA प्रदर्शन: यह प्रदर्शित करता है कि 50% की KV कैश कमी के साथ भी ट्रेनिंग लॉस और डाउनस्ट्रीम टास्क सटीकता दोनों में मानक अनकम्प्रेस्ड ट्रांसफॉर्मर से बेहतर प्रदर्शन करना संभव है।

6. महत्व (Significance)

यह कार्य LLM इन्फरेंस दक्षता और मॉडल क्षमता के बीच महत्वपूर्ण ट्रेड-ऑफ को संबोधित करता है। यह दिखाकर कि क्रॉस-लेयर कम्प्रेशन न केवल मानक ट्रांसफॉर्मर के प्रदर्शन के बराबर हो सकता है बल्कि बिना किसी I/O दंड (penalty) के उससे बेहतर भी हो सकता है, YOCO++ संसाधन-बाधित हार्डवेयर पर बड़े मॉडल्स को तैनात करने के लिए एक अत्यधिक व्यावहारिक समाधान प्रदान करता है। यह सुझाव देता है कि "रेसिडुअल कनेक्शन" को प्रभावी रूप से KV कैश डोमेन में लागू किया जा सकता है ताकि परतों के बीच सूचना प्रवाह (information flow) को संरक्षित किया जा सके, जो कुशल LLM डिज़ाइन के लिए एक नया मार्ग प्रदान करता है।