[अनुवाद] MoD(Mixture-of-Depths): Transformer-आधारित भाषा मॉडल गणना अनुकूलन के लिए एक दृष्टिकोण, और MoDE(MoD+MoE)

PyTorchKR

• हाल के समय में MoE(Mixture-of-Experts) तरीका, जिसमें Jamba, Qwen1.5-MoE, या DBRX जैसे मॉडलों की तरह कई 'expert' ब्लॉकों में से एक या अधिक को चुनकर input token को प्रोसेस किया जाता है, काफी उपयोग में आ रहा है. यह MoE तरीका विशाल LLM के कुछ हिस्सों को डायनेमिक रूप से नियंत्रित करके computational resources को कुशलता से आवंटित करता है और बेहतर प्रदर्शन दिखाता है.

• अगर MoE कई experts में से कुछ को चुनने के तरीके से 'मॉडल की चौड़ाई(width)' को डायनेमिक रूप से नियंत्रित करता है, तो यहाँ परिचित कराया जा रहा MoD(Mixture-of-Depths) 'मॉडल की गहराई(depth)' को डायनेमिक रूप से नियंत्रित करने का दृष्टिकोण अपनाता है. यानी, किसी विशेष token के लिए मॉडल की सभी layers पर computation करने के बजाय, कुछ layers को skip करके डायनेमिक computation किया जाता है. आइए इसे साथ में देखें.

---

MoD(Mixture-of-Depths): Transformer-आधारित भाषा मॉडल गणना अनुकूलन के लिए एक दृष्टिकोण (Mixture-of-Depths: Dynamically allocating compute in transformer-based language models)

परिचय

परंपरागत रूप से, transformer मॉडल input sequence के सभी हिस्सों को समान computational resources आवंटित करते हैं. लेकिन यहाँ परिचित कराया जा रहा MoD(Mixture-of-Depth) Transformer पेपर Transformer के कुछ specific sequence positions पर compute resources को डायनेमिक रूप से आवंटित करने और अलग-अलग layers के माध्यम से इस आवंटन को बदलने का प्रयास करता है. इस Mixture-of-Depths(MoD) रणनीति का उपयोग करने के लिए, लेखक पूरे computation की मात्रा को नियंत्रित करने हेतु यह तय करने के लिए top-k routing mechanism का उपयोग करते हैं कि किन tokens को प्रोसेस करना है. यह तकनीक मौजूदा मॉडलों के बराबर या उनसे बेहतर प्रदर्शन देती है, जबकि बहुत कम FLOPs की आवश्यकता होती है, और training के बाद sampling के दौरान अधिकतम 50% speedup प्रदान करती है.

Attention is All You Need शोध-पत्र में प्रस्तुत Transformer संरचना

हाल के transformer मॉडल और LLM(Large Language Model) लगातार आकार और जटिलता में बढ़ रहे हैं. ये मॉडल मानव-स्तरीय भाषा समझ और generation क्षमता के करीब पहुँचने के लिए अरबों parameters का उपयोग करते हैं, और परिणामस्वरूप उल्लेखनीय प्रदर्शन दिखाते हैं. लेकिन इस प्रगति के साथ भारी computational cost और energy consumption भी जुड़ा हुआ है, और मॉडल के आकार में वृद्धि training और inference समय को तेज़ी से बढ़ा देती है.

ऐसा इसलिए है क्योंकि Transformer input sequence के सभी tokens पर समान मात्रा में computation कर रहा है. लेकिन विचार यह है कि सभी tokens समान मात्रा में जानकारी नहीं दे सकते; कुछ tokens context को समझने में मुख्य भूमिका निभाते हैं, जबकि कुछ अपेक्षाकृत कम महत्वपूर्ण होते हैं.

