🤖 machine learning

SAKI: Score-Aware Low-Rank Key Indexing for Long-Context KV Retrieval

SAKI एक प्रशिक्षण-मुक्त (training-free), स्कोर-जागरूक (score-aware) लो-रैंक की इंडेक्सिंग विधि है जो एक क्लोज्ड-फॉर्म एसिमेट्रिक फैक्टराइजेशन के माध्यम से अटेंशन स्कोर डिस्टॉर्शन को सीधे कम करके KV कैश कंप्रेशन को अनुकूलित करती है, जिससे यह कई लार्ज लैंग्वेज मॉडल्स में लॉन्ग-कॉन्टेक्स्ट रिट्रीवल रिकॉल के मामले में PCA जैसे मौजूदा की रिकंस्ट्रक्शन-आधारित दृष्टिकोणों से काफी बेहतर प्रदर्शन करती है।

मूल लेखक: Lin Zhang

प्रकाशित 2026-08-05

📖 4 मिनट में पढ़ें☕ कॉफ़ी ब्रेक में पढ़ें

CC BY 4.0

मूल लेखक: Lin Zhang

मूल पेपर CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) के तहत लाइसेंस किया गया है। ✨ नीचे दिए गए पेपर की यह व्याख्या AI से तैयार की गई है। इसे लेखकों ने न तो लिखा है, न इसका समर्थन किया है। तकनीकी सटीकता के लिए मूल पेपर देखें। पूरा डिस्क्लेमर पढ़ें

कल्पना कीजिए कि आप घास के ढेर में एक विशिष्ट सुई खोजने की कोशिश कर रहे हैं, लेकिन वह घास का ढेर एक छोटे शहर के आकार का है, और आपको हर बार सवाल पूछने पर ऐसा ही करना पड़ता है। यह आधुनिक आर्टिफिशियल इंटेलिजेंस (AI) की दैनिक वास्तविकता है जब वह एक लंबी बातचीत या एक विशाल दस्तावेज़ को याद रखने की कोशिश करता है। AI की "याददाश्त" (जिसे KV कैश कहा जाता है) इतनी विशाल हो जाती है कि यह एक सामान्य लैपटॉप की क्षमता से भी अधिक जगह घेर लेती है, जिससे इसे खोजना धीमा और महंगा हो जाता है। इसे ठीक करने के लिए, इंजीनियर एक तरकीब का उपयोग करते हैं: घास के हर एक तिनके को देखने के बजाय, वे एक त्वरित, मोटा नक्शा (एक इंडेक्स) बनाते हैं ताकि यह अनुमान लगाया जा सके कि महत्वपूर्ण सुइयां कहाँ छिपी हो सकती हैं। यदि नक्शा अच्छा है, तो AI सुई को तेज़ी से ढूंढ लेता है; यदि नक्शा खराब है, तो AI भ्रमित हो जाता है और गलत उत्तर देता है।

लंबे समय तक, वैज्ञानिकों ने इन नक्शों को दो चीजों को देखकर बनाने की कोशिश की: या तो AI के मस्तिष्क के "ब्लूप्रिंट" (इसके वेट्स) को, या उस डेटा के "आकार" (कीज़ का वेरिएंस) को जो वह थामे हुए था। इसे ऐसे समझें जैसे कि बिना शीर्षक पढ़े, केवल किताब के स्पाइन के रंग को देखकर या केवल पन्नों की मोटाई को देखकर पुस्तकालय को व्यवस्थित करने की कोशिश करना। समस्या यह है कि AI वास्तव में स्पाइन के रंगों या पन्नों की मोटाई की परवाह नहीं करता है; उसे इस बात से फर्क पड़ता है कि एक विशिष्ट प्रश्न एक विशिष्ट उत्तर के साथ कितनी अच्छी तरह मेल खाता है। SAKI नामक यह पेपर तर्क देता है कि पुराने नक्शे महत्व को मापने के लिए गलत पैमाने का उपयोग कर रहे थे। लेखक ने महसूस किया कि एक आदर्श नक्शा बनाने के लिए, आपको यह मापना होगा कि एक प्रश्न और एक उत्तर वास्तव में आपस में कितनी अच्छी तरह "क्लिक" करते हैं, न कि केवल सामान्य आकारों या ब्लूप्रिंट के आधार पर अनुमान लगाना।

