RASA: Routing-Aware Safety Alignment for Mixture-of-Experts Models

यहाँ शोध पत्र "RASA: Routing-Aware Safety Alignment for Mixture-of-Experts Models" का विस्तृत तकनीकी सारांश दिया गया है।

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

Mixture-of-Experts (MoE) भाषा मॉडल स्केलेबिलिटी और दक्षता में सुधार के लिए स्पार्स रूटिंग तंत्र (sparse routing mechanisms) का उपयोग करते हैं, जिससे किसी भी दिए गए इनपुट के लिए केवल विशेषज्ञों (experts) के एक उपसमूह को सक्रिय किया जाता है। हालाँकि, लेखक इन आर्किटेटेक्चर के लिए मानक सुरक्षा संरेखण (safety alignment) में एक महत्वपूर्ण विफलता मोड की पहचान करते हैं: एलाइनमेंट शॉर्टकट्स (Alignment Shortcuts)।

• घटना (The Phenomenon): जब MoE मॉडलों पर मानक फुल-पैरामीटर सुरक्षा फाइन-ट्यूनिंग लागू की जाती है, तो सुरक्षा मेट्रिक्स इसलिए नहीं सुधरते क्योंकि असुरक्षित व्यवहारों के लिए जिम्मेदार अंतर्निहित "सेफ्टी-क्रिटिकल एक्सपर्ट्स" (SCEs) को सुधारा गया है, बल्कि इसलिए सुधरते हैं क्योंकि मॉडल असुरक्षित इनपुट्स को पहले से सुरक्षित विशेषज्ञों की ओर रूट करना सीख जाता है या राउटर को रूढ़िवादी पैटर्न की ओर पक्षपाती बना देता है।
• परिणाम (The Consequence): यह एक भंगुर रक्षा (brittle defense) बनाता है। असुरक्षित विशेषज्ञ बिना सुधारे ही रहते हैं और उन्हें नए या अनुकूलित जेलब्रेक हमलों द्वारा पुन: सक्रिय किया जा सकता है जो राउटर को अलग पथ चुनने के लिए मजबूर करते हैं। इसके अलावा, यह दृष्टिकोण अक्सर "ओवर-रिफ्यूज़ल" (over-refusal) की ओर ले जाता है, जहाँ सुरक्षित विशेषज्ञ अत्यधिक प्रभावी हो जाते हैं, जिससे सामान्य उपयोगिता (general utility) कम हो जाती है।
• परिकल्पना (Hypothesis): फुल-पैरामीटर फाइन-ट्यूनिंग रूटिंग हेरफेर के माध्यम से सुरक्षा उद्देश्यों को संतुष्ट करती है, न कि वास्तविक विशेषज्ञ-स्तरीय पैरामीटर मरम्मत के माध्यम से।

2. कार्यप्रणाली: RASA (Methodology: RASA)

लेखक RASA (Routing-Aware Safety Alignment) प्रस्तावित करते हैं, जो एक ऐसा फ्रेमवर्क है जो एलाइनमेंट शॉर्टकट्स को रोकने के लिए रूटिंग नियंत्रण से विशेषज्ञ सुधार (expert correction) को स्पष्ट रूप से अलग करता है। RASA एक वैकल्पिक रणनीति का उपयोग करते हुए एक द्वि-स्तरीय अनुकूलन समस्या (bi-level optimization problem) के रूप में कार्य करता है।

मुख्य घटक (Core Components):

1. सुरक्षा-महत्वपूर्ण विशेषज्ञों (Safety-Critical Experts - SCEs) की पहचान:

   • यह विधि एंकर प्रॉम्प्ट्स (ऐसी दुर्भावनापूर्ण इनपुट जिन्हें मॉडल सही ढंग से अस्वीकार करता है) और एडवर्सरियल प्रॉम्प्ट्स (जेलब्रेक वेरिएंट जो सुरक्षा को सफलतापूर्वक बायपास कर देते हैं) वाले मिश्रित बैच बनाती है।
   • प्रत्येक विशेषज्ञ के लिए एडवर्सरियल एक्टिवेशन डिसक्रेपेंसी (AAD) की गणना की जाती है: $\Delta_{l,e} = P(\text{active} | \text{Adversarial}) - P(\text{active} | \text{Anchor})$।
   • उच्च AAD वाले विशेषज्ञों को SCEs के रूप में पहचाना जाता है, क्योंकि वे सफल जेलब्रेक द्वारा असंगत रूप से सक्रिय किए जाते हैं।

