तकनीकी सारांश: HeadQ: KV-कैश क्वांटाइजेशन के लिए मॉडल-दृश्य विरूपण (Model-Visible Distortion) और स्कोर-स्पेस सुधार (Score-Space Correction)

1. समस्या विवरण

वर्तमान KV-कैश क्वांटाइजेशन विधियाँ मुख्य रूप से स्टोरेज-स्पेस पुनर्निर्माण (storage-space reconstruction) के लिए अनुकूलित होती हैं, जो आमतौर पर मूल कीज़ (keys) और वैल्यूज़ (values) के बीच मीन स्क्वेर्ड एरर (MSE) को कम करने का प्रयास करती हैं। हालाँकि, लेखक तर्क देते हैं कि यह उद्देश्य इस बात के साथ मेल नहीं खाता कि अटेंशन मैकेनिज्म वास्तव में कैश डेटा का उपयोग कैसे करता है।

• कीज़ (Keys) को लॉगिट्स (logits) (क्वेरी-की स्कोर) के माध्यम से पढ़ा जाता है, जिसका अर्थ है कि कीज़ के केवल विशिष्ट लीनियर प्रोजेक्शन ही अटेंशन डिस्ट्रीब्यूशन को प्रभावित करते हैं।
• वैल्यूज़ (Values) को अटेंशन-वेटेड रीडआउट ($\sum p_j v_j$) के माध्यम से पढ़ा जाता है।

पेपर का तर्क है कि निरंतर कैश एरर को कच्चे यूक्लिडियन स्टोरेज कोऑर्डिनेट्स के बजाय मॉडल-दृश्य (model-visible) कोऑर्डिनेट्स में मापा जाना चाहिए। एक गड़बड़ी (perturbation) का बड़ा स्टोरेज MSE हो सकता है लेकिन अटेंशन पर शून्य प्रभाव पड़ सकता है (यदि वह "सॉफ्टमैक्स-नल" दिशा में है), या छोटा स्टोरेज MSE हो सकता है लेकिन विनाशकारी अटेंशन क्षति (यदि वह "स्कोर-दृश्य" दिशा को प्रभावित करती है)। यह बेमेल (mismatch) इष्टतम क्वांटाइजेशन रणनीतियों को रोकता है, जिससे अप्रासंगिक घटकों पर बिट्स बर्बाद होते हैं जबकि व्यवहारिक रूप से महत्वपूर्ण घटकों को सुरक्षित रखने में विफलता होती है।

2. सैद्धांतिक आधार: मॉडल-दृश्य ज्यामिति (Model-Visible Geometry)

यह पेपर स्टोरेज एरर और व्यवहारिक एरर के बीच अंतर करने वाला एक ज्यामितीय ढांचा स्थापित करता है:

2.1 की स्कोर स्पेस (फिशर ज्योमेट्री)

एक निश्चित क्वेरी $q$ के लिए, एक की गड़बड़ी $\delta k_j$ एक स्कोर गड़बड़ी $\delta s_j = q^\top \delta k_j / \sqrt{d}$ उत्पन्न करती है।

• सॉफ्टमैक्स-नल दिशाएं (Softmax-Null Directions): एक टोकन-साझा ट्रांसलेशन (सभी कीज़ में एक स्थिर वेक्टर $\mu$ जोड़ना) सभी स्कोर में एक स्थिर मान जोड़ता है। चूंकि सॉफ्टमैक्स स्थिर बदलावों के प्रति अपरिवर्तनीय (invariant) है, इसलिए यह गड़बड़ी सॉफ्टमैक्स-नल है और शून्य अटेंशन क्षति का कारण बनती है, बावजूद इसके कि इसका स्टोरेज MSE बड़ा हो सकता है।
• स्कोर-दृश्य स्पेस (Score-Visible Space): अटेंशन केवल क्वेरी के सापेक्ष की गड़बड़ी के केंद्रित सक्रिय गुणांकों (centered active coefficients) पर निर्भर करता है।
• फिशर इंफॉर्मेशन (Fisher Information): सॉफ्टमैक्स डिस्ट्रीब्यूशन की स्थानीय द्वितीय-क्रम (second-order) ज्यामिति को फिशर इंफॉर्मेशन मैट्रिक्स द्वारा कैप्चर किया जाता है। पेपर व्युत्पन्न करता है कि मूल और गड़बड़ी वाली अटेंशन डिस्ट्रीब्यूशन के बीच KL डाइवर्जेंस स्थानीय रूप से रॉ की MSE के बजाय, एक फिशर-वेटेड स्कोर एरर के वेरिएंस (Eq. 6) के समानुपाती है।

2.2 वैल्यू रीडआउट ज्यामिति (Value Readout Geometry)