यह पेपर SAKI (स्कोर-अवेयर लो-रैंक की इंडेक्सिंग) नामक एक नई विधि पेश करता है। एक सामान्य पैमाने के बजाय, SAKI एक कस्टम, "स्कोर-अवेयर" नक्शा बनाता है जो सटीक रूप से भविष्यवाणी करता है कि AI के प्रश्न उसकी संग्रहीत यादों के साथ कितनी अच्छी तरह मेल खाएंगे। लेखक ने LLaMA-3.1-8B और Qwen2.5-7B सहित कई लोकप्रिय AI मॉडलों पर इस नए नक्शे का परीक्षण किया। उन्होंने पाया कि SAKI, पिछले सर्वोत्तम तरीकों की तुलना में सही सुइयों को खोजने में काफी बेहतर है। उदाहरण के लिए, जब नक्शे को एक छोटे आकार (रैंक 32) में संकुचित किया गया, तो SAKI ने AI की सही जानकारी को याद करने की क्षमता में उन त्रुटियों को 13% से 30% तक कम कर दिया जो पुराने तरीके अभी भी करते थे। LLaMA-3.1-8B मॉडल पर, इसने सफलता दर को 0.748 से बढ़ाकर 0.799 कर दिया, और Qwen2.5-7B पर, यह 0.786 से बढ़कर 0.850 हो गई।

लेखक बताते हैं कि पुराने तरीके विफल रहे क्योंकि उन्होंने AI की स्मृति को डेटा के एक स्थिर ढेर की तरह माना, यह अनदेखा करते हुए कि AI के प्रश्न उस डेटा के महत्व को बदल देते हैं। उन्होंने दिखाया कि AI की आंतरिक "स्कोरिंग मशीन" अजीब और असंतुलित (गणितीय रूप से, यह "नॉन-नॉर्मल" है) है, जिसका अर्थ है कि डेटा को सिकोड़ने के मानक तरीके (जैसे PCA) गलत हिस्सों को काट देते हैं। SAKI इसे एक विशेष गणितीय शॉर्टकट का उपयोग करके ठीक करता है जो प्रश्न और उत्तर दोनों को एक साथ ध्यान में रखता है। यह पेपर सिद्ध करता है कि यह नया दृष्टिकोण केवल एक भाग्यशाली अनुमान नहीं है; गणित लगभग पूर्ण सटीकता (0.997 का सहसंबंध) के साथ परिणामों की भविष्यवाणी करता है। हालांकि लेखक ने उल्लेख किया है कि उन्होंने अभी तक हर संभव प्रकार के टेक्स्ट या पूर्ण एंड-टू-एंड बातचीत पर इसका परीक्षण नहीं किया है, उनके माप दर्शाते हैं कि SAKI, मॉडलों को फिर से प्रशिक्षित किए बिना AI की स्मृति को तेज़ और स्मार्ट बनाने की दिशा में एक बड़ा कदम है।

तकनीकी सारांश: SAKI (स्कोर-अवेयर लो-रैंक की इंडेक्सिंग)

समस्या विवरण

लॉन्ग-कॉन्टेक्स्ट डिकोडिंग (जैसे, 1M टोकन) में, एक 8B-पैरामीटर मॉडल का KV कैश महत्वपूर्ण मेमोरी (bf16 में ~130 GB) का उपभोग करता है, जिससे रिट्रीवल (retrieval) एक महत्वपूर्ण बाधा बन जाता है। KV कैश के लिए मौजूदा लो-रैंक इंडेक्सिंग रणनीतियाँ दो श्रेणियों में आती हैं, जो दोनों ही अटेंशन मैकेनिज्म द्वारा उपयोग किए जाने वाले वास्तविक मेट्रिक को अनुकूलित करने में विफल रहती हैं:

