अटेंशन रेजिडुअल्स समझें: ट्रांसफॉर्मर की गहराई पर फिर से सोचना

आधुनिक LLMs और अधिक गहरे, चौड़े और कंप्यूट-भूखे हो गए हैं, फिर भी और अधिक ट्रांसफॉर्मर लेयर जोड़ना हमेशा अनुपातिक लाभ नहीं देता। एक वजह यह है कि मानक रेजिडुअल कनेक्शन लेयर आउटपुट्स को स्थिर इकाई वेट्स से समेकित करते हैं, जिससे हर लेयर अपने से पहले आई सारी जानकारी का एक समान योग विरासत में ले लेती है। इससे शुरुआती निरूपण पतले पड़ सकते हैं, हिडन-स्टेट का परिमाण बढ़ सकता है, और नेटवर्क के लिए सबसे उपयोगी मध्यवर्ती फीचर्स का चुनिंदा पुन: उपयोग करना कठिन हो जाता है।

गहराई को एक स्थिर योगात्मक पुनरावृत्ति मानने की बजाय, Kimi टीम का अटेंशन रेजिडुअल्स पेपर प्रत्येक लेयर को सीखी हुई सॉफ्टमैक्स वेट्स के जरिए पहले की लेयर्स के आउटपुट पर "अटेंड" करने देता है। इस लेख में, हम समझेंगे कि मानक रेजिडुअल संचय कैसे बॉटलनेक बन जाता है, अटेंशन रेजिडुअल्स कैसे काम करता है, ब्लॉक वेरिएंट क्यों मायने रखता है, और नतीजे वास्तव में गहरे भाषा मॉडलों के स्केलिंग के बारे में क्या सुझाते हैं।

यदि आप आधुनिक ट्रांसफॉर्मर आर्किटेक्चर के पीछे की कुछ अवधारणाओं के बारे में और सीखना चाहते हैं, तो मैं DataCamp का कोर्स Transformer Models Tutorial in PyTorch देखने की सलाह देता/देती हूँ।

समस्या: डूबता हुआ सिग्नल

रेजिडुअल कनेक्शन्स कई वर्षों से डीप लर्निंग की बुनियाद रहे हैं। ट्रांसफॉर्मर्स में, मानक रेजिडुअल अपडेट इस प्रकार है:

hl=hl−1​+fl−1​(hl−1​)

यह बहुत गहरे नेटवर्क्स में ग्रेडिएंट के प्रवाह में मदद करता है। लेकिन रेजिडुअल्स सिर्फ ग्रेडिएंट की चाल नहीं हैं। वे यह भी निर्धारित करते हैं कि गहराई में सूचना कैसे समेकित होती है। यदि हम इस पुनरावृत्ति को अनरोल करें, तो मिलता है:

hl​=h1​+∑i=1l−1​fi​(hi​)

इसका मतलब है कि लेयर l पर हिडन स्टेट सिर्फ एम्बेडिंग तथा सभी पूर्ववर्ती लेयर आउटपुट्स का समान वेट वाला योग है। यानी हर योगदान को प्रभावी रूप से एक ही वेट सौंपा जाता है।

स्केल पर यह समस्या बन जाती है। पेपर दर्शाता है कि PreNorm आर्किटेक्चर में, बिना वेट वाले रेजिडुअल संचय के कारण हिडन-स्टेट का परिमाण गहराई के साथ लगभग O(L) के रूप में बढ़ता है। 

जैसे-जैसे रेजिडुअल स्ट्रीम बढ़ती है, प्रारंभिक लेयर आउटपुट्स एक बड़े रनिंग सम के भीतर अधिकाधिक पतले पड़ते जाते हैं। एक बार जब शुरुआती लेयर का सिग्नल इस संचय में मिल जाता है, तो गहरी लेयर्स उसे चुनिंदा तौर पर वापस नहीं पा सकतीं। इसलिए वे केवल समेकित स्टेट पर ही काम करती हैं।

