तकनीकी सारांश: कार्य संरचना अनुक्रम मॉडलों में परतवार अवस्था एन्कोडिंग को उलट देती है

समस्या विवरण

अनुक्रम मॉडलों (sequence models) के यांत्रिक व्याख्यात्मकता (mechanistic interpretability) में अक्सर परतवार अवस्था एन्कोडिंग (layerwise state encodings) को एक निश्चित वास्तुशिल्प विशेषता (architectural trait) के रूप में माना जाता है। प्रचलित धारणा यह है कि आवर्ती मॉडल (जैसे LSTM, GRU) और स्टेट-स्पेस मॉडल (SSM जैसे Mamba) देर से आने वाली परतों में पठनीय अवस्था (readable state) को केंद्रित करते हैं, जबकि अटेंशन-आधारित मॉडल (Transformers) गहराई में अवस्था संचय (state accumulation) को वितरित करते हैं। यह शोध पत्र इस विचार को चुनौती देता है कि ये एन्कोडिंग रणनीतियाँ स्थिर वास्तुशिल्प अपरिवर्तनीय (architectural invariants) हैं। यह तर्क देता है कि समान वास्तुकला अपने परतवार अवस्था-एन्कोडिंग प्रोफाइल को तब बदल देती है जब कार्य की कम्प्यूटेशनल संरचना बदल जाती है। इसके अतिरिक्त, साहित्य अक्सर इन पैटर्न के दो संभावित चालकों को मिला देता है: बीजगणितीय संरचना (विशेष रूप से क्रमविनिमेयता/commutativity) बनाम कम्प्यूटेशनल संरचना (प्रिफिक्स अपडेट बनाम स्टैक ट्रैकिंग)। लेखक इन कारकों को अलग करने और यह निर्धारित करने का लक्ष्य रखते हैं कि क्या अवस्था एन्कोडिंग केवल वास्तुकला का गुण है या वास्तुकला और कार्य के बीच एक परस्पर क्रिया (interaction) है।

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

यह अध्ययन पाँच वास्तुकला परिवारों में एक नियंत्रित प्रयोगात्मक सूट का उपयोग करता है: Transformers, Mamba, Mamba-2, LSTMs, और GRUs।

कार्य (Tasks):
लेखक तीन औपचारिक अवस्था-ट्रैकिंग कार्यों और एक अर्ध-वास्तविक बेंचमार्क का उपयोग करते हैं:

1. Parity: एक क्रमविनिमेय (commutative), सिंगल-बिट प्रिफिक्स-अपडेट कार्य।
2. $S_3$ Permutation Composition: एक गैर-क्रमविनिमेय (non-commutative), 6-क्लास प्रिफिक्स-अपडेट कार्य। यह एक महत्वपूर्ण विभेदक के रूप में कार्य करता है: यदि बीजगणितीय क्रमविनिमेयता एन्कोडिंग प्रोफाइल को संचालित करती है, तो $S_3$ का व्यवहार Dyck जैसा होना चाहिए; यदि कम्प्यूटेशनल संरचना संचालित करती है, तो $S_3$ का व्यवहार Parity जैसा होना चाहिए।
3. Dyck-$k$: एक बाउंडेड-डेप्थ स्टैक-ट्रैकिंग (ब्रैकेट मिलान) कार्य जिसमें अधिकतम गहराई 10 है।
4. Code-Depth: Python कोड (CodeSearchNet) से लिया गया एक टोकन-लेवल नेस्टिंग-डेप्थ कार्य, जो एक अर्ध-वास्तविक सत्यापन के रूप में कार्य करता है।

प्रायोगिक सेटअप:

• औपचारिक कार्य: प्रयोग 4-परत वाले मॉडलों का उपयोग करते हैं जहाँ $d_{model}=128$ है। वास्तुकलाएँ पैरामीटर-मैच की गई हैं (लगभग 830K पैरामीटर्स)। प्रशिक्षण अनुक्रम लंबाई $L=40$ पर होता है, और मूल्यांकन $L \in \{40, 60, 80, 100, 120\}$ पर किया जाता है।
• प्री-ट्रेन्ड स्केल: Mamba-130M (24 परतें) और Pythia-160M/410M (12/24 परतें) पर फाइन-ट्यूनिंग की जाती है।
• मापन (Measurements):
  • लीनियर प्रोबिंग (Linear Probing): प्रत्येक परत में अवस्था की पठनीयता मापने के लिए फ्रोजन हिडन स्टेट्स पर लॉजिस्टिक रिग्रेशन क्लासिफायर प्रशिक्षित करना।
  • कारण संबंधी हस्तक्षेप (Causal Interventions): ज़ीरो-एब्लेशन (प्रोब दिशाओं को प्रोजेक्ट करना) और सबस्पेस एब्लेशन (शीर्ष PCA या सिंगुलर दिशाओं को हटाना) का उपयोग करके कार्यात्मक आवश्यकता का परीक्षण करना।
  • OOD सामान्यीकरण: अनदेखी लंबाई ($L=80, 100, 120$) पर $L=40$ पर प्रशिक्षित एब्लेशन लागू करना।
  • घटक एट्रिब्यूशन (Component Attribution): एक मैच की गई शुद्ध-PyTorch बैकएंड के तहत Mamba के Conv1D घटक को एब्लेट करना ताकि विभिन्न कार्यों के तहत इसके योगदान को अलग किया जा सके।
  • एक्टिवेशन पैचिंग (Activation Patching): लॉगिट रिकवरी को मापने के लिए क्लीन इनपुट्स से मिड-पोजीशन एक्टिवेशन्स को करप्टेड इनपुट्स में पैच करना।

मुख्य योगदान

1. कार्य-निर्भर उलटाव (Task-Dependent Reversal): यह शोध पत्र प्रदर्शित करता है कि "केंद्रित देर से" बनाम "वितरित निर्माण" का पैटर्न एक वास्तुशिल्प अपरिवर्तनीय नहीं है। यह कार्य के आधार पर बदल जाता है:
   • प्रिफिक्स-अपडेट कार्य (Parity, $S_3$): आवर्ती मॉडल और SSMs अंतिम परत में अवस्था को केंद्रित करते हैं, जबकि Transformers इसे परतों में धीरे-धीरे बनाते हैं।
   • बाउंडेड स्टैक-ट्रैकिंग (Dyck): पैटर्न उलट जाता है। Transformers तुरंत (शुरुआती परतों में) लगभग पूर्ण पठनीयता प्राप्त करते हैं, जबकि आवर्ती मॉडल और SSMs अवस्था का प्रतिनिधित्व प्रगतिशील रूप से बनाते हैं।
2. बीजगणित और कम्प्यूटेशन को अलग करना: गैर-क्रमविनिमेय $S_3$ कार्य को पेश करके, लेखक दिखाते हैं कि एन्कोडिंग प्रोफाइल कम्प्यूटेशनल संरचना (प्रिफिक्स अपडेट बनाम स्टैक) को ट्रैक करता है न कि बीजगणितीय संरचना (क्रमविनमेयता) को। $S_3$ सभी वास्तुकलाओं में Parity के साथ समूह बनाता है, जो इस परिकल्पना को खारिज करता है कि क्रमविनिमेयता मुख्य चालक है।
3. पठनीयता-कारणता अलगाव (Readability–Causality Decoupling): अध्ययन से पता चलता है कि जिस परत में अवस्था सबसे अधिक रैखिक रूप से पठनीय (सर्वश्रेष्ठ प्रोब परत) होती है, वह हमेशा उस परत से मेल नहीं खाती जो कार्य के लिए कारण संबंधी रूप से आवश्यक है। यह अलगाव प्री-ट्रेन्ड मॉडलों में विशेष रूप से स्पष्ट है, जहाँ प्रतिनिधित्व अधिक रेडंडेंट (redundant) और वितरित होते हैं।
4. घटक-स्तरीय कार्य निर्भरता: एक ही वास्तुकला (Mamba) के भीतर भी, विशिष्ट घटकों का महत्व कार्य के अनुसार भिन्न होता है। उदाहरण के लिए, Conv1D घटक $S_3$ और Parity के लिए महत्वपूर्ण है लेकिन Dyck प्रदर्शन पर इसका कोई महत्वपूर्ण प्रभाव नहीं पड़ता है, जो किसी वास्तुकला के लिए एकल "तंत्र" (mechanism) की धारणा को चुनौती देता है।