वैल्यूज़ को एक वेटेड सम (weighted sum) के माध्यम से उपभोग किया जाता है। निश्चित अटेंशन वेट्स $p$ के लिए, रीडआउट एरर $\Delta o = \sum p_j \Delta v_j$ है।

• पेपर एक $A^2$-वेटेड टोकन डिस्टॉर्शन का प्रस्ताव करता है जो वैल्यू एरर के लिए उपयुक्त सरोगेट है। रॉ वैल्यू MSE के विपरीत, यह मीट्रिक प्रत्येक टोकन के डिस्टॉर्शन को उसके अटेंशन प्रोबेबिलिटी ($p_j^2$) के वर्ग द्वारा वेट करता है, जो यह दर्शाता है कि अत्यधिक अटेंड किए गए टोकनों में त्रुटियां आउटपुट एरर में अधिक योगदान देती हैं।

3. कार्यप्रणाली: HeadQ

HeadQ एक की-साइड क्वांटाइजेशन विधि है जिसे कच्चे की वेक्टर के बजाय "स्कोर-दृश्य" अवशेष (residual) को संरक्षित करने के लिए डिज़ाइन किया गया है।

एल्गोरिदम अवलोकन

1. कैलिब्रेशन (Calibration): प्रत्येक लेयर और हेड के लिए, कैलिब्रेशन क्वेरीज (post-RoPE) से PCA के माध्यम से एक लो-रैंक बेसिस $U_h$ (रैंक $r$) सीखा जाता है।
2. कैश राइट (Cache Write):
   • एक मानक लो-बिट क्वांटाइज़र $B_b$ एक बेस की $\hat{k}^{base}_j$ जेनरेट करता है।
   • अवशेष $e_j = k_j - \hat{k}^{base}_j$ की गणना की जाती है।
   • अवशेष को सीखे गए बेसिस पर प्रोजेक्ट किया जाता है: $z_j = Q_q(U_h^\top e_j)$, जहाँ $Q_q$ एक लो-बिट साइड क्वांटाइज़र है।
   • कैश बेस की और साइड कोड $z_j$ को स्टोर करता है।
3. कैश रीड (Inference):
   • वर्तमान क्वेरी $q$ का उपयोग साइड कोड को डिक्वेंटाइज करने के लिए किया जाता है।
   • एक एडिटिव लोगिट करेक्शन (additive logit correction) लागू किया जाता है:
     $$\hat{s}_j = \frac{q^\top \hat{k}^{base}_j}{\sqrt{d}} + \frac{q^\top U_h \hat{z}_j}{\sqrt{d}}$$
   • यह करेक्शन सीधे स्कोर एरर को लक्षित करता है, पूर्ण की वेक्टर को पुनर्गठित करने की आवश्यकता को दरकिनार करता है।

मुख्य अंतर: HeadQ वर्तमान क्वेरी को स्टोर नहीं करता है; यह एक कैलिब्रेशन-लर्नड बेसिस में एक रेसिडुअल कोड स्टोर करता है जिसे रनटाइम क्वेरी एक लोगिट करेक्शन में डिकोड कर सकती है।

4. मुख्य योगदान

1. मॉडल-दृश्य विरूपण का फॉर्मूलेशन: पेपर की डिस्टॉर्शन को "स्कोर-दृश्य एरर मोड लो कॉन्स्टेंट शिफ्ट्स" के रूप में और वैल्यू डिस्टॉर्शन को "$A^2$ डायगोनल सरोगेट के साथ रीडआउट-दृश्य एरर" के रूप में परिभाषित करता है। यह समझाता है कि क्यों रॉ MSE एक खराब प्रॉक्सी है।
2. HeadQ एल्गोरिदम: एक व्यावहारिक कैश-राइट/रीड एल्गोरिदम जो एक क्वेरी बेसिस में लो-रैंक रेसिडुअल साइड कोड को स्टोर करता है और इसे एडिटिव लोगिट करेक्शन के रूप में लागू करता है।
3. मिथ्या-सिद्ध साक्ष्य श्रृंखला (Falsifiable Evidence Chain):
   • मीट्रिक श्रेष्ठता: फिशर/स्कोर-स्पेस एरर, छह मॉडलों में रॉ की MSE की तुलना में अटेंशन KL को काफी बेहतर तरीके से प्रेडिक्ट करता है (Spearman $\rho \approx 0.62\text{--}0.69$ गेन)।
   • काउंटर-एग्जेंपल्स: युग्मित प्रयोग दिखाते हैं कि कम स्टोरेज MSE वाला क्वांटाइज़र उच्च अटेंशन KL और आउटपुट एरर दे सकता है, जो सिद्ध करता है कि स्टोरेज MSE एक पर्याप्त प्रॉक्सी नहीं है।
   • नल-स्पेस हस्तक्षेप: एक "सॉफ्टमैक्स-नल" मीन कंपोनेंट (जिसे डेंस अटेंशन अनदेखा करता है) को बहाल करने से नॉन-इक्विवैरिएंट फिक्स्ड-ग्रिड क्वांटाइज़र्स के लिए फिशर डैमेज तेजी से बढ़ जाता है, जो दर्शाता है कि स्टोरेज-कोऑर्डिनेट वेस्ट व्यवहारिक प्रदर्शन को तोड़ देता है।
   • डिप्लॉयबिलिटी: कैलिब्रेशन डेटा से सीखे गए क्वेरी-PCA बेसिस सक्रिय अवशेष (active residual) को सफलतापूर्वक अनुमानित करते हैं, जबकि समान बजट वाले रैंडम बेसिस विफल रहते हैं।
   • डायरेक्शनैलिटी: HeadQ करेक्शन के साइन को उलटने ("wrong-sign" कंट्रोल) से विनाशकारी विफलता होती है, जो पुष्टि करती है कि यह विधि सही दिशा में विशिष्ट स्कोर ऑब्जेक्ट को ठीक करके काम करती है।

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

