तकनीकी सारांश: रीजनिंग मॉडल्स के लिए वैल्यू-अवेयर स्टोकेस्टिक KV कैश इविक्शन (VASE)

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

रीजनिंग मॉडल्स (जैसे Qwen3, OpenAI का o1) अंतिम उत्तर देने से पहले विस्तृत 'चेन ऑफ थॉट' (chain of thought) उत्पन्न करके उच्च सटीकता प्राप्त करते हैं। हालाँकि, यह क्षमता डिकोडिंग चरण के दौरान एक महत्वपूर्ण मेमोरी और कम्प्यूटेशनल बाधा (bottleneck) पैदा करती है। जैसे-जैसे अनुक्रम की लंबाई (sequence length) बढ़ती है, प्रत्येक पिछले टोकन के प्रतिनिधित्व को संग्रहीत करने के लिए आवश्यक की-वैल्यू (KV) कैश भारी ओवरहेड उत्पन्न करता है।

मौजूदा समाधान दो श्रेणियों में आते हैं:

• सिलेक्शन-आधारित विधियाँ (Selection-based methods): ये पूर्ण KV कैश को बनाए रखती हैं लेकिन अटेंशन कंप्यूटेशन के दौरान केवल टोकन्स के एक विरल (sparse) उपसमुच्चय को सक्रिय करती हैं। ये सटीक तो हैं, लेकिन इनका मेमोरी फुटप्रिंट अनुक्रम की लंबाई के साथ रैखिक रूप से ($O(T)$) बढ़ता है, जिससे मेमोरी बाधा का समाधान नहीं हो पाता।
• इविक्शन-आधारित विधियाँ (Eviction-based methods): ये एक पूर्व-निर्धारित बजट पूरा होने के बाद कम-महत्व वाले KV जोड़ों को स्थायी रूप से हटा (discard) देती हैं, जिससे एक स्थिर मेमोरी फुटप्रिंट और बेहतर थ्रूपुट मिलता है। हालाँकि, वर्तमान इविक्शन विधियाँ रीजनिंग कार्यों पर सिलेक्शन-आधारित विकल्पों की तुलना में महत्वपूर्ण सटीकता गिरावट का सामना करती हैं, जिससे मॉडल अक्सर दोहराव वाले रीजनिंग लूप में फंस जाते हैं या निरर्थक आउटपुट उत्पन्न करते हैं।

यह शोध पत्र पहचानता है कि वर्तमान इविक्शन रणनीतियाँ दो महत्वपूर्ण कारकों पर ध्यान देने में विफल रहती हैं: बड़े-परिमाण वाले वैल्यू स्टेट्स (large-magnitude value states) का असंगत प्रभाव और रिटेन किए गए टोकन्स में स्टोकेस्टिक विविधता (stochastic diversity) की आवश्यकता।

2. कार्यप्रणाली: VASE

लेखक वैल्यू-अवेयर स्टोकेस्टिक KV कैश इविक्शन (VASE) प्रस्तावित करते हैं, जो एक ट्रेनिंग-फ्री इविक्शन फ्रेमवर्क है जिसे दक्षता और सटीकता के बीच के अंतर को पाटने के लिए डिज़ाइन किया गया है। VASE एक आवधिक इविक्शन फ्रेमवर्क (एक निरंतर बजट $K$ और एक हालिया बफर $B$ का उपयोग करते हुए) के भीतर कार्य करता है और दो मुख्य तंत्र पेश करता है:

A. बड़े-परिमाण वाले वैल्यू स्टेट्स का संरक्षण (Protection of Large-Magnitude Value States)

लेखक देखते हैं कि रीजनिंग मॉडल्स में वैल्यू स्टेट्स एक अत्यधिक विषम वितरण (skewed distribution) प्रदर्शित करते हैं, जहाँ टोकन्स का एक छोटा हिस्सा असामान्य रूप से बड़े वेक्टर परिमाण (मैग्निट्यूड) रखता है (जिसे रेंज $Range(v) = \max(v) - \min(v)$ द्वारा मापा जाता है)।

• निष्कर्ष: इन उच्च-परिमाण वाले वैल्यूज को हटाने से विनाशकारी सटीकता पतन (जैसे, GSM8K पर ~88% से गिरकर 14% होना) होता है और ऐसे दोहराव वाले लूप उत्पन्न होते हैं जहाँ मॉडल निष्कर्ष तक पहुँचने के बजाय अंतहीन रूप से संदर्भ की पुन: जांच करता रहता है।
• तंत्र: VASE टोकन बजट का एक विशिष्ट हिस्सा ($N_v$) अनंतिम रूप से (unconditionally) उन $N_v$ टोकन्स को सुरक्षित रखने के लिए आरक्षित करता है जिनका वैल्यू मैग्निट्यूड सबसे अधिक है। यह सुनिश्चित करता है कि सबसे प्रभावशाली वैल्यू वेक्टर्स को कभी हटाया न जाए।

B. स्टोकेस्टिसिटी का परिचय (Introduction of Stochasticity)

वर्तमान इविक्शन विधियाँ अक्सर नियतात्मक (deterministic) टॉप-$k$ चयन का उपयोग करती हैं, जो रिटेन किए गए कैश में विविधता की कमी का कारण बन सकती है।