मुख्य परिणाम

• व्यवहार स्तर: आवर्ती बेसलाइन Parity और $S_3$ पर पूर्ण रूप से सामान्यीकरण करती हैं और Dyck पर SSMs और Transformers से बेहतर प्रदर्शन करती हैं। $S_3$ प्रदर्शन क्रम (Recurrent > SSM > Transformer) Parity के समान है, न कि Dyck के।
• प्रोबिंग प्रोफाइल:
  • Parity/$S_3$: Mamba और LSTMs अंतिम परत ($L_3$ में 4-परत मॉडल में) पर प्रोब सटीकता में एक तीव्र उछाल दिखाते हैं। Transformers $L_1$ से $L_3$ तक क्रमिक वृद्धि दिखाते हैं।
  • Dyck: Transformers $L_0$ या $L_1$ पर उच्च सटीकता प्राप्त करते हैं। Mamba और LSTMs क्रमिक वृद्धि दिखाते हैं, जो बाद की परतों में ही संतृप्ति (saturation) तक पहुँचते हैं।
  • $S_3$ पुष्टि: $S_3$, Parity प्रोफाइल का अनुसरण करता है (RNNs/SSMs के लिए देर से एकाग्रता, Transformers के लिए क्रमिक), जो कम्प्यूटेशनल संरचना परिकल्पना की पुष्टि करता है।
• कारण संबंधी हस्तक्षेप (Causal Interventions):
  • 4-परत वाले औपचारिक मॉडलों में, सर्वश्रेष्ठ प्रोब दिशा का एब्लेशन महत्वपूर्ण सटीकता गिरावट का कारण बनता है, और ये दिशाएं लंबी अनुक्रम लंबाई (OOD) पर भी आवश्यक रहती हैं।
  • प्री-ट्रेन्ड मॉडलों में, एक विभाजन होता है:
    • Pythia (Dyck): एक मजबूत मिड-लेयर बॉटलनेक मौजूद है ($L_6$–$L_7$ पर एब्लेशन 160M पर सटीकता को ~81% कम कर देता है)। सर्वश्रेष्ठ प्रोब परत ($L_{11}$) सबसे अधिक कारण संबंधी रूप से महत्वपूर्ण नहीं है।
    • Mamba (Parity/Dyck/$S_3$): अंतिम परत अत्यधिक पठनीय है, लेकिन एकल-दिशा एब्लेशन नगण्य गिरावट (<0.2%) का कारण बनता है, जो इंगित करता है कि अवस्था पर्याप्त है लेकिन कई दिशाओं में वितरित है, जिसे एक एकल रैखिक प्रोब द्वारा कैप्चर नहीं किया जा सकता है।
• Conv1D एट्रिब्यूशन: Mamba से Conv1D को हटाने से $S_3$ (-12.64 pp) और Parity (-6.83 pp) पर बड़ी गिरावट आती है, लेकिन Dyck पर कोई महत्वपूर्ण गिरावट नहीं होती है (+4.04 pp, जो महत्वपूर्ण नहीं है)। यह दर्शाता है कि घटक का महत्व कार्य-निर्भर है।

महत्व और दावे

यह शोध पत्र तर्क देता है कि यह पूछना कि "वास्तुकला X अवस्था को कैसे एनकोड करती है?" कार्य की कम्प्यूटेशनल संरचना को निर्दिष्ट किए बिना अधूरा है। लेखक दावा करते हैं कि:

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

लेखक निष्कर्ष निकालते हैं कि उनके निष्कर्ष औपचारिक कार्यों, प्री-ट्रेन्ड स्केल (410M पैरामीटर्स तक), और अर्ध-वास्तविक कोड बेंचमार्क में भी लागू होते हैं, जो अनुक्रम मॉडलिंग में वास्तुकला और कार्य संरचना के बीच एक मजबूत परस्पर क्रिया का सुझाव देते हैं।