Mixture-of-Depths: Transformer में computing resources को dynamically allocate करने की तकनीक

 (arxiv.org)
2 पॉइंट द्वारा GN⁺ 2024-04-08 | 1 टिप्पणियां | WhatsApp पर शेयर करें
  • Google DeepMind का Mixture-of-Depths(MoD) ऐसा तरीका है जिसमें transformer language model सभी tokens पर समान FLOPs खर्च नहीं करता, बल्कि हर layer में केवल ज़रूरी tokens को attention और MLP computation में शामिल करता है
  • हर layer का router, हर token के लिए scalar weight बनाता है, और पहले से तय top-k capacity के भीतर आने वाले tokens ही block से गुजरते हैं; बाकी residual connection के ज़रिये bypass हो जाते हैं
  • k को पहले से fixed रखने से computation graph और tensor size static बने रहते हैं, जबकि context के अनुसार token-स्तर पर computation allocation dynamically बदलता है
  • experiments में MoD ने isoFLOP आधार पर base transformer की तुलना में समान training time में कम loss दिया, या समान performance पर forward pass प्रति FLOPs घटाकर training और sampling steps को तेज़ बनाया
  • सबसे अच्छा configuration 12.5% capacity block को हर दूसरे block पर लागू करने वाला था, और auto-regressive sampling में predictor-based routing ने top-k की non-causal समस्या को कम किया

MoD किस compute waste को target करता है

  • सामान्य transformer language model, forward pass में सभी tokens पर एक जैसा computation इस्तेमाल करता है
  • MoD इस मान्यता से शुरू होता है कि सभी tokens और sequences को prediction के लिए समान समय और मेहनत नहीं चाहिए; इसलिए यह कुल compute budget घटाते हुए ज़रूरी tokens पर computation allocate करता है
  • मौजूदा conditional computation तरीके dynamic computation graph बना सकते हैं, जो आज के hardware द्वारा पसंद किए जाने वाले static computation graph से मेल नहीं खा सकते
  • अगर उपयोगकर्ता training से पहले कुल compute budget तय कर दे, तो model उसी के भीतर token-स्तर और layer-स्तर पर computation कहाँ खर्च करना है, यह सीखता है
    • computation, memory usage, और forward pass प्रति FLOPs का पहले से अनुमान लगाया जा सकता है
    • routing के लिए चुने जाने वाले tokens input context के अनुसार बदलते हैं

Routing structure और computation reduction का तरीका

  • MoD, MoE transformer की तरह router का उपयोग करता है, लेकिन कई experts में से एक चुनने के बजाय token को दो paths में से एक पर भेजता है
    • attention और MLP वाला standard block computation
    • value को जस का तस आगे बढ़ाने वाला residual connection
  • residual connection की computation cost कम होती है, और block output input value से तय होता है
  • अगर block की capacity को कुल token count T से कम रखा जाए, तो base transformer की तुलना में forward pass प्रति FLOPs घटते हैं
    • उदाहरण के लिए, अगर capacity को T/2 कर दिया जाए, तो attention की query-key matrix multiplication T² की जगह (T/2)² हो जाती है, यानी FLOPs लगभग 25% रह जाते हैं
  • capacity को 0 करने पर यह बहुत तेज़ हो जाता है, लेकिन transformer के ज़्यादातर parameters उपयोग नहीं होते, इसलिए performance गिरती है; और capacity को T रखने पर यह base transformer जैसा हो जाता है
  • लक्ष्य इन दोनों extremes के बीच ऐसा बिंदु खोजना है जहाँ base transformer के बराबर या बेहतर performance मिले और step speed भी बढ़े

top-k expert-choice routing

  • हर layer l में router, हर token embedding के लिए scalar weight बनाता है
  • उपयोगकर्ता द्वारा तय capacity C के अनुसार sequence tokens में से top-k वाले tokens ही attention और MLP computation में भाग लेते हैं
  • जो tokens चुने नहीं जाते, वे block computation skip करके residual connection के ज़रिये अगली layer में भेज दिए जाते हैं
  • यह तरीका expert-choice routing पर आधारित है
    • token-choice routing में token खुद path चुनता है, जिससे load balancing की समस्या आ सकती है
    • expert-choice routing में हर path top-k tokens चुनता है, इसलिए capacity ठीक-ठीक भरती है
  • MoD में paths असल में “block computation” और “bypass” ये दो ही होते हैं, इसलिए एक single top-k operation से tokens को दो groups में बाँटा जा सकता है
  • router weights को block output पर multiply किया जाता है, इसलिए वे gradient path पर रहते हैं और language modeling training के दौरान gradient descent से प्रभावित होते हैं