• निष्कर्ष: स्टोकेस्टिसिटी (यादृच्छिकता) को शामिल करने से सटीकता में सुधार होता है क्योंकि यह पूर्ण संदर्भ के अधिक प्रतिनिधि कवरेज को सुनिश्चित करता है।
• तंत्र: टोकन्स के टॉप-स्कोरिंग को नियत रूप से चुनने के बजाय, VASE वेटेड स्टोकेस्टिक सैंपलिंग का उपयोग करता है।
  • VASE-AttnV: यह अटेंशन स्कोर (जो SnapKV से प्राप्त होता है) के आधार पर स्टोकेस्टिक सैंपलिंग के साथ वैल्यू-अवेयर रिजर्वेशन को जोड़ता है।
  • VASE-DKV: यह CurDKV पद्धति (जो CUR मैट्रिक्स डीकंपोजिशन से लीवरेज स्कोर का उपयोग करती है) को हर इविक्शन स्टेप पर गाऊसी प्रोजेक्शन मैट्रिक्स $G$ को रीसैंपल करके अनुकूलित करता है। यह विशिष्ट प्रतिनिधित्व वाले टोकन्स को लगातार कम स्कोर दिए जाने और स्थायी रूप से हटाए जाने से रोकता है।

3. प्रमुख योगदान (Key Contributions)

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

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

लेखकों ने छह रीजनिंग कार्यों (AIME25/26, HMMT25, GPQA-Diamond, MATH, LiveCodeBench-v6) पर 4× KV कैश कंप्रेशन अनुपात के साथ Qwen3-4B और Qwen3-14B पर VASE का मूल्यांकन किया।

• सिलेक्शन मेथड्स बनाम सटीकता: VASE-AttnV ने दोनों मॉडल आकारों पर सबसे मजबूत सिलेक्शन-आधारित मेथड (SeerAttention-R) की तुलना में उच्च औसत सटीकता प्राप्त की, जबकि एक स्थिर मेमोरी फुटप्रिंट बनाए रखा।
  • Qwen3-4B: VASE-AttnV (59.09%) ने SeerAttention-R (58.81%) और सबसे मजबूत इविक्शन बेसलाइन R-KV (54.69%) को 4.4% से पछाड़ा।
  • Qwen3-14B: VASE-AttnV (65.81%) ने SeerAttention-R (65.37%) के बराबर प्रदर्शन किया और R-KV (60.90%) को 4.9% से पीछे छोड़ा।
• एब्लेशन स्टडीज (Ablation Studies):
  • वैल्यू अवेयरनेस: बड़े-मैग्निट्यूड वाले वैल्यूज के लिए स्लॉट्स आरक्षित करने से बेसलाइन्स की तुलना में GSM8K की सटीकता में 16.2% तक सुधार हुआ।
  • स्टोकेस्टिसिटी: CurDKV में स्टोकेस्टिक सैंपलिंग जोड़ने से Qwen3-14B पर सटीकता में 9.2% का सुधार हुआ।
• दक्षता: VASE-DKV ने उच्चतम थ्रूपुट (16K टोकन पर फुल मॉडल बेसलाइन से 3.1× तेज़) और परीक्षण किए गए सभी तरीकों में सबसे कम पीक मेमोरी उपयोग प्राप्त किया।
• कोड जनरेशन: LiveCodeBench पर, VASE मेथड्स ने SeerAttention-R जैसे सिलेक्शन-आधारित मेथड को काफी बेहतर प्रदर्शन किया, जिसे डोमेन शिफ्ट के साथ संघर्ष करना पड़ा।

5. महत्व और दावे (Significance and Claims)

शोध पत्र का दावा है कि VASE सफलतापूर्वक उस दक्षता-सटीकता अंतराल (efficiency-accuracy gap) को पाटता है जो ऐतिहासिक रूप से इविक्शन विधियों को प्रभावित करता रहा है। बड़े-परिमाण वाले वैल्यू स्टेट्स को प्राथमिकता देकर और स्टोकेस्टिसिटी पेश करके, VASE रीजनिंग मॉडल्स को पूर्ण-कैश या सिलेक्शन-आधारित दृष्टिकोणों से जुड़ी सटीकता से समझौता किए बिना एक स्थिर मेमोरी फुटप्रिंट के साथ संचालित करने में सक्षम बनाता है।

लेखक इस बात पर जोर देते हैं कि वैल्यू-स्टेट मैग्निट्यूड के महत्व के बारे में उनके निष्कर्ष इविक्शन से परे व्यापक निहितार्थ रखते हैं, विशेष रूप से KV कैश क्वांटाइजेशन के लिए, जहाँ बड़े-रेंज वाले वैल्यूज को त्रुटि के प्राथमिक स्रोत के रूप में पहचाना गया है। वे सुझाव देते हैं कि भविष्य के मेमोरी-एफिशिएंट इन्फरेंस मेथड्स को इन महत्वपूर्ण उच्च-परिमाण वाले स्टेट्स की रक्षा करने वाले मिक्स-प्रिसिजन अप्रोच पर विचार करना चाहिए।

अंततः, VASE फ्लैशअटेंशन2 (FlashAttention2) को सपोर्ट करने और लॉन्ग-चेन रीजनिंग मॉडल्स के लिए स्केलेबल इन्फरेंस को सक्षम करने के लिए एक सरल, प्रभावी और ट्रेनिंग-फ्री रेसिपी प्रदान करता है।