इसे पेपर में “ड्राउनिंग सिग्नल” प्रभाव कहा गया है। इस अक्षमता का एक ठोस अनुभवजन्य संकेत लेयर-प्रूनिंग अध्ययनों से आता है, यानी प्रशिक्षित मॉडलों की एक बड़ी हिस्सेदारी वाली लेयर्स को अक्सर प्रदर्शन पर न्यूनतम प्रभाव के साथ हटाया जा सकता है। यह संकेत देता है कि जबकि हम गहराई बढ़ाते जा रहे हैं, मॉडलों में इसे पूरी तरह उपयोग करने का प्रभावी तंत्र नहीं है। पदानुक्रमित तर्क शृंखला बनाने की बजाय, नेटवर्क एक दोहरावदार रिले की तरह व्यवहार करता है जहाँ शुरुआती सिग्नल क्रमशः पतले पड़ते जाते हैं।

फॉरवर्ड-पास का एक ट्रेड-ऑफ भी है। जैसे-जैसे रेजिडुअल स्ट्रीम बढ़ती है, बाद की लेयर्स को संचित स्टेट को अर्थपूर्ण रूप से प्रभावित करने के लिए अधिक परिमाण वाले आउटपुट उत्पन्न करने पड़ सकते हैं। पेपर इसे PreNorm व्यवहार से जोड़ता है, जहाँ मानक रेजिडुअल संचय के तहत हिडन-स्टेट का परिमाण गहराई के साथ एकरूपता से बढ़ता है।

रेजिडुअल कनेक्शन्स ग्रेडिएंट प्रवाह में मदद करते हैं, और किसी मध्यवर्ती हिडन स्टेट के सापेक्ष ग्रेडिएंट को इस प्रकार लिखा जा सकता है:

यहाँ पहचान (identity) पद प्रत्यक्ष ग्रेडिएंट पथ को संरक्षित करता है। लेकिन रेजिडुअल्स फिर भी फॉरवर्ड समेकन पथ को हर पूर्व लेयर के साथ 1.0 के स्थिर वेट से व्यवहार करने पर मजबूर करते हैं। यही संरचनात्मक सीमा है जिसे AttnRes सुधारने की कोशिश करता है।

योगात्मक रेजिडुअल्स से डेप्थ-वाइज अटेंशन तक

पेपर का सबसे दिलचस्प वैचारिक कदम है समय–गहराई द्वैतता। रेजिडुअल कनेक्शन्स गहराई में जानकारी को उसी तरह संकुचित करते हैं जैसे रिकरेंट नेटवर्क्स समय में जानकारी को संकुचित करते हैं। अनुक्रम मॉडलिंग में, अटेंशन ने पुनरावृत्ति की जगह ले ली क्योंकि हर पोजीशन पहले की पोजीशन्स तक चुनिंदा पहुँच बना सकती थी। AttnRes यही संक्रमण नेटवर्क की गहराई पर लागू करता है। अगला हिडन स्टेट पहले की लेयर्स पर स्थिर योग के रूप में परिभाषित करने के बजाय, AttnRes प्रत्येक लेयर को पिछली लेयर्स के आउटपुट्स पर अटेंड करने देता है:

जहाँ वेट्स αi→l गहराई पर सॉफ्टमैक्स अटेंशन वेट्स हैं और इनका योग 1 होता है। ये वेट्स इस प्रकार गणना किए जाते हैं:

जहाँ पेपर यह उपयोग करता है:

प्रत्येक लेयर को सीखा हुआ एक स्यूडो-क्वेरी वेक्टर wl मिलता है, और वह क्वेरी पहले की लेयर्स के आउटपुट्स से बने कीज़ और वैल्यूज़ पर अटेंड करती है। कीज़ पर RMSNorm लागू किया जाता है ताकि स्वाभाविक रूप से बड़े आउटपुट परिमाण वाली लेयर्स केवल स्केल अधिक होने के कारण सॉफ्टमैक्स पर हावी न हो जाएँ।

पेपर में मुख्य परिवर्तन यह है कि मानक रेजिडुअल्स को गहराई पर एक प्रकार के लीनियर अटेंशन की तरह माना गया है, जबकि AttnRes उसे गहराई पर सॉफ्टमैक्स अटेंशन में अपग्रेड करता है। समान संचय के बजाय, हमें गहराई में चुनिंदा पुनर्प्राप्ति मिलती है।

