FLEX-MoE: Federated Mixture-of-Experts with Load-balanced Expert Assignment for Edge Computing

तकनीकी सारांश: FLEX-MoE

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

यह शोध पत्र वायरलेस और IoT एज नेटवर्क के भीतर संसाधन-सीमित (resource-constrained) वातावरण में मिचर-ऑफ-एक्सपर्ट्स (Mixture-of-Experts - MoE) मॉडलों की तैनाती को संबोधित करता है। जबकि MoE सशर्त गणना (conditional computation) के माध्यम से स्केलेबिलिटी प्रदान करता है, और फेडरेटेड लर्निंग (FL) गोपनीयता और विकेंद्रीकृत प्रशिक्षण सुनिश्चित करती है, इन दोनों का संयोजन एज वातावरण में दो महत्वपूर्ण बाधाओं का सामना करता है:

1. संसाधन संबंधी सीमाएँ (Resource Constraints): एज क्लाइंट्स (जैसे कि UAVs, स्मार्टफोन, IoT नोड्स) के पास सीमित मेमोरी और प्रोसेसिंग पावर होती है, जिससे पारंपरिक वितरित MoE विधियों के लिए आवश्यक विशेषज्ञों (experts) के पूर्ण सेट को संग्रहीत या सिंक्रोनाइज़ करना असंभव हो जाता है।
2. विशेषज्ञ भार असंतुलन (Expert Load Imbalance): FL सेटिंग्स में non-IID (non-independent and identically distributed) डेटा के साथ, क्लाइंट स्वाभाविक रूप से उन विशिष्ट विशेषज्ञों के प्रति प्राथमिकता विकसित करते हैं जो उनके स्थानीय डेटा पैटर्न से मेल खाते हैं। मौजूदा दृष्टिकोण, जैसे कि ग्रीडी टॉप-$k$ चयन (greedy top-$k$ selection), स्थानीय वैयक्तिकरण (personalization) को प्राथमिकता देते हैं लेकिन वैश्विक भार वितरण (global load distribution) की अनदेखी करते हैं। इससे विशेषज्ञ उपयोग में गंभीर झुकाव (skew) पैदा होता है, जहाँ कुछ "लोकप्रिय" विशेषज्ञों का अत्यधिक उपयोग किया जाता है जबकि अन्य का उपयोग कम होता है, जिससे समग्र मॉडल सामान्यीकरण (generalization) और प्रदर्शन गिर जाता है।

मुख्य चुनौती एक ऐसा ढांचा डिजाइन करने की है जो सख्त क्लाइंट क्षमता सीमाओं के तहत विशेषज्ञ असाइनमेंट (क्लाइंट को उपयुक्त विशेषज्ञों से मिलाना) और लोड बैलेंसिंग (यह सुनिश्चित करना कि सभी विशेषज्ञों को पर्याप्त रूप से प्रशिक्षित किया जाए) को संयुक्त रूप से अनुकूलित (optimize) कर सके।

कार्यप्रणाली: FLEX-MoE

लेखक FLEX-MoE (Federated MoE with Load-balanced EXpert assignment) प्रस्तावित करते हैं, जो एक सर्वर-समन्वित ढांचा है जो क्लाइंट्स को विशेषज्ञों के उपसमुच्चयों (subsets) को गतिशील रूप से असाइन करता है। यह वर्कफ़्लो प्रति संचार राउंड (communication round) चार चरणों में कार्य करता है:

1. क्लाइंट-एक्सपर्ट फिटनेस स्कोरिंग (Client-Expert Fitness Scoring)

एक विशिष्ट क्लाइंट के स्थानीय डेटा के लिए किसी विशेषज्ञ की उपयुक्तता को मापने के लिए, सर्वर एक क्लाइंट-एक्सपर्ट फिटनेस स्कोर मैट्रिक्स $Q$ बनाए रखता है। यह स्कोर स्थानीय प्रशिक्षण दौरों से प्राप्त फीडबैक के आधार पर अपडेट किया जाता है, विशेष रूप से:

• सटीकता-संचालित (Accuracy-driven): विशेषज्ञ द्वारा प्राप्त स्थानीय वर्गीकरण सटीकता का उपयोग करके।
• लॉस-संचालित (Loss-driven): स्थानीय लॉस के एक्सपोनेंशियल फंक्शन ($\exp(-\alpha_L \ell)$) का उपयोग करके।
  समय के साथ उतार-चढ़ाव को सुचारू बनाने के लिए स्कोर को एक्सपोनेंशियल मूविंग एवरेज (EMA) का उपयोग करके अपडेट किया जाता है।

