Hierarchical vs. Flat Iteration in Shared-Weight Transformers

यहाँ संग-इल हान के शोध पत्र "Hierarchical vs. Flat Iteration in Shared-Weight Transformers" का विस्तृत तकनीकी सारांश दिया गया है।

1. समस्या विवरण (Problem Statement)

यह शोध पत्र डीप लर्निंग के एक मौलिक वास्तुशिल्प (architectural) ट्रेड-ऑफ को संबोधित करता है: स्टैकिंग बनाम इटरेशन (Stacking vs. Iteration)।

• मानक ट्रांसफॉर्मर (Standard Transformers) $L$ स्वतंत्र परतों को स्टैक करने पर निर्भर करते हैं, जिनमें से प्रत्येक के अलग पैरामीटर होते हैं ($O(Ld^2)$)। यह प्रतिनिधिक विविधता (सतही पैटर्न $\to$ सिंटैक्स $\to$ सिमेंटिक्स) की अनुमति देता है, लेकिन संग्रहीत भार (weights) और KV कैश के लिए उच्च मेमोरी लागत वहन करता है।
• साझा-भार इटरेशन (Shared-Weight Iteration) (जैसे, यूनिवर्सल ट्रांसफॉर्मर) स्टैकिंग को एक एकल साझा ब्लॉक द्वारा प्रतिस्थापित करने का प्रयास करता है जिसे $M$ बार लागू किया जाता है। जबकि यह संग्रहीत पैरामीटर्स को $O(d^2)$ तक कम कर देता है, पूर्ववर्ती कार्य सुझाव देते हैं कि "फ्लैट इटरेशन" (एक ही ब्लॉक को समान रूप से दोहराना) स्वतंत्र परतों के प्रदर्शन के साथ मेल करने में विफल रहता है, और अक्सर खराब वैलिडेशन क्रॉस-एन्ट्रॉपी (CE) पर स्थिर हो जाता है।

मुख्य प्रश्न: क्या एक पदानुक्रमित रूप से संरचित (hierarchically structured) साझा-भार पुनरावृत्ति (विशेष रूप से एक दो-गति प्रणाली) स्वतंत्र-परत स्टैकिंग की प्रतिनिधिक गुणवत्ता से मेल खा सकती है, जिससे प्रदर्शन से समझौता किए बिना बड़े पैमाने पर पैरामीटर संपीड़न प्राप्त किया जा सके?

2. कार्यप्रणाली: HRM-LM

लेखक HRM-LM (Language Modeling के लिए Hierarchical Reasoning Model) प्रस्तावित करते हैं, जो ऑटो-रिग्रेसिव लैंग्वेज मॉडलिंग के लिए HRM आर्किटेक्चर का एक अनुकूलन है।

आर्किटेक्चर डिज़ाइन

HRM-LM $L$ स्वतंत्र ट्रांसफॉर्मर परतों को एक पुनरावर्ती प्रणाली (recurrent system) से बदल देता है, जो $M = N \times T$ चरणों में साझा भार (shared weights) का उपयोग करता है, जिसे एक दो-गति पदानुक्रम (two-speed hierarchy) में व्यवस्थित किया गया है:

1. फास्ट-मॉड्यूल ($\theta_L$): प्रत्येक पुनरावर्ती चरण ($i=1 \dots M$) में अपडेट होता है।
   • कार्य: स्थानीय परिशोधन (local refinement) और सूक्ष्म-स्तरीय संदर्भ एकीकरण।
   • तंत्र: यह पिछले फास्ट स्टेट, स्लो स्टेट और इनपुट टोकन एनकोडिंग को संयोजित करने के लिए एक गेटेड अपडेट (GRU-शैली) का उपयोग करता है।
