专家权重平均：一种面向视觉Transformer的通用训练新方案

结构重参数化是一种面向卷积神经网络（CNNs）的通用训练方案，其能在不增加推理成本的情况下提升性能。随着视觉Transformer（ViTs）在各种视觉任务中逐渐超越CNNs，人们或许会问：是否存在一种专门针对ViTs的训练方案，同样能在不增加推理成本的情况下提升性能？近来，混合专家（MoE）受到越来越多的关注，因为它可通过稀疏激活的专家以固定成本高效扩展Transformer的容量。考虑到MoE也可视为一种多分支结构，我们能否利用MoE实现类似于结构重参数化的ViT训练方案？在本文中，我们肯定地回答了这些问题，提出了一种面向ViTs的通用训练新策略。具体而言，我们将ViT的训练与推理阶段解耦。在训练期间，我们用专门设计的、更高效的MoE替换ViT中的部分前馈网络（FFNs），这些MoE通过随机均匀划分将令牌分配给专家，并在每次迭代结束时对这些MoE执行专家权重平均（EWA）。训练完成后，我们通过平均专家将每个MoE转换为FFN，将模型还原为原始ViT用于推理。我们进一步提供了理论分析以说明其为何及如何生效。跨越多种2D和3D视觉任务、ViT架构和数据集的综合实验验证了所提训练方案的有效性与通用性。此外，我们的训练方案也可用于提升ViT微调时的性能。最后但同样重要的是，所提出的EWA技术可显著提升朴素MoE在各种2D视觉小数据集和3D视觉任务中的有效性。