MESA: Improving MoE Safety Alignment via Decentralized Expertise
Abstract
Mixture-of-Experts (MoE) architectures scale Large Language Models (LLMs) efficiently, enabling greater capacity with reduced computational cost by dynamically routing inputs to relevant experts, yet introduce a critical vulnerability: Safety Sparsity, where safety capabilities concentrate in few experts, making them susceptible to adversarial bypassing. Meanwhile, conventional alignment methods uniformly adapt all parameters, ignoring their functional differences and inadvertently degrading performances. To address these challenges, we propose MESA (MoE Safety Alignment), a targeted alignment framework for MoE-based LLMs that strategically decentralizes safety responsibility to maximize coverage while minimizing interference with utility. Based on Optimal Transport (OT) theory, MESA operates through two mechanisms: (1) Expert Capacity Reallocation uses a transport cost matrix to distribute safety duties to the most cost-effective experts, and (2) Dynamic Routing Refinement constrains the router to precisely activate these decentralized modules. Experiments show that MESA achieves robust defensive performance against varied harmful benchmarks while preserving helpfulness. Code is available at https://github.com/lorraine021/MESA.
Lay Summary
(1) Mixture-of-Experts (MoE) architectures have emerged as a popular paradigm for scaling Large Language Models (LLMs), enabling greater capacity with reduced computational cost by dynamically routing inputs to the most relevant experts. This makes the models efficient, but we discovered a critical vulnerability: safety capabilities—such as refusing harmful requests—cluster in just a handful of experts. If an attacker finds a way to bypass these few experts, the entire safety shield collapses. Meanwhile, conventional safety alignment methods for making models safer uniformly train the entire model, damaging both expert knowledge and the routing system that decides which experts to activate. (2) We propose MESA, a method that redistributes safety responsibility across a broader set of experts with minimal impact on useful abilities such as math, coding and reasoning. It uses optimal transport, which is a mathematical technique for finding the most cost-effective redistribution, to identify the optimal experts suited to share the safety burden, and then retrains the model's routing mechanism to activate this expanded defense network while maintaining the stability of routing distribution obtained from pre-training. (3) This work marks a paradigm shift from content-centric optimization (leverage safety samples to inject safety capability in all parameters) to structural resource allocation, selectively injecting safety capability to the most suitable experts while considering the structural robustness. This offers a scalable and mathematically rigorous solution for reconciling safety alignment with general capability preservation in MoE architectures.