2. स्लो-मॉड्यूल ($\theta_H$): केवल प्रत्येक $T$ चरणों के बाद (साइकिल सीमाओं पर) अपडेट होता है।
   • कार्य: वैश्विक संपीड़न (global compression) और अमूर्त संदर्भ एकत्रीकरण।
   • तंत्र: यह स्लो स्टेट को अपडेट करने के लिए वर्तमान फास्ट स्टेट और पिछले स्लो स्टेट का उपयोग करता है। यह सीधे कच्चे टोकन इनपुट को नहीं देखता है, बल्कि केवल डिस्टिल्ड (distilled) फास्ट स्टेट को देखता है।
3. इनपुट/आउटपुट: एक निश्चित इनपुट एनकोडर ($\theta_I$) और एक फ्यूजन तंत्र जो अंतिम फास्ट स्टेट, स्लो स्टेट और इनपुट एनकोडिंग को संयोजित करता है।

प्रमुख तकनीकी घटक

• गेटेड अपडेट्स (Gated Updates): मानक रेसिडुअल कनेक्शन के विपरीत, फास्ट-मॉड्यूल स्टेट को संरक्षित करने और नया कंटेंट इंजेक्ट करने के बीच इंटरपोलेशन करने के लिए एक सीखे गए गेट $g \in (0,1)$ का उपयोग करता है, जिससे लंबी पुनरावृत्ति श्रृंखलाओं में स्थिरता सुनिश्चित होती है।
• ट्रंकेटेड बैकप्रोपैगेशन (TBPTT): मेमोरी प्रबंधित करने के लिए, ग्रेडिएंट्स केवल अंतिम $K$ चरणों ($K < M$) के माध्यम से बैकप्रोपैगेट किए जाते हैं, जबकि पिछले $M-K$ चरण बिना ग्रेडिएंट संचय के एक "वार्मअप" के रूप में कार्य करते हैं।
• मल्टी-पास सुपरविजन (Multi-Pass Supervision): मॉडल डेटा पर $S$ पास चलाता है, जो ग्रेडिएंट विस्फोट को रोकने के लिए पास के बीच स्टेट्स को अलग (detach) करता है, जिससे लर्निंग सिग्नल स्मूथ होता है।

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

1. वास्तुशिल्प नवाचार: HRM-LM का परिचय, जो एक साझा-भार वाला ट्रांसफॉर्मर वेरिएंट है जो विशिष्ट प्रति-परत पैरामीटर्स के बिना प्रतिनिधिक विविधता उत्पन्न करने के लिए फास्ट/स्लो दो-गति पदानुक्रम का उपयोग करता है।
2. अनुभवजन्य खोज (The "3.4 Nat Gap"): एक मजबूत निष्कर्ष कि पदानुक्रमित इटरेशन (hierarchical iteration) समान पैरामीटर संख्या पर फ्लैट इटरेशन की तुलना में काफी बेहतर प्रदर्शन करता है।
   • फ्लैट इटरेशन (UniTF): एक पैरामीटर-मैच्ड यूनिवर्सल ट्रांसफॉर्मर (1.2B पैरामीटर्स) हाइपरपैरामीटर ट्यूनिंग के बावजूद, 5 स्वतंत्र रन में $\approx 7.6$ नैट्स पर स्थिर (plateau) हो जाता है।
   • Hierarchical Iteration (HRM): HRM मॉडल (1.2B पैरामीटर्स) $\approx 4.18$ नैट्स पर अभिसरित (converge) होता है।
   • निष्कर्ष: यह अंतर ($\approx 3.4$ नैट्स) संरचनात्मक है, न कि पैरामीटर गणना या ऑप्टिमाइज़ेशन आर्टिफैक्ट्स के कारण। साझा-भार पुनरावृत्ति के सफल होने के लिए पदानुक्रम आवश्यक है।
3. मेमोरी दक्षता: HRM-LM संग्रहीत पैरामीटर्स को $O(Ld^2)$ से घटाकर $O(d^2)$ कर देता है ($L$-गुना अनकंडीशनल कमी)। जब $L > M$ होता है, तो यह KV कैश के आकार को भी $M/L$ के कारक से कम कर देता है।
4. ईमानदार सीमाएँ: पेपर स्पष्ट रूप से उन परिदृश्यों का विवरण देता है जहाँ HRM कम प्रदर्शन करता है (जैसे, बिना ट्यूनिंग के मानकीकृत हाइपरपैरामीटर के तहत) और स्वीकार करता है कि परिणाम वर्तमान में OpenWebText पर $\sim 1.2$B पैरामीटर्स तक सीमित हैं।

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