एक छोटा पर महत्वपूर्ण अमल-सम्बंधी विवरण है इनिशियलाइज़ेशन। लेखक सुझाव देते हैं कि सभी स्यूडो-क्वेरी वेक्टर्स को शून्य से इनिशियलाइज़ किया जाना चाहिए, जिससे शुरुआती अटेंशन वेट्स समान रहते हैं और मॉडल प्रशिक्षण एक समान-वेट औसत के रूप में शुरू होता है, न कि किसी यादृच्छिक रूप से पक्षपाती अटेंशन तंत्र के रूप में।

चित्र 1: अटेंशन रेजिडुअल्स का अवलोकन (Attention Residuals पेपर)

फुल अटेंशन रेजिडुअल्स

फुल अटेंशन रेजिडुअल्स में, प्रत्येक लेयर सभी पूर्ववर्ती लेयर आउटपुट्स पर अटेंड करती है, जिससे वह गहराई में फैली जानकारी को चुनिंदा रूप से समेकित कर सकती है। यह मॉडल को अधिकतम लचीलापन देता है, ताकि कोई गहरी लेयर अपने ठीक पहले की लेयर, मूल एम्बेडिंग, या किसी भी पूर्व लेयर पर जोर दे सके—जहाँ भी सबसे उपयोगी सिग्नल मौजूद हो।

मानक प्रशिक्षण में, फुल AttnRes की मेमोरी ओवरहेड उतनी बड़ी नहीं होती जितनी पहली नज़र में लग सकती है, क्योंकि कई लेयर आउटपुट्स पहले से ही बैकप्रोपेगेशन के लिए रखे जाते हैं। लेकिन बड़े पैमाने के प्रशिक्षण में तस्वीर बदल जाती है। एक बार जब आप एक्टिवेशन रिकम्प्यूटेशन और पाइपलाइन पैरेललिज्म लाते हैं, तो पहले के आउटपुट्स को स्पष्ट रूप से संरक्षित और संप्रेषित करना पड़ता है ताकि बाद की लेयर्स उन पर अटेंड कर सकें—और यह महंगा हो जाता है।

इसलिए फुल AttnRes मूल अवधारणा को समझने का सर्वोत्तम तरीका है, लेकिन यह वह संस्करण नहीं है जिसे अधिकांश टीमें स्केल पर तैनात करना चाहेंगी।

ब्लॉक अटेंशन रेजिडुअल्स

पेपर ने मेमोरी और संचार ओवरहेड को कम करने के लिए ब्लॉक अटेंशन रेजिडुअल्स का विचार प्रस्तुत किया। मॉडल की लेयर्स को ब्लॉक्स में विभाजित किया जाता है ताकि प्रत्येक ब्लॉक के भीतर आउटपुट्स को मानक योगात्मक रेजिडुअल संचय से जोड़ा जाए, लेकिन ब्लॉक्स के पार, मॉडल हर व्यक्तिगत पूर्व लेयर के बजाय ब्लॉक-स्तरीय सारांशों पर अटेंशन करता है।

यदि Bn ब्लॉक n में लेयर्स का सेट है, तो ब्लॉक निरूपण इस प्रकार दिया जाता है:

इसके बाद मॉडल एम्बेडिंग b0=h1 पूर्व ब्लॉक सारांशों और गणना के प्रगति करते समय वर्तमान ब्लॉक के आंशिक योग पर अटेंड करता है। इससे मेमोरी और संचार O(Ld) से घटकर O(Nd) हो जाता है, जहाँ N ब्लॉक्स की संख्या है।

पेपर बताता है कि लगभग आठ ब्लॉक्स का उपयोग फुल संस्करण के अधिकांश लाभ वापस ले आता है, और स्केल बढ़ने पर फुल AttnRes और ब्लॉक AttnRes के बीच का प्रदर्शन अंतर घटता जाता है। S=2,4,8 जैसे ब्लॉक आकार फुल संस्करण के क़रीब बने रहते हैं, जबकि बहुत मोटे समूह फिर से बेसलाइन व्यवहार की ओर लौटते हैं।

