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

TL;DR

This paper introduces a dual-model approach to improve the robustness of Mixture of Experts models against adversarial attacks while maintaining high accuracy, supported by theoretical bounds and experimental validation.

Contribution

It proposes a novel robust training technique and a dual-model strategy for MoEs, along with theoretical robustness bounds and a joint training method to enhance robustness and accuracy.

Findings

Enhanced adversarial robustness of MoE models demonstrated on CIFAR-10 and TinyImageNet.

The dual-model approach allows flexible robustness-accuracy trade-offs.

Theoretical certified robustness bounds are derived for the proposed models.

Abstract

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. The code is publicly available at https://github.com/TIML-Group/Robust-MoE-Dual-Model.