Optimizing Robustness and Accuracy in Mixture of Experts: A Dual-Model Approach

Mixture of Experts (MoE) have shown remarkable success in leveraging specialized expert networks for complex machine learning tasks. However, their susceptibility to adversarial attacks presents a critical challenge for deployment in robust applications. This paper addresses the critical question of how to incorporate robustness into MoEs while maintaining high natural accuracy. We begin by analyzing the vulnerability of MoE components, finding that expert networks are notably more susceptible to adversarial attacks than the router. Based on this insight, we propose a targeted robust training technique that integrates a novel loss function to enhance the adversarial robustness of MoE, requiring only the robustification of one additional expert without compromising training or inference efficiency. Building on this, we introduce a dual-model strategy that linearly combines a standard MoE model with our robustified MoE model using a smoothing parameter. This approach allows for flexible control over the robustness-accuracy trade-off. We further provide theoretical foundations by deriving certified robustness bounds for both the single MoE and the dual-model. To push the boundaries of robustness and accuracy, we propose a novel joint training strategy JTDMoE for the dual-model. This joint training enhances both robustness and accuracy beyond what is achievable with separate models. Experimental results on CIFAR-10 and TinyImageNet datasets using ResNet18 and Vision Transformer (ViT) architectures demonstrate the effectiveness of our proposed methods.

Modern AI models are becoming more powerful—but also more vulnerable to tiny, invisible changes in input data. This is especially true for “Mixture-of-Experts” models, which divide work among several smaller networks called experts. Our research found that these expert networks are much easier to fool than the part of the system that assigns tasks to them. To fix this, we designed a new way to selectively train just one extra expert at a time, making the entire system more robust without slowing it down. We also created a two-part model that combines a normal version with a more defensive one, giving users a way to balance accuracy and safety. To support this, we developed new mathematical tools to show when and why our system can resist attacks. Our experiments on image datasets confirm that our method makes AI models both stronger and smarter. This work helps pave the way for safer AI systems in real-world applications like self-driving cars, medical imaging, or content moderation—where mistakes from small changes could be costly.