Sampling में आने वाली non-causality समस्या

  • expert-choice routing का top-k operation non-causal है
    • कोई token top-k में आएगा या नहीं, यह तय करने के लिए बाद के tokens के router weights की जानकारी भी चाहिए होती है
    • auto-regressive sampling में future token information का उपयोग नहीं किया जा सकता
  • इससे बचने के लिए दो तरीकों का परीक्षण किया गया
    • binary cross-entropy auxiliary loss, ताकि router output का sigmoid यह predict करे कि token top-k में चुना जाएगा या नहीं
    • router जैसे ही input लेने वाला, लेकिन stop-gradient लागू किया गया एक छोटा auxiliary MLP predictor
  • पहले तरीके का मुख्य language modeling objective पर लगभग 0.2~0.3% असर पड़ा
  • दूसरे predictor-based तरीके का language modeling objective पर कोई प्रभाव नहीं पड़ा, और step speed पर भी कोई meaningful असर नहीं था
  • auxiliary prediction task ने experimentally जल्दी ही 99% accuracy हासिल कर ली

Experimental results और observations

  • शुरुआती experiments कुल 6e18 FLOPs budget के भीतर MoD hyperparameters खोजने के रूप में किए गए
  • सबसे अच्छा MoD setup वह था जिसमें हर दूसरे block पर routing लागू की गई और top-k को 256 रखा गया
    • sequence length 2048 के आधार पर 256 tokens, यानी सिर्फ 12.5%, attention और MLP से गुजरते हैं
    • बाकी 1792 tokens, यानी 87.5%, block को bypass करते हैं
  • 220M parameter वाला MoD variant, isoFLOP-optimal baseline से थोड़ा बेहतर perform करता था, जबकि training के दौरान step speed 60% से अधिक तेज़ थी, और उसी hardware पर कुल training wall-clock time लगभग समान था
  • 6e18, 2e19, 1e20 FLOPs के isoFLOP analysis में भी MoD के FLOP-optimal models में baseline की तुलना में ज़्यादा parameters होने की प्रवृत्ति दिखी
  • ऐसे MoD variants मौजूद थे जो समान performance या कम loss देते हुए forward pass प्रति कम FLOPs इस्तेमाल करते थे
  • step speed improvement दो कारणों से मिला
    • कुछ tokens block bypass कर देते हैं, इसलिए समान model size पर FLOPs per parameter घट जाते हैं
    • baseline जैसी performance देने वाले छोटे MoD variants चुने जा सकते हैं
  • stochastic routing वाला MoD, base transformer और सामान्य MoD दोनों की तुलना में काफ़ी कमज़ोर performance दिखाता था, जिससे learned routing का महत्व सामने आया
  • auto-regressive evaluation में 256,000 sequences और 500M tokens के held-out data का उपयोग किया गया, और predictor-based routing पर स्विच करने पर भी performance drop छोटा था
  • MoD को MoE के साथ जोड़कर Mixture-of-Depths-and-Experts(MoDE) के रूप में भी बनाया जा सकता है
    • staged MoDE में attention से पहले तय किया जाता है कि token को block में भेजना है या bypass करना है
    • integrated MoDE में मौजूदा MLP experts के बीच एक “no-op” expert जोड़कर MoD routing को integrate किया जाता है
  • routing analysis में देखा गया कि कुछ tokens कई blocks से बार-बार गुजरते हैं, जबकि दूसरे tokens मौका मिलते ही bypass करते हैं
  • जो tokens अक्सर blocks से गुजरते हैं, उनके बारे में प्रारंभिक analysis में पाया गया कि वे output prediction की अधिक entropy से संबंधित थे, जो संभवतः यह दिखाता है कि वे predict करने में अधिक कठिन tokens हैं

संबंधित पढ़ाई

