एक नया “Attention Sink Detector” उन टोकन का पता लगाता है जो लार्ज-लैंग्वेज-मॉडल (LLM) इन्फरेंस के प्रीफिल स्टेज (prefill stage) के दौरान अधिकांश अटेंशन मास (attention mass) को इकट्ठा करते हैं। यह टूल दिखाता है कि की-वैल्यू (KV) कैश से इन सिंक टोकन को हटाने से मॉडल का प्रदर्शन गंभीर रूप से प्रभावित होता है, जिससे वे स्थिर जनरेशन के लिए आवश्यक एंकर बन जाते हैं।
प्रीफिल स्टेज क्यों महत्वपूर्ण है
अधिकांश LLM शोध टोकन-दर-टोकन जनरेशन लूप पर केंद्रित होते हैं, लेकिन पहले टोकन से पहले होने वाला काम—प्रीफिल—आगे होने वाली हर चीज़ की दिशा तय करता है। प्रीफिल के दौरान मॉडल पूरे प्रॉम्प्ट को प्रोसेस करता है, KV कैश बनाता है, और हर पोजीशन पर अटेंशन फैलाता है। चूंकि अटेंशन वेट्स (attention weights) उत्पन्न करने वाले सॉफ्टमैक्स (softmax) का योग 1 होना चाहिए, इसलिए मॉडल यह नहीं कह सकता कि "कुछ भी प्रासंगिक नहीं है।" यह बचे हुए प्रोबेबिलिटी मास (probability mass) को एक ऐसे टोकन पर डाल देता है जिसे एड्रेस करना आसान हो, जो आमतौर पर सीक्वेंस का पहला टोकन होता है। वह टोकन एक attention sink बन जाता है।
सरल शब्दों में अटेंशन सिंक क्या है
एक ट्रांसफॉर्मर (Transformer) में, प्रत्येक हेड हर टोकन पेयर के लिए एक वेट (weight) की गणना करता है। जब वितरण (distribution) बहुत अधिक झुक जाता है, तो एक एकल टोकन को अटेंशन वेट का असमान हिस्सा मिल सकता है। यह घटना कोई बग नहीं है; यह सॉफ्टमैक्स बाधा (softmax constraint) का परिणाम है। पहला टोकन (अक्सर वाक्य की शुरुआत का मार्कर) उस अवशिष्ट प्रोबेबिलिटी (residual probability) के लिए एक "रिलीज़ वाल्व" के रूप में कार्य करता है जिसे कहीं और असाइन नहीं किया जा सकता है।
KV-cache प्रबंधन के लिए जोखिम
KV कैश हर उस टोकन के लिए की (key) और वैल्यू (value) वेक्टर्स को स्टोर करते हैं जिसे प्रोसेस किया गया है, जिससे जनरेशन के दौरान तेज़ लुक-अप संभव हो पाता है। लंबे कॉन्टेक्स्ट वाले परिदृश्यों में, विशेषज्ञ GPU मेमोरी की सीमाओं के भीतर रहने के लिए कैश को प्रून (prune) करते हैं। नया डिटेक्टर दिखाता है कि बिना सोचे-समझे उन टोकन को हटाना जो "महत्वहीन" लगते हैं, वास्तव में उन्हीं सिंक को भी हटा देता है जिन पर मॉडल निर्भर करता है, जिससे प्रदर्शन अचानक गिर जाता है। कैश में सिंक टोकन को पिन करने से मॉडल का आंतरिक संतुलन बना रहता है और जनरेशन स्थिर रहती है।
डिटेक्टर कैसे काम करता है
- Eager attention: यह कार्यान्वयन FlashAttention जैसे फास्ट कर्नेल को बायपास करता है, जो पूर्ण अटेंशन मैट्रिक्स को सिकोड़ देते हैं, और इसके बजाय प्रत्येक हेड के लिए रॉ अटेंशन स्कोर (raw attention scores) रिकॉर्ड करता है।
- Layer-wide aggregation: प्रत्येक टोकन के लिए एक एकल अटेंशन-मास प्रोफाइल बनाने के लिए सभी लेयर्स और हेड्स में स्कोर का औसत निकाला जाता है।
- Log-median absolute deviation (MAD): चूंकि अटेंशन वेट्स अत्यधिक राइट-स्केव्ड (right-skewed) होते हैं, इसलिए एक साधारण मीन-स्टैंडर्ड-डेविएशन टेस्ट सामान्य वेरिएंस को आउटलायर्स के रूप में गलत वर्गीकृत कर देता है। वेट्स को लॉग-ट्रांसफॉर्म करने से वितरण सामान्य हो जाता है; इसके बाद MAD लागू करने से उन टोकन को फ्लैग किया जाता है जिनका अटेंशन मास सामान्य प्रसार से अधिक होता है।
उजागर किए गए सिंक के दो वर्ग
- True sinks – वाक्य की शुरुआत (BOS) वाला टोकन प्रॉम्प्ट की सामग्री की परवाह किए बिना लगातार अटेंशन का एक बड़ा हिस्सा सोख लेता है।
- Structural sinks – वे टोकन जो सिस्टम-लेवल प्रॉम्प्ट से संबंधित होते हैं, जैसे ChatML में उपयोग किए जाने वाले मार्कर (
<im_start>,<im_end>), छोटे क्लस्टर बनाते हैं जो अटेंशन को आकर्षित करते हैं। वे तार्किक सीमाओं (logical boundaries) के रूप में कार्य करते हैं, जिससे मॉडल को यूजर मैसेज को सिस्टम निर्देशों से अलग करने में मदद मिलती है।
काउंटर-पॉइंट: क्या हमें वास्तव में रॉ मैट्रिक्स की आवश्यकता है?
लेखक बताते हैं कि कच्चे आंकड़ों (raw numbers) के बिना सिंक घटना छिपी रहती है, और अधूरे डेटा पर आधारित कोई भी कैश-प्रूनिंग पॉलिसी मॉडल को अस्थिर करने का जोखिम उठाती है।
आगे क्या देखें
यह डिटेक्टर चार भागों वाली श्रृंखला का दूसरा हिस्सा है जो जनरेशन क्वालिटी के लिए एंट्रॉपी ट्रैकर के साथ शुरू हुई थी। आगामी भाग सिंक डिटेक्शन को entropy-guided speculative decoding के साथ एकीकृत करेगा। अंतिम अध्ययन एक अनुभवजन्य (empirical) अध्ययन होगा।
निष्कर्ष: अटेंशन सिंक कोई अस्पष्ट विचित्रता नहीं हैं; वे संरचनात्मक स्तंभ हैं जो प्रीफिल के दौरान ट्रांसफॉर्मर के अटेंशन वितरण को व्यवस्थित रखते हैं। कोई भी KV-cache प्रूनिंग रणनीति जो उन्हें अनदेखा करती है, वह मॉडल की स्थिरता को खतरे में डाल देगी। इन टोकन का पता लगाना और उन्हें सुरक्षित रखना एक कम लागत वाला सुरक्षा उपाय है जो लंबे कॉन्टेक्स्ट इन्फरेंस को विश्वसनीय और कुशल दोनों बना सकता है।