MoD Transformer हर token के लिए computation को डायनेमिक रूप से आवंटित करके मॉडल को उन जगहों पर resources केंद्रित करने देता है जहाँ उनकी आवश्यकता हो. इससे मॉडल की कुल computational cost कम होती है, और महत्वपूर्ण tokens पर अधिक ध्यान देकर information processing की efficiency बढ़ाई जा सकती है. खासकर LLM जैसे बड़े मॉडलों में, MoD का दृष्टिकोण training और inference के दौरान समय और resources बचाने का एक महत्वपूर्ण तरीका बन सकता है.

MoD(Mixture-of-Depths) तकनीक का परिचय

MoD तकनीक transformer layer की प्रत्येक layer में किन tokens को प्रोसेस करना है, इसे डायनेमिक रूप से चुनने के लिए top-k routing mechanism को एकीकृत करती है. इसके माध्यम से केवल चुने गए कुछ tokens ही पूर्ण computation process से गुजरते हैं, जबकि अन्य tokens layer(s) को skip कर देते हैं.

चित्र 1 / बायाँ: MoD(Mixture-of-Depths) Transformer का overview: MoE(Mixture-of-Experts) संरचना की तरह router के माध्यम से computation path चुना जाता है. लेकिन MoE से अलग, यहाँ यह तय किया जाता है कि standard block(Self-Attention और MLP) पर computation करना है या नहीं. इस समय यदि input token($X_i$) router के निर्णय के अनुसार पूरा computation activate नहीं करता, तो वह उस layer की computation से नहीं गुजरता.
चित्र 1 / दायाँ: 64 tokens तक घटाए गए छोटे sequence के लिए routing decision का परिणाम. (X-अक्ष: sequence, Y-अक्ष: layer) बैंगनी रंग का अर्थ है कि उस layer की पूरी computation की गई, और नारंगी रंग का अर्थ है कि उस layer की computation को skip किया गया.

ऊपर दिए गए चित्र के विवरण की तरह, MoD transformer architecture के भीतर predictive router लागू करता है जो प्रत्येक token के लिए यह तय करता है कि layer computation चलानी है या नहीं. इससे computational resources को अपेक्षाकृत "महत्वपूर्ण" tokens पर केंद्रित किया जा सकता है.

MoD Transformer का training

MoD Transformer की training process मूल रूप से standard transformer के समान है. Input sequence मॉडल को दिया जाता है, और प्रत्येक layer से गुजरते हुए आवश्यक computations किए जाते हैं. लेकिन MoD Transformer में router के माध्यम से यह तय करने का अतिरिक्त चरण शामिल होता है कि हर token कौन-सा computation path अपनाएगा.

1. Router weight calculation: Router मॉडल की प्रत्येक layer में प्रत्येक token के लिए weights की गणना करता है. ये weights उस token को मिलने वाले computation की मात्रा को दर्शाते हैं और token के महत्व पर आधारित होते हैं.

2. Top tokens selection: Router गणना किए गए weights के आधार पर प्रत्येक layer के लिए top tokens का चयन करता है. इन tokens को full computation path की ओर भेजा जाता है.

3. Computation path assignment: चुने गए tokens को standard computation(Self-Attention और MLP) करने वाले path पर assign किया जाता है. बाकी tokens residual connection के माध्यम से आगे बढ़ते हैं और बिना अतिरिक्त computation के अगली layer में चले जाते हैं.

4. Loss calculation and backpropagation: मॉडल के output का मूल्यांकन अंतिम loss function के माध्यम से किया जाता है, और gradients मॉडल में backpropagate होकर weights को update करते हैं. इस प्रक्रिया में router भी train होता है, जिससे समय के साथ वह अधिक कुशल routing decisions ले सकता है.

MoD Transformer की routing schemes

Routing implementation MoD Transformer का मुख्य हिस्सा है. लेखकों ने 3 routing schemes की तुलना की और पाया कि Expert-Choice MoD तरीके की token routing सबसे अच्छा प्रदर्शन दिखाती है. आइए परिचित कराई गई अन्य routing schemes के साथ इसकी तुलना देखें:

Token-choice routing

Token-choice routing में हर token स्वयं तय करता है कि वह किस path से गुज़रेगा. मॉडल प्रत्येक token के लिए विभिन्न computation paths (जैसे अलग layers या processing blocks) के लिए probability distribution बनाता है, और उसी के अनुसार token को संबंधित path पर भेजता है.

Token-choice routing में हर token अपने लिए सबसे उपयुक्त computation path चुन सकता है, इसलिए routing काफी flexible होती है. साथ ही, मॉडल प्रत्येक token के महत्व या context को अधिक सटीकता से ध्यान में रख सकता है, जिससे individual token स्तर पर optimization संभव होता है.

Expert-choice routing

Expert-choice routing में मॉडल सीधे यह चुनता है कि प्रत्येक path (या "expert") किन tokens को प्रोसेस करेगा. यह प्रत्येक token के router weights के आधार पर top tokens चुनने वाले top-k mechanism के माध्यम से किया जाता है.

इस तरह expert-choice routing करने पर tokens को प्रत्येक path में अपेक्षाकृत समान रूप से वितरित किया जा सकता है, जिससे computational load संतुलित रहता है और मॉडल अपने computational resources का अधिक कुशल उपयोग कर पाता है. साथ ही, top tokens स्पष्ट रूप से तय होने के कारण मॉडल की predictability और stability बेहतर होती है.

Expert-choice MoD

Expert-choice MoD, expert-choice routing की अवधारणा को MoD Transformer की संरचना पर लागू करता है. इस तरीके में मॉडल प्रत्येक layer या processing step के लिए top tokens चुनता है, और केवल चुने गए tokens ही computation करने वाले path पर प्रोसेस होते हैं. बाकी tokens residual connection के माध्यम से कम लागत वाले path का अनुसरण करते हैं.

इस तरीके में केवल चुने गए tokens ही computation करते हैं, इसलिए पूरे मॉडल का FLOP उपयोग काफी कम हो जाता है. साथ ही, महत्वपूर्ण tokens के लिए पूरी layer computation सुनिश्चित करते हुए भी अनावश्यक computation को न्यूनतम किया जा सकता है, जिससे मॉडल का प्रदर्शन बना रह सकता है या बेहतर हो सकता है. अंत में, router के माध्यम से dynamic routing का लाभ यह है कि मॉडल input sequence की विशेषताओं के अनुसार computational resources को लचीले ढंग से आवंटित कर सकता है.

Routing implementation method

Routing का implementation नीचे दिए गए 3 मुख्य चरणों से संभव है:

1. Token embedding और router weight determination: मॉडल की प्रत्येक layer के लिए दिए गए sequence के token embeddings का सेट $X^i$ होता है. प्रत्येक token embedding $x^l_i$ के लिए router linear transformation के माध्यम से scalar weight $r^l_i$ की गणना करता है, जो यह दर्शाता है कि संबंधित token computation में भाग लेगा या नहीं.

2. Top weight determination: Token के router weight set $R^l$ में $\beta$-वें weight की percentile probability $P_\beta(R^l)$ खोजी जाती है. यहाँ $\beta$ यह तय करता है कि मॉडल प्रत्येक layer में tokens के किस अनुपात को प्रोसेस करेगा. इसलिए जिन tokens के router weights $P_\beta(R^l)$ से बड़े होते हैं, वे computation में भाग लेते हैं.

3. Computation path selection: प्रत्येक token का router weight $r^l_i$ यदि $P_\beta(R^l)$ से बड़ा या बराबर है, तो वह संबंधित layer की computation को प्रोसेस करता है; यदि छोटा है, तो वह सीधे अगली layer में भेज दिया जाता है. इससे केवल चुने गए tokens पर computation किया जाता है, जबकि बाकी tokens को ऐसे path से प्रोसेस किया जा सकता है जो computational cost उत्पन्न नहीं करता.