2. आंतरिक अनुकूलन: चयनात्मक विशेषज्ञ फाइन-ट्यूनिंग (Inner Optimization: Selective Expert Fine-tuning - SCE-FT):

   • लक्ष्य: पहचाने गए SCEs के मापदंडों (parameters) की मरम्मत करना।
   • तंत्र: राउटर को फ्रीज (freeze) कर दिया जाता है। केवल SCEs के मापदंडों को अपडेट किया जाता है ताकि रिफ्यूज़ल व्यवहार (जो विविध रिफ्यूज़ल प्रतिक्रियाओं पर प्रशिक्षित है) को इंजेक्ट किया जा सके। सामान्य क्षमताओं को बनाए रखने के लिए बेनिन (benign) विशेषज्ञों को अछूता रखा जाता है।
   • लॉस (Loss): सुरक्षा-संरेखित रिफ्यूज़ल लक्ष्यों के विरुद्ध मानक क्रॉस-एन्ट्रॉपी लॉस।

3. बाहरी अनुकूलन: राउटर निरंतरता (Outer Optimization: Router Consistency):

   • लक्ष्य: यह सुनिश्चित करना कि राउटर नए पथों के माध्यम से मरम्मत किए गए विशेषज्ञों को बायपास करने के लिए न सीखे।
   • तंत्र: विशेषज्ञ मापदंडों को फ्रीज कर दिया जाता है। राउटर को एडवर्सरियल इनपुट्स के रूटिंग वितरण को सुरक्षित एंकर इनपुट्स के रूटिंग वितरण के साथ संरेखित करने के लिए अनुकूलित किया जाता है।
   • लॉस: एक फॉरवर्ड KL डाइवर्जेंस लॉस ($L_{router}$) एडवर्सरियल रूटिंग वितरण और एंकर रूटिंग वितरण के बीच के अंतर को कम करता है। यह जेलब्रेक इनपुट्स को उसी सुरक्षित रूटिंग पथों का पालन करने के लिए मजबूर करता है जिसका उपयोग बेनिन इनपुट्स करते हैं।

ट्रेनिंग लूप (Training Loop): यह प्रक्रिया राउटर को फिक्स करने और फिर SCEs को अपडेट करने, और फिर SCEs को फिक्स करने और राउटर को अपडेट करने के बीच बारी-बारी से चलती है, जिससे यह सुनिश्चित होता है कि दोनों घटक मजबूती से संरेखित हैं।

3. प्रमुख योगदान (Key Contributions)

1. एलाइनमेंट शॉर्टकट्स की अनुभवजन्य खोज (Empirical Discovery of Alignment Shortcuts): यह पेपर साक्ष्य प्रदान करता है कि MoE मॉडलों में मानक फुल-पैरामीटर फाइन-ट्यूनिंग अक्सर रूटिंग शॉर्टकट्स पर निर्भर करती है, न कि असुरक्षित विशेषज्ञों की मरम्मत पर।
2. RASA फ्रेमवर्क: एक नवीन, आर्किटेक्चर-संरक्षण एलाइनमेंट विधि जो अतिरिक्त विशेषज्ञों या आर्किटेक्चरल परिवर्तनों की आवश्यकता के बिना, स्पष्ट रूप से SCEs को लक्षित करती है।
3. व्यापक मूल्यांकन: यह प्रदर्शित करता है कि लक्षित विशेषज्ञ मरम्मत, ग्लोबल फाइन-ट्यूनिंग या इन्फरेंस-टाइम स्टीयरिंग की तुलना में बेहतर मजबूती प्रदान करती है, जिसमें नियंत्रित ओवर-रिफ्यूज़ल और संरक्षित सामान्य क्षमताएं शामिल हैं।

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

