A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs
JoLT टोकन और फीचर अक्षों पर एक आंशिक टकर डिकंपोजिशन (Tucker decomposition) लागू करके और विसर्जित की गई जानकारी को जॉनसन-लिंडेनस्ट्रॉस रोटेटेड लो-बिट रेसिडुअल के माध्यम से पुनर्स्थापित करके, लार्ज लैंग्वेज मॉडल्स के KV कैश का लगभग लॉसलेस 2-3x संपीड़न प्राप्त करता है, जो पर्प्लेक्सिटी और डाउनस्ट्रीम कार्यों पर बेसलाइन प्रदर्शन बनाए रखने के लिए एक एकीकृत बाइट बजट के तहत अनुकूलित है।
मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें
कल्पना कीजिए कि आप एक विशाल, सुपर-स्मार्ट रोबोट दिमाग (एक लार्ज लैंग्वेज मॉडल) चला रहे हैं जो एक कहानी सुनाने की कोशिश कर रहा है। कहानी को जारी रखने के लिए, रोबोट को वर्तमान बातचीत में कही गई हर एक बात को याद रखना होता है। यह अपनी याददाश्त को एक विशेष "नोटबुक" में रखता है जिसे KV Cache कहा जाता है।
समस्या यह है कि जैसे-जैसे कहानी लंबी होती जाती है, यह नोटबुक बहुत बड़ी हो जाती है। यह इतनी बड़ी हो जाती है कि यह रोबोट की सारी मेमोरी खा जाने लगती है, जिससे सब कुछ धीमा हो जाता है। यह ऐसा है जैसे आप मैराथन दौड़ते समय अपने बैकपैक में एक पूरी लाइब्रेरी ले जाने की कोशिश कर रहे हों; अंततः, आप रुक जाते हैं।
वैज्ञानिकों ने पहले भी इस नोटबुक को सिकोड़ने की कोशिश की है। कुछ ने पूरी चीज़ को एक छोटे से बॉक्स में दबाने (क्वांटाइजेशन) की कोशिश की, तो कुछ ने पंक्तियों या कॉलमों में पैटर्न ढूंढकर पन्नों का सारांश निकालने (लो-रैंक मेथड्स) की कोशिश की। लेकिन इस पेपर के लेखक, राहुल कृष्णन और वोल्कर शुल्ज़ ने देखा कि इन अन्य तरीकों ने कुछ मिस कर दिया: नोटबुक केवल कागज़ का एक सपाट ढेर नहीं है। यह एक 3D ब्लॉक है जिसके तीन अलग-अलग पक्ष हैं: Heads (सोचने के अलग तरीके), Tokens (शब्द), और Features (विवरण)।
उन्होंने पाया कि इनमें से दो पक्ष उबाऊ, दोहराव वाली चीज़ों से भरे हुए हैं जिन्हें आसानी से सिकोड़ा जा सकता है, लेकिन अन्य पक्ष अद्वितीय हैं और उन्हें बिना रोबोट की दिमागी शक्ति खोए नहीं सिकोड़ा जा सकता।
बड़ी खोज: "JoLT" विधि
टीम ने एक नया तरीका निकाला जिसे JoLT (जॉइंट टकर एंड JL-रेसिडुअल एलोकेशन) कहा जाता है। इसे आपके बैकपैक के लिए एक सुपर-स्मार्ट पैकिंग सर्विस की तरह समझें।
- स्मार्ट स्क्वीज़ (पार्शियल टकर): पूरे 3D ब्लॉक को सिकोड़ने के बजाय, JoLT डेटा को देखता है और कहता है, "ठीक है, 'Heads' और 'Layers' अद्वितीय और कीमती हैं; इन्हें अकेला छोड़ दो। लेकिन 'Tokens' और 'Features' फालतू चीज़ों से भरे हैं।" इसलिए, यह केवल उन दो विशिष्ट पक्षों को सिकोड़ता है। यह एक विशाल, फूले हुए तकिए को लेने और केवल उसके बीच की हवा को निकालने जैसा है, जबकि मजबूत किनारों को बरकरार रखा जाता है।
- सेफ्टी नेट (JL-रेसिडुअल): जब आप एक तकिए को दबाते हैं, तो कुछ हवा बाहर निकल जाती है। यदि आप इसे ऐसे ही छोड़ देते हैं, तो तकिया चपटा और बेकार हो जाता है। JoLT उस "निकाली गई हवा" (खोई हुई जानकारी) को पकड़ लेता है और उसे एक JL-रेसिडुअल नामक एक छोटे, सुपर-कुशल सेफ्टी नेट में स्टोर करता है। यह नेट पैकिंग में इतना अच्छा है कि यह कम से कम बिट्स में भी गायब विवरणों को रख सकता है।
- परफेक्ट बैलेंस (लैग्रेंजियन ड्यूल): यहाँ जादू वाला हिस्सा है। रोबोट के पास अपनी जगह के लिए एक सख्त बजट होता है (मान लीजिए 1 बाइट)। JoLT एक गणितीय "स्मार्ट एलोकेटर" का उपयोग यह तय करने के लिए करता है कि तकिए को कितना सिकोड़ना है और सेफ्टी नेट को कितनी जगह देनी है। यह महसूस करता है कि मेमोरी के कुछ हिस्से (जैसे "Keys") को सिकोड़ना आसान है, जबकि अन्य (जैसे "Values") जिद्दी हैं और उन्हें अधिक सेफ्टी नेट स्पेस की आवश्यकता है। यह सर्वोत्तम परिणाम प्राप्त करने के लिए बजट को गतिशील रूप से इधर-उधर घुमाता है।
उन्होंने क्या साबित किया (और क्या नहीं)
लेखकों ने परीक्षण के लिए दो प्रसिद्ध रोबोट दिमागों का उपयोग किया: Mistral-7B (जो "ग्रुपड-क्वेरी" शैली का उपयोग करता है) और LLaMA-2-13B (जो "मल्टी-हेड" शैली का उपयोग करता है)।
"फ्री ज़ोन": उन्होंने एक स्वीट स्पॉट पाया जहाँ वे मेमोरी को 2 से 3 गुना (2–3×) सिकोड़ सकते थे और रोबोट का प्रदर्शन बिल्कुल भी नहीं गिरा। यह "नियर-लॉसलेस" (लगभग बिना नुकसान के) था।
- GSM8K (गणित की समस्याएं) और RULER (टेक्स्ट के ढेर में सुई खोजने जैसा काम) जैसे परीक्षणों पर, संकुचित रोबोट का स्कोर अनकंप्रेस्ड वाले के समान ही था, जो सांख्यिकीय शोर (statistical noise) के भीतर था।
- मेमोरी को पुनर्गठित करने में त्रुटि बहुत कम थी: कीज़ (keys) के लिए लगभग 0.009 और वैल्यूज़ (values) के लिए 0.006। यह पिछले तरीकों जैसे 4-बिट क्वांटाइजेशन या क्रॉस-लेयर SVD की तुलना में लगभग 10 गुना बेहतर (एक ऑर्डर ऑफ मैग्नीट्यूड) है।
"क्लिफ" (चट्टान): उन्होंने एक सीमा भी पाई। यदि आप मेमोरी को बहुत अधिक दबाने की कोशिश करते हैं (3× से अधिक), तो चीजें गड़बड़ा जाती हैं।
- Mistral रोबोट धीरे-धीरे प्रदर्शन खोते हुए गरिमापूर्ण तरीके से कम हुआ।
- हालाँकि, LLaMA रोबोट 4× और 5× संपीड़न के बीच एक "क्लिफ" से टकराया। इसका प्रदर्शन तेजी से गिरा, जो 5.39 के स्कोर से 9.07 (परप्लेक्सिटी में एक बड़ा उछाल, जिसका अर्थ है कि यह शब्दों की भविष्यवाणी करने में बहुत खराब हो गया) तक पहुँच गया।
तेज़ संस्करण: FlashJoLT
परफेक्ट स्क्वीज़ की गणना करने में समय लगता है। इसे ठीक करने के लिए, उन्होंने FlashJoLT बनाया। हर बार भारी गणित करने के बजाय, यह मुख्य पैटर्न का अनुमान लगाने के लिए एक "रैंडमाइज्ड" शॉर्टकट का उपयोग करता है।
- परिणाम: यह मेमोरी को कंप्रेस करने में 5 से 13 गुना तेज़ है, लेकिन इसकी गुणवत्ता धीमी, परफेक्ट वर्शन के समान ही रहती है।
उन्होंने किसे खारिज किया
पेपर स्पष्ट रूप से बताता है कि इस विशिष्ट समस्या के लिए क्या काम नहीं करता है:
- सब कुछ सिकोड़ना: तीनों पक्षों (Heads, Tokens, और Features) को सिकोड़ने की कोशिश करना एक बुरा विचार है। "Heads" और "Layers" बहुत अद्वितीय हैं; उन्हें सिकोड़ने से रोबोट की बुद्धि को नुकसान पहुँचता है।
- फिक्स्ड-बिट क्वांटाइजेशन: हर संख्या के लिए बिट्स की संख्या को कम करना (जैसे हर चीज़ को 4 बिट्स बनाना) 2–3× संपीड़न के "स्वीट स्पॉट" तक नहीं पहुँच सकता। या तो यह पर्याप्त संपीड़न नहीं करता या बहुत अधिक गुणवत्ता खो देता है।
- एक ही आकार सबके लिए (One-size-fits-all): आप "Keys" और "Values" के साथ एक जैसा व्यवहार नहीं कर सकते। "Values" को कंप्रेस करना बहुत कठिन है (2–3× अधिक कठिन), इसलिए उन्हें अलग बजट की आवश्यकता होती है।
निचोड़ (The Bottom Line)
लेखकों ने वास्तविक हार्डवेयर (एक A100 GPU) पर इसका माप किया और पाया कि JoLT एक नियर-लॉसलेस तरीका बनाता है जिससे रोबोट की बुद्धिमत्ता को नुकसान पहुँचाए बिना मेमोरी को 2–3 गुना कम किया जा सकता है।
हालाँकि, वे सावधान करते हैं कि यह हर चीज़ के लिए जादुई समाधान नहीं है।
- यह "फ्री ज़ोन" (2–3×) के लिए बहुत अच्छा काम करता है, लेकिन कुछ रोबोट प्रकारों (जैसे LLaMA) पर अधिक दबाव डालने से गुणवत्ता में भारी गिरावट आती है।
- जबकि मेमोरी स्टोरेज छोटी है, रोबोट को बोलने के दौरान हर बार मेमोरी को "अन-स्क्वीज़" (अनपैक) करने के लिए कुछ गणित करना पड़ता है। वे सुझाव देते हैं कि इसे वास्तविक दुनिया के उपयोग के लिए वास्तव में व्यावहारिक बनाने के लिए, इंजीनियरों को विशेष कंप्यूटर चिप्स (फ्यूज्ड कर्नेल) बनाने की आवश्यकता है जो अनपैक किए बिना सीधे संकुचित मेमोरी को पढ़ सकें।
संक्षेप में, JoLT एक शानदार, गणित-आधारित पैकिंग ट्रिक है जो लंबी बातचीत के लिए बहुत सारा स्थान बचाती है, लेकिन इसकी एक सीमा है, और इसे पूरी गति से चलाने के लिए भविष्य के हार्डवेयर से थोड़े सहयोग की आवश्यकता है।
अपने क्षेत्र के पेपरों की भीड़ में उलझे हुए हैं?
आपके रिसर्च कीवर्ड से मेल खाने वाले सबसे नए और अलग सोच वाले पेपरों का रोज़ाना Digest पाएँ—तकनीकी सारांश के साथ, आपकी भाषा में।