वेट-साइड ट्रंकेशन (Weight-side truncation): प्रोजेक्शन मैट्रिसेस को ट्रंकेट करने वाली विधियाँ (जैसे, Weight-SVD) ऑपरेटर ज्योमेट्री को अनुकूलित करती हैं लेकिन डेटा डिस्ट्रीब्यूशन को अनदेखा करती हैं, जिसके परिणामस्वरूप खराब रिकॉल (रैंक $r=32$ पर 0.366) होता है।
डेटा-अवेयर रिकंस्ट्रक्शन (Data-aware reconstruction): Key-PCA जैसी विधियाँ कैश किए गए कीज़ (keys) के वेरिएंस को संरक्षित करती हैं। हालाँकि, अटेंशन स्कोर $s = q^\top k$ के रूप में कंप्यूट किए जाते हैं। एक दिशा जिसमें की वेरिएंस (key variance) अधिक है, उसे क्वेरी (query) द्वारा कभी नहीं पढ़ा जा सकता है, जबकि एक दिशा जिसमें मध्यम वेरिएंस है लेकिन जिसे क्वेरी प्रोजेक्शन ( $W_Q$ ) द्वारा प्रवर्धित (amplify) किया जाता है, वह स्कोर पर हावी हो सकती है। PCA रिकंस्ट्रक्शन एरर ( $\|k - Pk\|^2$ ) को कम करता है, न कि स्कोर डिस्टॉर्शन (score distortion) को।

मूल समस्या यह है कि दोनों में से कोई भी वर्ग अटेंशन स्कोर रैंकिंग को अनुकूलित नहीं करता है, जो KV रिट्रीवल का वास्तविक उद्देश्य है।

कार्यप्रणाली: SAKI

यह शोध पत्र रैंक- $r$ की (key) कंप्रेशन से प्रेरित अपेक्षित स्कोर डिस्टॉर्शन (expected score distortion) को व्युत्पन्न करता है और SAKI (स्कोर-अवेयर लो-रैंक की इंडेक्सिंग) का प्रस्ताव करता है, जो एक ट्रेनिंग-फ्री इंडेक्स है जो सीधे इस डिस्टॉर्शन को अनुकूलित करता है।

1. सैद्धांतिक व्युत्पत्ति (Theoretical Derivation)

लेखक रैंक- $r$ लीनियर मैप $P$ द्वारा कीज़ पर लागू किए गए अपेक्षित स्कोर डिस्टॉर्शन $L(P)$ को परिभाषित करते हैं। $\Sigma_q$ और $\Sigma_k$ के सेकंड मोमेंट्स के साथ एक इंडिपेंडेंट क्वेरी/की मॉडल और अटेंशन ऑपरेटर $A = W_Q^\top W_K$ के तहत:
$L(P) = \mathbb{E}[(q^\top A (I-P) k)^2] = \|\Sigma_q^{1/2} A (I-P) \Sigma_k^{1/2}\|_F^2$
यह फॉर्मूलेशन प्रकट करता है कि इष्टतम इंडेक्स एक टू-साइडेड कोवेरिएंस-वेटेड लो-रैंक एप्रोक्सिमेशन है। PCA के विपरीत, जो सिंगल-साइडेड वेट ( $\Sigma_k$ ) का उपयोग करता है, SAKI अपने एप्रोक्सिमेशन को मॉडल के स्कोरिंग ऑपरेटर $A$ द्वारा मध्यस्थता करते हुए, क्वेरी स्टैटिस्टिक्स ( $\Sigma_q$ ) और की स्टैटिस्टिक्स ( $\Sigma_k$ ) दोनों द्वारा वेट करता है।

2. सटीक समाधान (Exact Solution)

यह शोध पत्र सिद्ध करता है कि $L(P)$ के लिए रैंक- $r$ इष्टतम एक प्रोजेक्टर (जैसा कि PCA में है) नहीं बल्कि एक असिमेट्रिक लो-रैंक बाइलिनियर फैक्टराइजेशन है।