प्रयोग WikiText-103 पर छह मॉडलों (GPT-2, Pythia, Qwen2, Qwen2.5, TinyLlama, Mistral) पर किए गए।

• K-Only प्रयोग (Dense Values):
  • सबसे अधिक क्षतिग्रस्त 2-बिट रोज़ पर, HeadQ (रैंक 4) Qwen2, Pythia, और GPT-2 जैसे मॉडलों के लिए KIVI-जैसे बेसलाइन की तुलना में अतिरिक्त परप्लेक्सिटी (PPL) को लगभग 84–94% तक हटा देता है।
  • TinyLella महत्वपूर्ण सुधार दिखाता है; Mistral इस प्रोटोकॉल में पहले से ही डेंस परफॉरमेंस के करीब है।
• Full-KV संरचना (2-bit):
  • HeadQ (की-साइड) को $A^2$-वेटेड वैल्यू पॉलिसी के साथ जोड़ने से सभी छह मॉडलों में 2-बिट प्रिसिजन पर प्रदर्शन में सुधार होता है।
• विसंगति विश्लेषण (Anomaly Analysis):
  • Pythia-160M ने विशिष्ट लो-एन्ट्रॉपी "रूट-फ्लिप" रिजीम (जहाँ छोटे व्यवधान टॉप-1 टोकन को बदल देते हैं) में फिशर और KL के बीच कमजोर सहसंबंध दिखाया। यह विसंगति एक विशिष्ट फैमिली-स्पेसिफिक बाउंड्री तक सीमित थी और बड़े Pythia चेकपॉइंट्स (410M, 1B) में कमजोर हो गई।

6. महत्व और दावे

पेपर का दावा है कि ऑपरेटर मैचिंग (operator matching) KV-कैश क्वांटाइजेशन के लिए महत्वपूर्ण डिज़ाइन सबक है:

• कीज़ स्कोर ऑब्जेक्ट्स हैं; संपीड़न को क्वेरी-दृश्य अवशेषों को संरक्षित करना चाहिए और सॉफ्टमैक्स-नल शिफ्ट्स पर संसाधन बर्बाद करने से बचना चाहिए।
• वैल्यूज़ रीडआउट ऑब्जेक्ट्स हैं; संपीड़न को $A^2$-वेटेड एरर को प्राथमिकता देनी चाहिए।

लेखक स्पष्ट रूप से कहते हैं कि यह कार्य मैकेनिस्टिक और डायग्नोस्टिक है, न कि एक प्रोडक्शन-रेडी सर्विंग सिस्टम का दावा। वे यह दावा नहीं करते कि उन्होंने पैक्ड कर्नेल्स, लेटेंसी या थ्रूपुट को ऑप्टिमाइज़ किया है। इसका महत्व यहाँ है:

1. KV-कैश क्वांटाइजेशन के सैद्धांतिक उद्देश्य को स्टोरेज पुनर्निर्माण से बदलकर मॉडल-दृश्य विरूपण में सुधारना।
2. एक व्यावहारिक विधि (HeadQ) प्रदान करना जो सही एरर ज्योमेट्री को लक्षित करके लो-बिट रिजीम में पर्याप्त PPL रिकवरी प्राप्त करती है।
3. यह प्रदर्शित करना कि मानक स्टोरेज मेट्रिक्स (MSE) व्यवहारिक रूप से भ्रामक हो सकते हैं, जिसके लिए स्कोर-स्पेस प्रॉक्सी की आवश्यकता है।

पेपर निष्कर्ष निकालता है कि अगला कदम एक पूर्ण-KV कोडेक में इस ऑपरेटर-मैच्ड दृष्टिकोण को विस्तारित करना है, न कि स्टोरेज MSE को डिफॉल्ट लक्ष्य के रूप में मानना है।