2. लोड-बैलेंस्ड एक्सपर्ट असाइनमेंट एल्गोरिदम (Load-Balanced Expert Assignment Algorithm)

ग्रीडी विधियों के विपरीत जो प्रत्येक क्लाइंट को केवल शीर्ष-$k$ विशेषज्ञ सौंपती हैं, FLEX-MoE असाइनमेंट को एक प्रतिबंधित अनुकूलन समस्या (constrained optimization problem) के रूप में तैयार करता है (जो जनरलाइज्ड असाइनमेंट प्रॉब्लम से अनुकूलित है)।

• उद्देश्य: कुल संचित क्लाइंट-एक्सपर्ट फिटनेस स्कोर ($\sum Q(c,e) \cdot X_{c,e}$) को अधिकतम करना।
• बाधाएं (Constraints):
  • क्लाइंट क्षमता: प्रत्येक क्लाइंट $c$ को ठीक $k_c$ विशेषज्ञ आवंटित किए जाते हैं।
  • लोड सीमाएं: प्रत्येक विशेषज्ञ $e$ को सौंपा गया कुल डेटा लोड गतिशील रूप से गणना की गई ऊपरी ($\Gamma$) और निचली ($L$) सीमाओं के भीतर होना चाहिए।
• गतिशील सीमाएं (Dynamic Bounds): सर्वर कुल सिस्टम क्षमता के आधार पर एक लक्षित लोड $\tau$ की गणना करता है और पिछले दौर के असंतुलन को ठीक करने के लिए (एक्सपोनेंशियल स्मूथिंग का उपयोग करके) सीमाओं को समायोजित करता है।
• समाधान: सर्वर इष्टतम बाइनरी असाइनमेंट मैट्रिक्स $X_{c,e}$ निर्धारित करने के लिए इस लीनियर प्रोग्रामिंग समस्या को हल करता है।

3. प्रशिक्षण और एकत्रीकरण (Training and Aggregation)

• स्थानीय प्रशिक्षण: क्लाइंट अपने आवंटित विशेषज्ञ मापदंडों (parameters) और एक साझा फीचर एक्सट्रैक्टर को डाउनलोड करते हैं। वे अपने इनपुट को अपने आवंटित विशेषज्ञों तक रूट करने के लिए एक व्यक्तिगत गेटिंग नेटवर्क का उपयोग करके स्थानीय रूप से प्रशिक्षण लेते हैं।
• फीडबैक: क्लाइंट सर्वर को पैरामीटर अपडेट और फीडबैक मेट्रिक्स (उपयोग गणना, सटीकता, लॉस) अपलोड करते हैं।
• सर्वर एकत्रीकरण: साझा फीचर एक्सट्रैक्टर को FedAvg के माध्यम से एकत्रित किया जाता है, और विशेषज्ञ मापदंडों को उपयोग-भारित औसत (usage-weighted averaging) के माध्यम से अपडेट किया जाता है, यह सुनिश्चित करता है कि अधिक डेटा पर प्रशिक्षित विशेषज्ञों का वैश्विक अपडेट में अधिक योगदान हो।

मुख्य योगदान (Key Contributions)

लेखक स्पष्ट रूप से निम्नलिखित योगदानों का दावा करते हैं:

1. नवीन समस्या सूत्रीकरण (Novel Problem Formulation): यह संसाधन-सीमित वातावरण में फेडरेटेड MoE के तहत विशेषज्ञ असाइनमेंट और लोड बैलेंसिंग की संयुक्त समस्या को स्पष्ट रूप से संबोधित करने वाला पहला कार्य है।
2. फिटनेस स्कोरिंग तंत्र: एक मात्रात्मक क्लाइंट-एक्सपर्ट फिटनेस स्कोर का परिचय जो प्रशिक्षण फीडबैक को शामिल करता है ताकि अनुकूलित विशेषज्ञ चयन सक्षम किया जा सके, जिससे व्यक्तिगतकरण और प्रशिक्षण दक्षता दोनों में सुधार होता है।
3. अनुकूलन-आधारित असाइनमेंट: एक लोड-बैलेंस्ड असाइनमेंट एल्गोरिदम का डिज़ाइन जो विशेषज्ञता (specialization) और सिस्टम-व्यापी लोड बैलेंस को संयुक्त रूप से अनुकूलित करता है, जो ग्रीडी दृष्टिकोणों में देखी जाने वाली विशेषज्ञ असमान प्रशिक्षण को स्पष्ट रूप से रोकता है।
4. अनुभवजन्य सत्यापन (Empirical Validation): यह प्रदर्शन का प्रदर्शन कि लोड बैलेंस लागू करना प्रदर्शन के लिए आवश्यक है, विशेष रूप से अत्यधिक विषम (skewed) डेटा वितरण के तहत, जहाँ यह विधि सटीकता और लोड बैलेंस मेट्रिक्स दोनों में मौजूदा रणनीतियों से बेहतर प्रदर्शन करती है।