पेपर से सबसे उपयोगी निष्कर्ष यह है कि लाभ का अधिकांश हिस्सा हासिल करने के लिए आपको हर लेयर पर फुल डेप्थ-वाइज अटेंशन की आवश्यकता नहीं है। जहाँ ब्लॉक AttnRes डेप्थ-वाइज अटेंशन को संगणकीय रूप से संभव बनाता है, वहीं इसे स्केल पर कुशलतापूर्वक तैनात करने के लिए अभी भी सावधानीपूर्वक सिस्टम डिज़ाइन की जरूरत होती है। पेपर ने कई ऐसे अनुकूलन भी प्रस्तुत किए हैं जो AttnRes को वास्तविक दुनिया के प्रशिक्षण और इन्फरेंस के लिए व्यावहारिक बनाते हैं।

चित्र 2: 4 फिजिकल रैंक्स और प्रति रैंक 2 वर्चुअल स्टेज के साथ कैश-आधारित पाइपलाइन संचार उदाहरण, जहाँ हैच्ड बॉक्स AttnRes ब्लॉक्स के अंत को दर्शाते हैं (Attention Residuals पेपर)

यहाँ तीन प्रमुख सिस्टम ऑप्टिमाइज़ेशन्स हैं जो AttnRes को मानक रेजिडुअल्स के लिए ड्रॉप-इन रिप्लेसमेंट बनाते हैं:

• टू-फेज कंप्यूटेशन: स्यूडो-क्वेरी को हिडन स्टेट से अलग करके, हम किसी ब्लॉक के भीतर सभी लेयर्स की क्वेरीज़ को बैच कर सकते हैं। फेज 1 में इंटर-ब्लॉक अटेंशन समानांतर में कंप्यूट होता है, जबकि फेज 2 में ऑनलाइन सॉफ्टमैक्स मर्ज का उपयोग करते हुए क्रमिक इंट्रा-ब्लॉक निर्भरताएँ संभाली जाती हैं। इससे mHC जैसे पूर्व मल्टी-स्ट्रीम तरीकों की तुलना में प्रति-लेयर I/O में उल्लेखनीय कमी आती है।
• क्रॉस-स्टेज कैशिंग: पाइपलाइन पैरेललिज्म में, पहले से प्राप्त ब्लॉक्स को लोकल कैश करके अनावश्यक डेटा ट्रांसफर्स को समाप्त किया जाता है। यह पीक प्रति-ट्रांज़िशन लागत में V× सुधार देता है (जहाँ V वर्चुअल स्टेज की संख्या है, जैसा कि चित्र 2 में दिखाया गया है), जिससे संचार को पूरी तरह गणना के साथ ओवरलैप किया जा सकता है।
• मेमोरी-एफिशिएंट प्रीफिलिंग: ब्लॉक निरूपणों को टेन्सर-पैरेलल डिवाइसेज़ में शार्ड करके, लंबी संदर्भ (128K+) अनुक्रमों के लिए मेमोरी फ़ुटप्रिंट ऑर्डर्स ऑफ़ मैग्निट्यूड तक घट जाता है (उदाहरण के लिए, प्रति डिवाइस 15GB से 0.3GB तक)।

परिणाम और प्रदर्शन विश्लेषण

पेपर ने AttnRes को कई स्तरों पर सत्यापित किया है, जैसे स्केलिंग लॉज़, एब्लेशन्स, प्रशिक्षण डायनेमिक्स, और डाउनस्ट्रीम बेंचमार्क्स।

चित्र 3: अटेंशन रेजिडुअल्स के लिए स्केलिंग लॉ कर्व्स। (Attention Residuals पेपर)

हालाँकि, शीर्ष-स्तरीय परिणाम स्केलिंग-लॉ प्रयोगों से आता है। पाँच मॉडल आकारों में, फुल AttnRes और ब्लॉक AttnRes दोनों ही लगातार बेसलाइन की तुलना में कम वैलिडेशन लॉस प्राप्त करते हैं। फिटेड कर्व्स के आधार पर, ब्लॉक AttnRes वही लॉस हासिल करता है जो बेसलाइन लगभग 1.25× अधिक कंप्यूट के साथ प्रशिक्षण देकर पाता है। यह संकेत देता है कि AttnRes सिर्फ सिद्धांत में बेहतर नहीं है, बल्कि कंप्यूट के लिहाज़ से भी अधिक दक्ष है।

