Experts Weights Averaging: A New General Training Scheme for Vision Transformers

TL;DR

This work introduces Experts Weights Averaging (EWA), a general training scheme for Vision Transformers (ViTs) that decouples training from inference by replacing select FFNs with Random Uniform Partition MoE blocks and performing expert-weight averaging after each update. After training, each MoE is collapsed into a single FFN, yielding a dense ViT for inference with no additional latency or parameters. The authors provide theoretical analysis showing weight-decay-like dynamics and historical weight aggregation, and validate EWA across diverse 2D and 3D vision tasks, architectures, and datasets, including successful fine-tuning scenarios. They also demonstrate that EWA improves naive MoE performance on small datasets when applied in Early EWA, highlighting broad applicability and robustness. Overall, EWA offers a practical, general approach to boost ViT performance without changing deployment costs, with potential impact on ViT training and small-data MoE applications.

Abstract

Structural re-parameterization is a general training scheme for Convolutional Neural Networks (CNNs), which achieves performance improvement without increasing inference cost. As Vision Transformers (ViTs) are gradually surpassing CNNs in various visual tasks, one may question: if a training scheme specifically for ViTs exists that can also achieve performance improvement without increasing inference cost? Recently, Mixture-of-Experts (MoE) has attracted increasing attention, as it can efficiently scale up the capacity of Transformers at a fixed cost through sparsely activated experts. Considering that MoE can also be viewed as a multi-branch structure, can we utilize MoE to implement a ViT training scheme similar to structural re-parameterization? In this paper, we affirmatively answer these questions, with a new general training strategy for ViTs. Specifically, we decouple the training and inference phases of ViTs. During training, we replace some Feed-Forward Networks (FFNs) of the ViT with specially designed, more efficient MoEs that assign tokens to experts by random uniform partition, and perform Experts Weights Averaging (EWA) on these MoEs at the end of each iteration. After training, we convert each MoE into an FFN by averaging the experts, transforming the model back into original ViT for inference. We further provide a theoretical analysis to show why and how it works. Comprehensive experiments across various 2D and 3D visual tasks, ViT architectures, and datasets validate the effectiveness and generalizability of the proposed training scheme. Besides, our training scheme can also be applied to improve performance when fine-tuning ViTs. Lastly, but equally important, the proposed EWA technique can significantly improve the effectiveness of naive MoE in various 2D visual small datasets and 3D visual tasks.

Figures (2)

• Figure 1: Training-time & inference-time comparison of (a) recent MoE block, (b) typical structural re-parameterization block, and (c) our EWA training block. Each FFN denotes an expert of MoE.
• Figure 2: Ablation study on share schedules and share rates of the proposed EWA scheme.