प्रयोगात्मक परिणाम (Experimental Results)

ढांचे का मूल्यांकन तीन डेटासेट्स (CIFAR-10, EMNIST, GTSRB) पर IID और non-IID (Dirichlet और Class Partition) वितरणों के तहत किया गया।

• सटीकता (Accuracy): FLEX-MoE ने अधिकांश कॉन्फ़िगरेशन में लगातार उच्चतम या उच्चतम के करीब टेस्ट सटीकता प्राप्त की। जबकि ग्रीडी विधियों ने विशिष्ट परिदृश्यों (जैसे लॉस-ड्रिवन मेट्रिक्स के साथ GTSRB) में कभी-कभी सटीकता के बराबर प्रदर्शन किया, FLEX-MoE ने विशेष रूप से non-IID सेटिंग्स में बेहतर मजबूती (robustness) प्रदर्शित की।
• लोड बैलेंस (Load Balance): FLEX-MoE ने बेसलाइन की तुलना में विशेषज्ञ उपयोग में काफी कम विचरण गुणांक (Coefficient of Variation - CV) प्राप्त किया। जबकि ग्रीडी विधियों में अक्सर CV मान 0.2 से अधिक हो जाता था (जो गंभीर असंतुलन का संकेत देता है), FLEX-MoE ने लगभग समान विशेषज्ञ प्रशिक्षण का संकेत देते हुए 0.003–0.005 के करीब CV मान बनाए रखा।
• गंभीर डेटा विषमता (Severe Data Skew): अत्यधिक विषम परिदृश्यों (जैसे 2-Class partition या Dirichlet $\alpha=0.1$) में, ग्रीडी विधियों को लोड असंतुलन के कारण महत्वपूर्ण सटीकता गिरावट का सामना करना पड़ा। FLEX-MOE ने विशेषज्ञ भुखमरी (expert starvation) को रोककर उच्च सटीकता (जैसे 2-Class CIFAR-10 के लिए Greedy के 51.28% बनाम 61.00%) बनाए रखी।
• ओवरहेड (Overhead): सर्वर-साइड अनुकूलन के लिए कम्प्यूटेशनल ओवरहेड मामूली था, जिसके परिणामस्वरूप फेडरेटेड प्रशिक्षण की तुलना में कुल प्रशिक्षण समय में केवल ~6% की वृद्धि हुई, जिसे व्यावहारिक माना गया।

महत्व और दावे (Significance and Claims)

यह शोध पत्र FLEX-MoE को AI-नेटिव वायरलेस संचार और एज कंप्यूटिंग के लिए एक आवश्यक विकास के रूप में प्रस्तुत करता है। लेखक तर्क देते हैं कि:

• लोड बैलेंसिंग वैकल्पिक नहीं है: फेडरेटेड MoE में, लोड असंतुलन की अनदेखी करने से मॉडल की गुणवत्ता गिर जाती है, विशेष रूप से जब डेटा विषमता (heterogeneity) उच्च होती है।
• वैश्विक अनुकूलन स्थानीय लालच (local greediness) से बेहतर है: जबकि स्थानीय वैयक्तिकरण महत्वपूर्ण है, एक एकीकृत MoE मॉडल के लिए यह आवश्यक है कि सभी विशेषज्ञों को पर्याप्त रूप से संलग्न किया जाए ताकि सामान्यीकरण सुनिश्चित हो सके।
• स्केलेबिलिटी: क्लाइंट्स को विशेषज्ञों के केवल उपसमुच्चयों को स्टोर करने की अनुमति देकर, FLEX-MoE बड़े पैमाने के MoE मॉडलों को संसाधन-सीमित एज डिवाइस पर तैनात करना सक्षम बनाता है, जो अगली पीढ़ी के वायरलेस नेटवर्क के लिए एक पूर्व शर्त है।

कार्य इस निष्कर्ष पर पहुँचता है कि विषम, संसाधन-सीमित वातावरण में प्रभावी फेडरेटेड MoE की तैनाती के लिए विशेषज्ञ असाइनमेंट और लोड बैलेंसिंग को संयुक्त रूप से अनुकूलित करना महत्वपूर्ण है।