तकनीकी सारांश: कॉज़ल सर्किट डिस्कवरी के लिए मल्टी-ग्रैनुलर नोड प्रूनिंग

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

कॉज़ल सर्किट डिस्कवरी का लक्ष्य उन न्यूनतम उप-नेटवर्कों (subnetworks) की पहचान करना है जो बड़े भाषा मॉडलों (LLMs) के भीतर विशिष्ट व्यवहारों को कॉज़ली (causally) संचालित करते हैं। जबकि मौजूदा विधियों ने मैकेनिस्टिक इंटरप्रिटेबिलिटी (mechanistic interpretability) में प्रगति की है, वे दो प्राथमिक सीमाओं का सामना करती हैं:

1. कंप्यूटेशनल स्केलेबिलिटी (Computational Scalability): वर्तमान दृष्टिकोण मुख्य रूप से एज प्रूनिंग (घटकों के बीच कनेक्शन हटाना) पर निर्भर करते हैं। चूंकि ट्रांसफॉर्मर में नोड्स की संख्या के साथ किनारों (edges) की संख्या द्विघाती रूप से बढ़ती है ($O(|N|^2)$), इसलिए एज-स्तरीय प्रूनिंग मॉडल के आकार बढ़ने के साथ गणनात्मक रूप से अव्यवहार्य हो जाती है। इसके अलावा, इन विधियों के लिए विसंयोजित (disentangled) रेसिडुअल-स्ट्रीम रिप्रेजेंटेशन और इंटरमीडिएट एक्टिवेशन को स्टोर करने की आवश्यकता होती है, जिससे अत्यधिक मेमोरी उपयोग होता है।
2. ग्रैनुलैरिटी सीमाएँ (Granularity Limitations): मौजूदा विधियाँ आमतौर पर एक मोटे स्तर (coarse granularity) पर कार्य करती हैं, जहाँ पूरे अटेंशन हेड्स या MLP ब्लॉक्स को परमाणु इकाइयों (atomic units) के रूप में माना जाता है। यह इस संभावना की अनदेखी करता है कि इन घटकों के भीतर विशिष्ट न्यूरॉन्स (नोड्स) लक्षित व्यवहारों के लिए जिम्मेदार हो सकते हैं, जिसके परिणामस्वरूप प्राप्त सर्किट सूक्ष्म संरचनाओं को खो देते हैं और अनावश्यक घटकों को बनाए रखते हैं।

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

लेखक एक मल्टी-ग्रैनुलर नोड-प्रूनिंग फ्रेमवर्क प्रस्तावित करते हैं जो एक ही एकीकृत अनुकूलन चरण (optimization step) के भीतर विभिन्न स्तरों की ग्रैनुलैरिटी पर न्यूनतम सर्किट की खोज करता है।

2.1 पदानुक्रमित ग्रैनुलैरिटी डिज़ाइन (Hierarchical Granularity Design)

मोटे घटकों के बीच किनारों (edges) को प्रून करने के बजाय, यह फ्रेमवर्क एक ही समय में चार अलग-अलग स्तरों पर नोड्स पर मास्क सीखता है:

• ब्लॉक स्तर (Block Level): संपूर्ण अटेंशन ब्लॉक्स और MLP ब्लॉक्स।
• हेड स्तर (Head Level): व्यक्तिगत अटेंशन हेड्स।
• न्यूरॉन स्तर (Neuron Level): अटेंशन लेयर्स और MLP ब्लॉक्स के दो लीनियर लेयर्स (MLP1/Keys और MLP2/Values) के भीतर व्यक्तिगत न्यूरॉन्स।

2.2 टू-स्ट्रीम कॉज़ल प्रूनिंग सेटअप

यह पहचानना कि बिना डिस्ट्रीब्यूशन शिफ्ट (जो सरल ज़ीरो-एब्लेशन के साथ एक जोखिम है) पेश किए, कॉज़ल महत्व क्या है, इसके लिए विधि एक टू-स्ट्रीम फॉरवर्ड पास का उपयोग करती है:

• क्लीन स्ट्रीम (Clean Stream): मॉडल एक ऐसे इनपुट को प्रोसेस करता है जो एक विशिष्ट व्यवहार को प्रेरित करने के लिए डिज़ाइन किया गया है (जैसे, एक वर्ष का अनुमान लगाना > 50)।
• करप्टेड स्ट्रीम (Corrupted Stream): मॉडल इनपुट के एक न्यूनतम परिवर्तित संस्करण को प्रोसेस करता है (जैसे, जेंडर प्रोनोन्स को उलटना या शुरुआती वर्ष को बदलना) ताकि एक अलग भविष्यवाणी की जा सके।
• इंटरपोलेशन (Interpolation): प्रत्येक नोड $n_i$ के लिए, एक लर्नबल मास्क $m_i \in [0, 1]$ क्लीन एक्टिवेशन ($h^{clean}$) और करप्टेड एक्टिवेशन ($h^{corrupted}$) के बीच इंटरपोलेट करता है:
  $$h^{node}_i = m_i \cdot h^{clean}_i + (1 - m_i) \cdot h^{corrupted}_i$$
  अनुकूलन (Optimization) ऐसे मास्क खोजने की कोशिश करता है जो मॉडल को सही आउटपुट (क्लीन व्यवहार से मेल खाता हुआ) देने की अनुमति दे, जबकि स्पर्सिटी (sparsity) को अधिकतम करे।

2.3 अनुकूलन और मास्क पैरामीट्रिज़ेशन

• हार्ड कॉन्क्रीट डिस्ट्रीब्यूशन (Hard Concrete Distribution): मास्क को हार्ड कॉन्क्रीट डिस्ट्रीब्यूशन का उपयोग करके पैरामीटराइज़ किया गया है ताकि ग्रेडिएंट-आधारित अनुकूलन को सक्षम बनाया जा सके और लगभग बाइनरी व्यवहार (0 या 1) प्रेरित किया जा सके।
• एकीकृत उद्देश्य (Unified Objective): प्रशिक्षण उद्देश्य तीन घटकों को जोड़ता है:
  1. फेथफुलनेस लॉस ($L_{faith}$): एक KL डाइवर्जेंस जो सुनिश्चित करता है कि प्रून किया गया सर्किट बेस मॉडल के आउटपुट डिस्ट्रीब्यूशन से मेल खाता है।
  2. टास्क लॉस ($L_{task}$): एक मार्जिन-आधारित रैंकिंग लॉस जो सुनिश्चित करता है कि सर्किट ग्राउंड-ट्रुथ और डिस्ट्रैक्टर विकल्पों के बीच सही ढंग से अंतर करता है।
  3. स्पर्सिटी रेगुलराइज़र ($L_0$): मास्क को शून्य होने के लिए प्रोत्साहित करता है, जिससे सर्किट का आकार कम होता है।
• पदानुक्रमित निरंतरता (Hierarchical Consistency): प्रशिक्षण के बाद, मास्क को एक ऐसे प्रतिबंध के साथ बाइनराइज़ किया जाता है जो यह सुनिश्चित करता है कि यदि एक पैरेंट नोड (जैसे, एक MLP ब्लॉक) को प्रून किया जाता है, तो उसके सभी चाइल्ड नोड्स (जैसे, उस ब्लॉक के भीतर के न्यूरॉन्स) भी निष्क्रिय हो जाते हैं।

3. प्रमुख योगदान

1. मल्टी-ग्रैनुलर फ्रेमवर्क: एक एकीकृत अनुकूलन फ्रेमवर्क का परिचय जो ट्रांसफॉर्मर ब्लॉक्स, अटेंशन हेड्स, MLP ब्लॉक्स और व्यक्तिगत न्यूरॉन्स में एक साथ मास्क सीखता है।
2. स्केलेबिलिटी और दक्षता: यह विधि एज प्रूनिंग की द्विघाती स्केलिंग और सभी किनारों के लिए मध्यवर्ती सक्रियणों को स्टोर करने की आवश्यकता से बचती है। यह पूर्व विधियों की तुलना में GPU मेमोरी उपयोग को 3× से 11× तक कम कर देती है, जिससे एक सिंगल 40GB GPU का उपयोग करके Llama 3.1-8B पर सर्किट डिस्कवरी संभव हो पाती है।
3. फाइन-ग्रेन्ड डिस्कवरी: यह दृष्टिकोण बेसलाइन की तुलना में काफी अधिक कॉम्पैक्ट सर्किट की पहचान करता है, विशेष रूप से MLP ब्लॉक और न्यूरॉन स्तर पर उच्च स्पर्सिटी प्राप्त करता है, जिसे पिछली विधियाँ हासिल करने में विफल रही थीं।

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

विधि का मूल्यांकन GPT-2, GPT-2-XL, और Llama मॉडलों का उपयोग करके तीन मानक सर्किट डिस्कवरी कार्यों (इंडायरेक्ट ऑब्जेक्ट आइडेंटिफिकेशन, जेंडरड प्रोनोन्स, और ग्रेटर देन) पर किया गया।

4.1 स्पर्सिटी और कॉम्पैक्टनेस