लेखकों ने दो प्रतिनिधि MoE आर्किटेक्चर (Qwen3-30B-A3B और OLMoE-1B-7B-0125) पर विविध जेलब्रेक हमलों (FlipAttack, DeepInception, Persuasion) के विरुद्ध RASA का मूल्यांकन किया।

• मजबूती (Robustness): RASA ने संरेखित हमलों पर लगभग पूर्णतः हानिरहितता (≈1.0) प्राप्त की और मिश्रित एवं अनदेखे हमले वेक्टर्स पर बेसलाइन (Original, SteerMoE, SafeMoE-FT) से काफी बेहतर प्रदर्शन किया।
• सामान्यीकरण (Generalization): RASA के साथ प्रशिक्षित मॉडलों ने मजबूत क्रॉस-अटैक जनरलाइजेशन दिखाया। विभिन्न प्रकार के हमलों पर प्रशिक्षण से मॉडल को स्पष्ट मल्टी-अटैक पर्यवेक्षण के बिना विजातीय (disjoint) डेटासेट (जैसे, StrongReject) पर सामान्य करने में मदद मिली।
• उपयोगिता संरक्षण (Utility Preservation): फुल-पैरामीटर फाइन-ट्यूनिंग के विपरीत, जो अक्सर सामान्य बेंचमार्क (MMLU, GSM8K, TruthfulQA) में महत्वपूर्ण गिरावट और उच्च ओवर-रिफ्यूज़ल दर का कारण बनती है, RASA ने सामान्य क्षमताओं को बनाए रखा और कम ओवर-रिफ्यूज़ल दर बनाए रखी।
• कारण सत्यापन (Causal Validation):
  • रूटिंग बहाली (Routing Restoration): जब फुल-पैरामीटर फाइन-ट्यून किए गए मॉडलों में मूल रूटिंग पैटर्न को बहाल किया गया, तो सुरक्षा विफल हो गई (हानिरहितता 0.40 गिर गई)। RASA मॉडल उच्च सुरक्षा (0.94–0.96) बनाए रखते हैं, जो यह सिद्ध करता है कि उनकी सुरक्षा विशेषज्ञ मरम्मत में निहित है, न कि रूटिंग ट्रिक्स में।
  • रूटिंग वितरण (Routing Distribution): फुल-पैरामीटर फाइन-ट्यूनिंग के कारण रूटिंग वितरण में भारी बदलाव आया (उच्च KL डाइवर्जेंस), जबकि RASA ने मूल मॉडल के करीब रूटिंग संरचनाओं को बनाए रखा।

5. महत्व और प्रभाव (Significance and Impact)

• प्रतिमान परिवर्तन (Paradigm Shift): RASA इस धारणा को चुनौती देता है कि MoE मॉडलों की सुरक्षा के लिए ग्लोबल पैरामीटर अपडेट आवश्यक या पर्याप्त हैं। यह दर्शाता है कि लक्षित, विशेषज्ञ-स्तरीय हस्तक्षेप अधिक प्रभावी और कुशल है।
• दक्षता (Efficiency): यह विधि अत्यधिक डेटा-कुशल है, जिसे मजबूत संरेखण प्राप्त करने के लिए केवल एक अंश (25–50%) एडवर्सरियल नमूनों की आवश्यकता होती है।
• व्यावहारिकता (Practicality): यह अपने आर्किटेक्चर को संशोधित किए बिना या इन्फरेंस-टाइम हस्तक्षेप (SteerMoE के विपरीत) की आवश्यकता के बिना, बड़े पैमाने के MoE मॉडलों को संरेखित करने के लिए एक व्यावहारिक समाधान प्रदान करता है, जो इसे प्रोडक्शन वातावरण में तैनाती के लिए उपयुक्त बनाता है।
• भविष्य की दिशा (Future Direction): यह कार्य सुरक्षा संरेखण में आंतरिक रूटिंग तंत्र के विश्लेषण के महत्व को उजागर करता है, जो सुझाव देता है कि स्पार्स मॉडल्स के लिए भविष्य की संरेखण रणनीतियों को रूटिंग डायनेमिक्स और विशेषज्ञ विशेषज्ञता (expert specialization) के बीच परस्पर क्रिया को स्पष्ट रूप से संबोधित करना चाहिए।