इस routing implementation में नीचे दिए गए बिंदुओं पर विशेष ध्यान देना चाहिए:

• Weights पर आधारित dynamic decision: Routing process हर token के weights के आधार पर डायनेमिक रूप से तय होती है, जो token के महत्व को दर्शाती है. महत्वपूर्ण tokens अधिक computation प्राप्त करते हैं और मॉडल के प्रदर्शन को optimize करते हैं, जबकि कम महत्वपूर्ण tokens residual path के माध्यम से cost बचाते हैं.

• Computational efficiency और performance optimization: इस routing mechanism के माध्यम से MoD Transformer computational resources का कुशल उपयोग करते हुए भी मॉडल का प्रदर्शन बनाए रख सकता है या बेहतर कर सकता है. खासकर, उन layers में जहाँ computation cost अधिक होती है, केवल महत्वपूर्ण tokens को प्रोसेस करके कुल FLOP उपयोग कम किया जा सकता है.

• Implementation की सरलता और efficiency: Routing को linear transformation और percentile calculation के माध्यम से implement किया जाता है, इसलिए यह अपेक्षाकृत सरल और कुशल है. इससे मॉडल training और inference के दौरान अतिरिक्त computational burden बढ़ाए बिना computational resources का dynamic allocation संभव हो पाता है.

MoD Transformer का कार्य उदाहरण

अब तक देखी गई बातों के आधार पर, आइए एक उदाहरण से MoD Transformer के काम करने का तरीका समझें.

उदाहरण के लिए, मान लें कि input sequence में 100 tokens हैं और router ने क्षमता को 10 पर सेट किया है. इस स्थिति में, प्रत्येक layer में router सबसे महत्वपूर्ण 10 tokens का चयन करेगा, और केवल ये tokens standard computation path के अनुसार प्रोसेस होंगे.

बाकी 90 tokens computation को skip करेंगे और residual connection के माध्यम से सीधे अगली layer में भेज दिए जाएँगे. इस प्रक्रिया से MoD Transformer कुल FLOP उपयोग को कम करते हुए भी महत्वपूर्ण tokens के लिए आवश्यक computation सुनिश्चित करता है.

MoD Transformer का performance comparison

Training performance

Hyperparameter optimization: 12.5% capacity routing blocks को full attention blocks के साथ alternating तरीके से रखने वाले MoD Transformer variants को train करने वाले प्रयोगों के माध्यम से optimal hyperparameter settings खोजी गईं. इस setting ने isoFLOP analysis में optimal baseline से बेहतर प्रदर्शन दिया.

• Computational efficiency: कुछ MoD variants ने मौजूदा मॉडलों की तुलना में कम FLOP के साथ समान प्रदर्शन हासिल किया और 66% तेज़ processing speed दिखाई.

Routing analysis

Routing decisions की sparsity: 12.5% capacity routing का उपयोग करने से अधिकांश tokens blocks को bypass कर गए और केवल कुछ महत्वपूर्ण tokens ही blocks से गुज़रे. इसे router weight distribution में भी देखा जा सकता है, और network ने depth के अनुसार कुछ specific tokens को प्राथमिकता देकर route भी किया.

Inference performance

Autoregressive inference: Training के समय non-causal top-k routing scheme से inference के समय causal prediction-based approach पर स्विच करने पर प्रदर्शन में लगभग कोई गिरावट नहीं थी. ऐसा संभवतः इसलिए था क्योंकि इस prediction problem को सीखना आसान था, और training के शुरुआती चरण से ही 97% से अधिक accuracy हासिल हो गई थी.

Mixture-of-Depths-and-Experts(MoDE) = MoD + MoE