मान लीजिए $C = \Sigma_q^{1/2} \Sigma_k^{1/2}$ (उस हेड कोऑर्डिनेट्स में जहाँ $A$ को अवशोषित किया गया है)।
इष्टतम मैप $M_r$ को $C$ के ट्रंकेटेड SVD से प्राप्त किया जाता है: $C_r = U_r \Lambda_r V_r^\top$ ।
समाधान है $M_r = \Sigma_q^{-1/2} C_r \Sigma_k^{-1/2}$ ।
मुख्य अंतर: $M_r$ सामान्यतः इडेम्पोटेंट (idempotent) नहीं है ( $M_r^2 \neq M_r$ )। यह एक इष्टतम लीनियर स्कोर मैप है, न कि एक ओब्लिक प्रोजेक्टर। यह उन दिशाओं को कैप्चर करने की अनुमति देता है जिन्हें स्पैन-आधारित विधियाँ (PCA सहित) व्यक्त नहीं कर पाती हैं।

3. एल्गोरिदम (SAKI)

यह विधि ट्रेनिंग-फ्री है और इसके लिए केवल एक सिंगल कैलिब्रेशन पास की आवश्यकता होती है:

प्रति-हेड कोवेरिएंस $\Sigma_q$ (pre-RoPE क्वेरीज का अनसेंटर्ड सेकंड मोमेंट) और $\Sigma_k$ (सेंटर्ड कीज़) की गणना करें।
मैट्रिक्स $C = \Sigma_q^{1/2} \Sigma_k^{1/2}$ की गणना करें।
$U_r, \Lambda_r, V_r$ प्राप्त करने के लिए $C$ पर SVD करें।
कंप्रेशन मैट्रिसेस $B_k$ और $B_q$ का निर्माण करें ताकि कीज़ को कोड्स $c_j = B_k^\top (k_j - \mu)$ के रूप में स्टोर किया जा सके और स्कोर को $(B_q^\top q)^\top c_j$ के माध्यम से रिकंस्ट्रक्ट किया जा सके।
इन मैप्स को ऑफलाइन मॉडल के $W_Q$ और $W_K$ वेट्स में फोल्ड किया जा सकता है।

मुख्य योगदान

समस्या का फॉर्मूलेशन: KV रिट्रीवल के लिए सही ऑब्जेक्टिव के रूप में अपेक्षित स्कोर डिस्टॉर्शन की पहचान करना, जिससे टू-साइडेड कोवेरिएंस-वेटेड एप्रोक्सिमेशन समस्या सामने आई।
सटीक समाधान: एक क्लोज्ड-फॉर्म, असिमेट्रिक रैंक- $r$ समाधान का व्युत्पन्न, जो किसी भी स्पैन-प्रतिबंधित (प्रोजेक्टर-आधारित) विधि से स्पष्ट रूप से बेहतर प्रदर्शन करता है।
अनुभवजन्य सत्यापन (Empirical Validation): यह प्रदर्शन कि सैद्धांतिक स्वतंत्रता धारणा (जो क्लोज्ड फॉर्म को व्युत्पन्न करने के लिए उपयोग की गई थी) लगभग पूर्ण सटीकता के साथ काम करती है ( $r=0.997$ प्रेडिक्टेड और एक्चुअल स्कोर-MSE रिडक्शन के बीच पियर्सन कोरिलेशन)।
ऑपरेटर ज्योमेट्री डायग्नोसिस: यह विश्लेषण दिखाना कि अटेंशन स्कोर ऑपरेटर्स अत्यधिक नॉन-नॉर्मल (median Henrici departure 0.95–0.98) हैं और भारी-पूंछ वाले स्पेक्ट्रा (heavy-tailed spectra) रखते हैं। यह समझाता है कि वेट-ओनली ट्रंकेशन क्यों विफल होता है (लाइव दिशाओं को हटा देता है) और क्यों इनवेरिएंट-सबस्पेस विधियाँ विफल होती हैं (हब-आकार के क्रॉस-सबस्पेस कपलिंग्स को काट देती हैं)।

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

लेखक ने चार मॉडल्स (LLaMA-3.1-8B, Qwen2.5-7B, Mistral-7B-v0.1, Llama-3.2-3B) पर SAKI का मूल्यांकन किया, जिसमें 4,096-टोकन नेचुरल टेक्स्ट और सटीक क्वेरी का प्रोटोकॉल उपयोग किया गया, और टॉप-64 रिकॉल को मापा गया।