• MLP और न्यूरॉन प्रूनिंग: प्रस्तावित विधि एज-प्रूनिंग बेसलाइन्स (EAP, EP) और अनस्ट्रक्चर्ड वेट प्रूनिंग (DiscoGP) से काफी बेहतर प्रदर्शन करती है। सबसे प्रतिकूल स्थिति में, इसने पूर्व विधियों की तुलना में 33.34% अधिक MLP ब्लॉक्स और 59.8% अधिक न्यूरॉन्स को प्रून किया।
• टास्क-डिपेंडेंट स्ट्रक्चर: डिस्कवर किए गए सर्किट कार्य के आधार पर विशिष्ट टोपोलॉजी प्रदर्शित करते हैं:
  • IOI: वितरित MLP कंप्यूटेशन द्वारा संचालित; अटेंशन को अंतिम परतों को छोड़कर काफी हद तक प्रून कर दिया गया है।
  • जेंडरड प्रोनोन्स (GP): अत्यधिक स्थानीयकृत और विच्छिन्न (discontinuous); MLP कंप्यूटेशन कुछ परतों तक सीमित है, जबकि अटेंशन स्पर्स है।
  • ग्रेटर देन (GT): निष्क्रिय परतों के निरंतर ब्लॉकों और "स्किप-लाइक" निष्पादन पैटर्न के साथ अत्यधिक स्पर्सिटी प्रदर्शित करता है।

4.2 फेथफुलनेस और कंसिस्टेंसी

• परफॉरमेंस: प्रून किए गए सर्किट उच्च टास्क सटीकता (अक्सर >95%) और फुल मॉडल के सापेक्ष कम KL डाइवर्जेंस बनाए रखते हैं।
• कंसिस्टेंसी: पांच रैंडम सीड्स के माध्यम से प्रयोग दिखाते हैं कि डिस्कवर किए गए सर्किट (फेथफुलनेस मेट्रिक्स और स्पर्सिटी प्रोफाइल) अत्यधिक स्थिर हैं।
• गुणात्मक विश्लेषण: Llama 3.2-1B में, IOI कार्य के लिए जीवित बचे अटेंशन हेड्स को विशिष्ट कार्यात्मक भूमिकाएं (IO-Movers, S-Inhibitors) निभाने में पाया गया, जो GPT-2 में मैन्युअल रूप से डिस्कवर किए गए सर्किटों की नकल करते हैं, जो मेथड की मैकेनिस्टिक अंतर्दृष्टि प्राप्त करने की क्षमता को प्रमाणित करता है।

4.3 कंप्यूट एफिशिएंसी

• मेमोरी: इस विधि को एज-प्रूनिंग बेसलाइन्स (जैसे, GPT-2 पर EAP के लिए 72.79 GB बनाम 6.27 GB) की तुलना में कम से कम 3× कम मेमोरी की आवश्यकता होती है।
• समय: प्रशिक्षण 6 मिनट से कम समय में पूरा हो जाता है, जो एज-प्रूनिंग विधियों (लगभग 50 मिनट) की तुलना में काफी तेज़ है, हालांकि यह सिंगल-पास EAP विधि की तुलना में थोड़ा धीमा है (जो भारी मेमोरी उपयोग के बदले समय का त्याग करती है)।

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

पेपर का दावा है कि इसका प्राथमिक महत्व वर्तमान सर्किट डिस्कवरी की स्केलेबिलिटी और ग्रैनुलैरिटी बाधाओं को दूर करने में निहित है। एज प्रूनिंग से मल्टी-ग्रैनुलर नोड प्रूनिंग की ओर स्थानांतरित होकर, लेखक प्रदर्शित करते हैं कि:

• कई न्यूरॉन्स जिन्हें मोटे-स्तर (coarse-grained) की विधियों द्वारा बनाए रखा गया है, वे अनावश्यक हैं और प्रदर्शन पर नगण्य प्रभाव के साथ हटाए जा सकते हैं।
• बड़े मॉडल (जैसे, Llama 8B) छोटे मॉडलों के समान या उससे भी मजबूत प्रूनिंग का समर्थन करते हैं, जो यह सुझाव देता है कि बढ़ता हुआ स्केल विशिष्ट कार्यों के लिए सघन (dense) सर्किट की आवश्यकता के बजाय रेडंडेंसी (redundancy) लाता है।
• यह विधि बड़े पैमाने के मॉडलों में फाइन-ग्रेन्ड कॉज़ल सर्किट की खोज के लिए एक कम्प्यूटेशनल रूप से व्यवहार्य मार्ग प्रदान करती है, जो यह प्रकट करती है कि विभिन्न व्यवहार अटेंशन, MLP और न्यूरॉन गतिविधि के विशिष्ट पैटर्न द्वारा एनकोड किए जाते हैं।

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