MoDE(Mixture-of-Depths-and-Experts) का अर्थ MoD(Mixture-of-Depths) और MoE(Mixture of Experts) तकनीकों का एकीकरण है. इस तरीके का उद्देश्य दोनों approaches के फायदों को जोड़कर transformer मॉडल की computational efficiency और performance दोनों को साथ में optimize करना है. जहाँ MoD तरीका हर token के लिए अधिक गहराई या कम गहराई से computation करने का चयन करता है, वहीं MoE तरीका विभिन्न "experts" में से सबसे उपयुक्त expert के माध्यम से token को प्रोसेस करता है. MoDE इन दोनों approaches को एकीकृत करके tokens को कुछ specific blocks skip करने या चुनिंदा रूप से अधिक गहराई वाली computation करने की अनुमति देता है, और साथ ही विभिन्न experts में से optimal processing method चुनने में सक्षम बनाता है.

MoDE का implementation method

MoDE के दो मुख्य implementation तरीके हैं: Staged MoDE और Integrated MoDE.

1. Staged MoDE(चरणबद्ध MoDE): Staged MoDE में पहले MoD तकनीक लागू की जाती है ताकि प्रत्येक token का computation path तय हो सके, और उसके बाद MoE तकनीक लागू की जाती है ताकि चुने गए path में प्रोसेस होने वाले tokens को अलग-अलग expert blocks की ओर route किया जा सके. इस तरह मॉडल पहले computation की depth नियंत्रित करता है और फिर प्रत्येक चरण में सबसे उपयुक्त expert चुन सकता है.

2. Integrated MoDE(एकीकृत MoDE): Integrated MoDE, MoD और MoE तकनीकों को एक ही routing step में एकीकृत करके implement करता है. यानी token को route करते समय एक ही निर्णय में computation path की depth और प्रोसेस करने वाले expert दोनों का चयन किया जाता है. यह तरीका implementation के लिहाज़ से अधिक सरल और कुशल हो सकता है, और हर token के लिए सबसे उपयुक्त computation path और expert को एक साथ निर्धारित करने का लाभ देता है.

MoDE के लाभ

• Computational efficiency: MoDE हर token के लिए आवश्यक computation की मात्रा और प्रोसेस करने वाले expert को डायनेमिक रूप से तय करके मॉडल की कुल computational cost को प्रभावी ढंग से नियंत्रित करता है.
• Performance optimization: प्रत्येक token को सबसे उपयुक्त computation path और expert पर आवंटित करके MoDE मॉडल का प्रदर्शन अधिकतम करता है. यह खासकर जटिल समस्याओं को हल करने या विभिन्न प्रकार के data को प्रोसेस करने में उपयोगी हो सकता है.
• Flexibility और scalability: MoDE विभिन्न model structures और tasks पर लागू की जा सकने वाली flexibility प्रदान करता है, और ऐसा scalable framework देता है जो computational efficiency और performance दोनों को एक साथ बेहतर कर सकता है.

MoDE transformer मॉडलों के भविष्य के विकास की दिशा प्रस्तुत करता है, और computational cost को न्यूनतम रखते हुए मॉडल के प्रदर्शन को अधिकतम करने के तरीकों की खोज करने वाले शोध में महत्वपूर्ण योगदान देता है.

आगे पढ़ें

https://arxiv.org/abs/2404.02258

---

यह लेख GPT मॉडल से तैयार किए गए सार पर आधारित है, इसलिए संभव है कि इसमें मूल लेख की सामग्री या आशय से अलग तरह से व्यवस्थित कुछ बातें हों. यदि विषय में रुचि हो तो कृपया मूल लेख भी साथ में देखें. पढ़ते समय यदि आपको कोई अटपटा या गलत हिस्सा मिले, तो कृपया टिप्पणी में बताएं.

⚠️विज्ञापन⚠️: क्या PyTorch Korea User Group द्वारा तैयार किया गया यह लेख आपके लिए उपयोगी था? सदस्य के रूप में जुड़ें, तो हम आपको प्रमुख लेख ईमेल से भेजेंगे! (डिफ़ॉल्ट रूप से Weekly, लेकिन Daily में बदलना भी संभव है.)