A. फ्लैट इटरेशन (UniTF) के विरुद्ध प्रदर्शन

• सेटअप: HRM (1.2B पैरामीटर्स) बनाम UniTF (1.2B पैरामीटर्स)।
• परिणाम: UniTF सीखने में विफल रहता है, 7.6 नैट्स पर स्थिर हो जाता है। HRM 4.18 नैट्स पर अभिसरित होता है।
• महत्व: यह सिद्ध करता है कि आंतरिक पदानुक्रम (Fast/Slow) ही वह महत्वपूर्ण कारक है जो साझा-भार वाले मॉडलों को सफल बनाता है, न कि केवल इटरेशन की क्रिया।

B. मानक ट्रांसफॉर्मर के विरुद्ध प्रदर्शन

• समान पैरामीटर्स (HRM बनाम T-L4):
  • मैच की गई हाइपरपैरामीटर सर्च (HPSearch) के तहत, HRM (4.31 नैट्स) Transformer L=4 (4.54 नैट्स) से 0.23 नैट्स से बेहतर प्रदर्शन करता है।
  • मानकीकृत समान-हाइपरपैरामीटर सेटिंग्स (MultiSeed) के तहत, T-L4 (4.70 नैट्स) थोड़ा HRM (4.90 नैट्स) से बेहतर प्रदर्शन करता है। यह सुझाव देता है कि HRM विशिष्ट ट्यूनिंग (लर्निंग रेट, वार्मअप) के प्रति संवेदनशील है, लेकिन सही ढंग से ट्यून किए जाने पर यह एक संरचनात्मक लाभ रखता है।
• समान FLOPs (HRM बनाम T-L12):
  • HRM (1.2B पैरामीटर्स) एक Transformer L=12 (3.4B पैरामीटर्स) से काफी बेहतर प्रदर्शन करता है, जिसे समान FLOPs के लिए प्रशिक्षित किया गया है (HRM: 3.90 नैट्स बनाम T-L12: 20k स्टेप्स पर 4.25 नैट्स)।
  • T-L12 को 60k स्टेप्स तक बढ़ाने के बाद भी (4.25 पर अभिसरण), यह 10k स्टेप्स पर HRM (4.18) से बदतर रहता है, जबकि इसके पास 2.8$\times$ अधिक पैरामीटर और 6$\times$ अधिक ट्रेनिंग स्टेप्स हैं।

C. दक्षता और विलंबता (Efficiency and Latency)

• मेमोरी: HRM, T-L4 की तुलना में संग्रहीत भारों में 2.5$\times$ की कमी और T-L12 की तुलना में 4$\times$ की कमी प्राप्त करता है।
• विलंबता (Latency): अनुक्रमिक पुनरावृत्ति (sequential recurrence) के कारण, HRM टोकन जनरेशन में 2–5$\times$ धीमा है (जैसे, T-L4 के 318 टोकन/सेकंड के मुकाबले 61 टोकन/सेकंड)।
• ट्रेड-ऑफ: HRM कंप्यूट समय (विलंबता) के बदले मेमोरी फुटप्रिंट का व्यापार करता है, जो इसे लैटेंसी-क्रिटिकल डेटा सेंटरों के बजाय मेमोरी-प्रतिबंधित एज डिवाइसेस (रोबोटिक्स, मोबाइल) के लिए आदर्श बनाता है।

D. एब्लेशन अध्ययन (Ablation Studies)