डाउनस्ट्रीम कार्यों पर, AttnRes सभी मूल्यांकित बेंचमार्क्स में बेसलाइन से बेहतर है। कुछ सबसे बड़े लाभ तर्क-प्रधान कार्यों पर हैं, जैसा कि तालिका 1 में दिखाया गया है:

तालिका 1: समान प्री-ट्रेनिंग रेसिपी के बाद AttnRes और बेसलाइन का प्रदर्शन तुलना (Attention Residuals पेपर)

मल्टी-स्टेप रीजनिंग, गणित और कोड पर लाभ विशेष रूप से मजबूत हैं, जो पेपर की मूल परिकल्पना के अनुकूल है कि यदि बाद की लेयर्स धुँधले समेकन को विरासत में लेने के बजाय पहले की निरूपणों को चुनिंदा रूप से प्राप्त कर सकें, तो संयोजी तर्क बेहतर होना चाहिए।

बेसलाइन की तुलना में, AttnRes पूरे प्रशिक्षण के दौरान कम वैलिडेशन लॉस दिखाता है, गहराई में अधिक सीमाबद्ध आउटपुट परिमाण, और लेयर्स में अधिक समान ग्रेडिएंट परिमाण। बेसलाइन में सामान्य PreNorm डाइल्यूशन पैटर्न दिखाई देता है, जहाँ हिडन-स्टेट का परिमाण गहराई के साथ एकरूपता से बढ़ता है, जबकि ब्लॉक AttnRes ब्लॉक सीमाओं पर चुनिंदा समेकन की बदौलत अधिक नियंत्रित, सीमाबद्ध पैटर्न उत्पन्न करता है।

आर्किटेक्चर-स्तर पर एक व्यापक निष्कर्ष यह भी है कि एक फिक्स्ड-कंप्यूट, फिक्स्ड-पैरामीटर स्वीप में, AttnRes के लिए इष्टतम कॉन्फ़िगरेशन बेसलाइन की तुलना में अधिक गहरे, संकरे मॉडल की ओर शिफ्ट होता है। यह सुझाव देता है कि मानक रेजिडुअल संचय की तुलना में AttnRes अतिरिक्त गहराई को अधिक उपयोगी बना सकता है।

निष्कर्ष

रेजिडुअल्स को अक्सर एक जरूरी प्रशिक्षण चाल के रूप में देखा जाता है, यानी एक आइडेंटिटी शॉर्टकट जो ग्रेडिएंट्स को जीवित रखता है। AttnRes सुझाव देता है कि यह दृष्टिकोण बहुत संकीर्ण है। रेजिडुअल पाथ वे माध्यम भी हैं जिनसे सूचना गहराई के पार रूट होती है, और स्थिर योगात्मक संचय बहुत गहरे मॉडलों के लिए शायद अत्यधिक आदिम है।

अटेंशन रेजिडुअल्स किसी ज्ञात समस्या पर पट्टी बाँधने से कम और एक उपेक्षित डिज़ाइन विकल्प को अपग्रेड करने के बारे में अधिक हैं। अनुक्रम मॉडलिंग पुनरावृत्ति से अटेंशन पर इसलिए आई क्योंकि स्थिर पुनरावृत्ति बहुत सीमित थी। AttnRes का तर्क है कि अब गहराई भी उसी संक्रमण के लिए तैयार हो सकती है।

मॉडल फिर भी मजबूत स्थानीयता दिखाता है, जहाँ लेयर्स अक्सर निकटवर्ती पूर्ववर्तियों पर सबसे अधिक अटेंड करती हैं, लेकिन यह गैर-तुच्छ स्किप पैटर्न भी सीखता है, एम्बेडिंग पर लगातार वेट बनाए रखता है, और प्री-अटेंशन तथा प्री-MLP लेयर्स के बीच भिन्न व्यवहार बनाए रखता है। 

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