प्रदर्शन: SAKI हर रैंक ( $r=16, 32, 64$ $r = 16, 32, 64$ ) पर और हर मॉडल पर Key-PCA से बेहतर प्रदर्शन करता है।
- $r=32$ पर, SAKI PCA द्वारा छोड़े गए शेष रिकॉल एरर में से 13–30% को कम करता है।
- उदाहरण (LLaMA-3.1-8B): रिकॉल 0.748 (PCA) से सुधरकर 0.799 (SAKI) हो जाता है।
- उदाहरण (Qwen2.5-7B): रिकॉल 0.786 (PCA) से सुधरकर 0.850 (SAKI) हो जाता है।
हेड-वाइज सुधार: प्रति मॉडल 68–89% अटेंशन हेड्स में सुधार दिखता है।
लेयर-वाइज डिस्ट्रीब्यूशन: लाभ गहरे लेयर्स (जैसे, लेयर्स 16–31) में केंद्रित हैं, जहाँ क्वेरी और की ज्योमेट्री सबसे अधिक अलग होती है।
एब्लेशन (Ablation): पेपर लाभ के स्रोत को अलग करता है। एक वेरिएंट जो "वेटेड स्पैन" (SAP-map) का उपयोग करता है, वह रॉ स्पैन (PCA) की तुलना में बेहतर प्रदर्शन करता है, लेकिन सटीक असिमेट्रिक ऑप्टिमम (SAKI-opt) स्पष्ट रूप से श्रेष्ठ है। यह पुष्टि करता है कि लाभ ऑब्जेक्टिव फंक्शन और असिमेट्रिक फैक्टराइजेशन से आता है, न कि केवल कोवेरिएंस मैट्रिसेस के इंसर्शन से।

महत्व और दावे

पेपर का दावा है कि क्षेत्र की पिछली दृष्टिकोण (वेट-साइड या डेटा-रिकंस्ट्रक्शन) विफल होती हैं क्योंकि वे गलत मात्रा को अनुकूलित करती हैं। अटेंशन की (key) वेरिएंस के बजाय क्वेरी और की स्टैटिस्टिक्स के बीच की इंटरेक्शन पर निर्भरता होती है।

सैद्धांतिक वैधता: व्युत्पन्न क्लोज्ड-फॉर्म समाधान केवल एक ह्यूरिस्टिक नहीं है; यह स्कोर-डिस्टॉर्शन ऑब्जेक्टिव के लिए सटीक इष्टतम है।
व्यावहारिक प्रभाव: SAKI मौजूदा KV इंडेक्स के लिए एक ट्रेनिंग-फ्री, ड्रॉप-इन रिप्लेसमेंट प्रदान करता है जो रिट्रीवल फिडेलिटी को काफी बेहतर बनाता है, विशेष रूप से गहरी लेयर्स में जहाँ रिट्रीवल सबसे कठिन होता है।
सीमाएँ: लेखक नोट करते हैं कि हालांकि SAKI रिकॉल में सुधार करता है, लेकिन यह अभी तक $r \le 32$ पर 0.95 रिकॉल तक नहीं पहुँच पाया है। एंड-टू-एंड जनरेशन क्वालिटी और लंबे कॉन्टेक्स्ट (4K से परे) या विभिन्न डोमेन पर प्रदर्शन का पूर्ण सत्यापन किया जाना बाकी है। स्वतंत्रता धारणा, हालांकि अनुभवजन्य रूप से मान्य है, RoPE रोटेशन के संबंध में एक अनुमान है।

पेपर निष्कर्ष निकालता है कि भविष्य के इंडेक्सिंग को अटेंशन मैकेनिज्म की विशिष्ट "ऑपरेटर ज्योमेट्री" को संरक्षित करना चाहिए, विशेष रूप से नॉन-नॉर्मल, क्रॉस-सबस्पेस कपलिंग्स को जिन्हें SAKI कैप्चर करने के लिए डिज़ाइन किया गया है।