1 टिप्पणियां

 
GN⁺ 2024-04-08
Hacker News की राय

  • ज़्यादा जटिल routing निश्चित रूप से और आम हो जाएगी, ऐसा लगता है
    खासकर मुझे लगता है कि किसी दिन हम recursive routing तक पहुँचेंगे, जहाँ विशेषज्ञों के सेट से फिर से गुज़ारा जाएगा। आगे चलकर 'chain-of-thought' मॉडल के अंदर recursive तरीके से होगा

    • ऐसे काल्पनिक ऑब्जेक्ट को Recursive Neural Networks कहा जा सकता है
    • यहाँ जो बताया गया है, वह Universal Transformers वाली रिसर्च लाइन से थोड़ा मिलता-जुलता लगता है। इसमें input embeddings को एक ही transformer block से कई बार गुज़ारा जाता है, और एक अलग module तय करता है कि embedding पर्याप्त रूप से “पक” गई है या नहीं, फिर उसे बाहर निकालता है
      “experts” के विचार के ज्यादा करीब पिछले साल का Sparse Universal Transformers पेपर है, जो Universal Transformer और sparse mixture of experts को मिलाता है, ताकि gating mechanism तय करे कि कौन-से transformer blocks किस क्रम में इस्तेमाल होंगे
      यह मेरा विशेषज्ञता क्षेत्र नहीं है, लेकिन मेरी समझ में इसे सही से train करना मुश्किल है, और सामान्य transformer जैसे नतीजे पाने के लिए inference के समय कुल computation ज्यादा चाहिए होता है। फिर भी यह दिलचस्प दिशा है, और प्रति token computational steps की संख्या पर upper bound होना classical transformer architecture की बड़ी कमियों में से एक है
    • मेरे हिसाब से यह अभी तक इसलिए नहीं हुआ क्योंकि training के समय यह तय करने का तरीका नहीं है कि recursion कितनी बार करनी है
      random संख्या चुनने या कई recursion depths आज़माने से output “धुंधला” हो जाता है। यानी यह पता नहीं चलता कि किसी layer के output को final result के लिए महत्वपूर्ण जानकारी देनी चाहिए या अगले recursion round के लिए सबसे अच्छा input
    • attention मूल रूप से routing ही है, और ऐसे दूसरे routing तरीके model को ज्यादा मोटे विकल्प देते हैं, जिससे training संभावित रूप से आसान हो सकती है
    • रुझान निश्चित रूप से dynamic routing बढ़ने की ओर है, लेकिन मुझे लगता है कि MoE/MoD/MoDE deeper reasoning को संभव बनाने से ज्यादा weights में overlap घटाते हुए अतिरिक्त facts समाने में मदद करते हैं
      deeper reasoning शायद layer-level dynamicity से ज्यादा token-level dynamicity से आएगी। उदाहरण के लिए, हाल का Quiet-STaR पेपर है जो model द्वारा discard किए जाने वाले reasoning tokens output करता है: https://arxiv.org/abs/2403.09629

  • यह 2024 का सबसे महत्वपूर्ण पेपर भी हो सकता है
    यह विचार कि हमें ऐसे model चाहिए जिनमें सभी tokens पर समान मात्रा में computation खर्च न हो, काफी पहले से था, लेकिन इसके लिए पहली बार कोई convincing mechanism दिखा है

    Equipped with these new methods, we can sample autoregressively by choosing to route tokens to or around a block based on the router’s output, which does not depend on any information from future tokens. We provide empirical evidence that this is a relatively easy auxiliary task that quickly achieves 99% accuracy.
    यह थोड़ा हैरान करने वाला नहीं है?

    • Sparse Universal Transformer इससे पुराना है, और उसने routing-based early exit पहले ही किया था
    • सबसे महत्वपूर्ण तक? सभी tokens को पूरा context window चाहिए ही नहीं, यह तो obvious optimization होना चाहिए

  • आसान शुरुआती version:
    कल्पना कीजिए कि आपके पास एक smart assistant है जो भाषा समझ और process कर सकता है। आम तौर पर यह assistant हर शब्द पर समान ध्यान देता है, चाहे वह शब्द पूरे अर्थ के लिए कितना भी महत्वपूर्ण हो
    अब मान लीजिए कि हमने assistant को अपने “brain resources” ज्यादा smart तरीके से इस्तेमाल करना सिखाने का तरीका खोज लिया। हर शब्द पर समान ध्यान देने के बजाय, वह अर्थ समझने के लिए सबसे महत्वपूर्ण शब्दों पर ज्यादा focus करता है, और context के हिसाब से उस focus को तुरंत adjust करता है
    assistant overload न हो, इसके लिए किसी खास समय पर इस्तेमाल किए जा सकने वाले कुल “brain resources” की भी सीमा रखी जाती है। जैसे उसे budget देकर कहा जाए, “एक बार में केवल निश्चित संख्या के शब्दों पर ही resources इस्तेमाल कर सकते हो।” तब assistant को तय करना पड़ता है कि कौन-से शब्द सबसे महत्वपूर्ण हैं
    इस सीमा के बावजूद assistant resource usage को लचीले ढंग से adjust करता है। कुछ शब्दों पर ज्यादा लगाता है, कुछ पर कम, और स्थिति के हिसाब से priorities तय करता है
    इस तरह train होने पर assistant समझदारी से ध्यान देते हुए भी efficient तरीके से काम करता है। वह उन assistants जितना ही अच्छा समझता है जो हर शब्द पर समान ध्यान देते हैं, लेकिन कुल brain resources कम खर्च करता है, इसलिए responses और नई जानकारी process करना काफी तेज हो जाता है

    • समझता हूँ कि यह ELI5 है, लेकिन जिस तरीके को समझाया गया है, क्या वह वही नहीं है जो attention पहले से करता है? क्योंकि यह पिछले sequence में contextually महत्वपूर्ण शब्दों पर खास focus करता है

  • अपनी समझ के आधार पर मैंने यहाँ थोड़ा整理 किया है:
    https://lifeinthesingularity.com/p/googles-breakthroughs-in-...

    • लेख अच्छा है। New Scientist style की याद आती है। “समझाते हुए आगे बढ़ना, लेकिन छोटा रखना” जैसा एहसास देता है, इसलिए इस क्षेत्र का अंदाज़ा लगाने के लिए अच्छा है

  • यह mixture of experts से बहुत मिलता-जुलता है। बस tokens को कई experts में route करने के बजाय, “एक single expert में batch करना जिसे dynamically skip किया जा सकता है” जैसा है

    • दोनों को मिलाया जाए तो काफी शानदार हो सकता है। performance बनाए रखते हुए MoE computation को और कम किया जा सकता है

  • “This is more computationally efficient than performing a full content-based lookup across an entire memory buffer for each step in the future, and could be one step towards drastically increasing the context-length available for making a prediction.”
    क्या यही वह तरीका है जो 10 million token context window संभव बनाता है? या बात भविष्य में और लंबे context windows की हो रही है?

  • कुछ RASP algorithms को समझने और implement करने की कोशिश के बाद मेरा निष्कर्ष यह था कि कुछ functions को काम करने के लिए निश्चित संख्या में transformer layers चाहिए होती हैं
    इस logic को follow करें तो पता चलता है कि transformer द्वारा सीखे गए functions कई heads में फैल सकते हैं। ऐसे function को repeat करना problem को समझने और हल करने में बहुत valuable हो सकता है, लेकिन मौजूदा inference तरीका लगातार heads के group को बार-बार run नहीं कर सकता। यह पेपर सच में promising direction लगता है
    [1] https://arxiv.org/pdf/2106.06981.pdf
    [2] https://www.youtube.com/watch?v=t5LjgczaS80

  • LLMs के बहुत तेजी से आगे बढ़ने का शायद एकमात्र नुकसान यह है कि papers उस रफ्तार से तेज निकल रहे हैं जिस रफ्तार से कोई भी—कम से कम Google के बाहर—improvements को सीख और test कर सके
    जब मैंने deep learning शुरू की थी, तब ReLU और dropout का दौर था, और consumer 1080 पर code की एक-दो lines बदलकर कुछ घंटों में test कर सकते थे कि improvement हुआ या नहीं। अब लगता है कि Mistral जैसी जगहों के try करने तक कुछ हफ्ते इंतजार करना पड़ेगा

    • GPU poor वर्ग में आपका स्वागत है
      मैं quantization approaches पर focus करते हुए पुराने previous-generation GPUs से test कर रहा हूँ

  • hu-po AI papers की in-depth live streaming reviews करता है
    जोरदार recommendation है, और यहाँ वह चर्चा वाले mixture-of-depths paper को cover करता है: https://www.youtube.com/watch?v=Teru_qIdB8Y