• ग्रेडिएंट विंडो ($K$): बड़ा $K$ (अधिक ग्रेडिएंट कवरेज) लगातार प्रदर्शन में सुधार करता है। स्थिरता के लिए $K/M \ge 15\text{--}20\%$ की आवश्यकता होती है।
• स्लो-मॉड्यूल का महत्व: स्लो-मॉड्यूल को अक्षम करने से प्रदर्शन में भारी गिरावट आती है (अप टू +4.0 नैट्स फॉर $N \times T = 12$), जो इसकी आवश्यक भूमिका की पुष्टि करता है।
• स्केलिंग ($M=48$): HRM 4$\times$ इटरेशन डेप्थ ($M=48$) पर भी प्रदर्शन बनाए रखता है, जो बताता है कि पदानुक्रम अच्छी तरह से स्केल करता है, जबकि फ्लैट इटरेशन अस्थिर रहता है।

5. महत्व और निहितार्थ

• सैद्धांतिक अंतर्दृष्टि: यह पेपर इस धारणा को चुनौती देता है कि प्रत्येक गहराई पर अलग पैरामीटर्स के लिए प्रतिनिधिक विविधता की आवश्यकता होती है। यह प्रदर्शित करता है कि साझा पैरामीटर सेट के भीतर कालिक पदानिर्भरता (temporal hierarchy) (फास्ट/स्लो टाइमस्केल) स्थानिक गहराई (स्टैकिंग) का स्थान ले सकती है।
• डिप्लॉयमेंट संदर्भ: HRM-LM को एज AI, रोबोटिक्स और लंबे संदर्भ वाले परिदृश्यों के लिए एक मेमोरी-कुशल विकल्प के रूप में रखा गया है जहाँ VRAM सीमित है लेकिन विलंबता बजट लचीला (सॉफ्ट रियल-टाइम) है।
• अन्य मॉडलों से संबंध:
  • Mamba/RWKV के विपरीत (जो लीनियर अटेंशन के माध्यम से KV कैश बाधाओं को हल करते हैं), HRM-LM साझा-भार के माध्यम से संग्रहीत भार (stored weight) की बाधा को हल करता है।
  • पेपर स्पष्ट करता है कि जबकि Mamba/RWKV इन्फरेंस के दौरान मेमोरी-कुशल हैं, वे गहराई के सापेक्ष संग्रहीत पैरामीटर गणनाओं को कम नहीं करते हैं, जबकि HRM ऐसा करता है।
• सीमाएँ और भविष्य का कार्य:
  • वर्तमान परिणाम $\sim 1.2$B पैरामीटर्स तक सीमित हैं। 7B+ तक स्केल करने के लिए ग्रेडिएंट कवरेज बॉटलनेक (जहाँ $M$ बढ़ने पर $K/M$ बहुत छोटा हो जाता है) को हल करना आवश्यक है।
  • "समान पैरामीटर" वाला तुलनात्मक अध्ययन बड़े पैमाने पर कम सार्थक हो जाता है जहाँ $L \gg M$; भविष्य के कार्य को "समान FLOPs" तुलनाओं पर ध्यान केंद्रित करना चाहिए।
  • पेपर यह दावा नहीं करता है कि HRM ट्रांसफॉर्मर का सार्वभौमिक विकल्प है, बल्कि यह मेमोरी-प्रतिबंधित शासन के लिए एक विशिष्ट संरचनात्मक लाभ को उजागर करता है।

निष्कर्ष

यह पेपर इस बात का पुख्ता अनुभवजन्य प्रमाण प्रदान करता है कि आंतरिक पदानिर्भरता (internal hierarchy) ही सफल साझा-भार वाले ट्रांसफॉर्मर के लिए गायब कड़ी है। फास्ट लोकल रिफाइनमेंट को स्लो ग्लोबल कम्प्रेशन से अलग करके, HRM-LM काफी कम संग्रहीत पैरामीटर्स के साथ मानक ट्रांसफॉर्मर के समकक्ष (और कभी-कभी बेहतर) प्रदर्शन प्राप्त करता है, जो मेमोरी-प्रतिबंधित हार्डवेयर पर बड़े भाषा मॉडलों को तैनात करने के लिए एक व्यवहार्य मार्ग प्रदान करता है।