Swin-Transformer中的Swin-MoE小型模型配置解析

模型背景与概述

Swin-Transformer是微软研究院提出的一种基于Transformer的视觉骨干网络，它通过引入分层特征图和移位窗口机制，在计算机视觉任务中取得了显著效果。而Swin-MoE(Swin Mixture of Experts)是在此基础上引入专家混合(Mixture of Experts)机制的变体，旨在提升模型容量和效率。

本文要解析的配置文件定义了一个小型(Small)规模的Swin-MoE模型，专为ImageNet-22K数据集设计，使用192x192输入分辨率，并在32个GPU上进行训练。

数据配置解析

DATA:
  DATASET: imagenet22K
  IMG_SIZE: 192

这部分配置指定了训练使用的数据集和输入图像尺寸：

• DATASET: 使用ImageNet-22K数据集，这是一个包含约1400万张图像和21841个类别的大规模图像分类数据集
• IMG_SIZE: 输入图像将被调整为192x192像素的分辨率

模型架构详解

MODEL:
  TYPE: swin_moe
  NAME: swin_moe_small_patch4_window12_192_32expert_32gpu_22k
  DROP_PATH_RATE: 0.2

• TYPE: 指定模型类型为swin_moe，即带有专家混合机制的Swin-Transformer
• NAME: 模型名称包含了关键参数信息：小型规模、4x4的patch大小、12的窗口大小、192输入分辨率、32个专家、32个GPU训练、22K数据集
• DROP_PATH_RATE: 0.2的DropPath率，这是一种正则化技术，有助于防止过拟合

Swin-MoE核心参数

SWIN_MOE:
  EMBED_DIM: 96
  DEPTHS: [ 2, 2, 18, 2 ]
  NUM_HEADS: [ 3, 6, 12, 24 ]
  WINDOW_SIZE: 12

• EMBED_DIM: 96维的嵌入维度，这是模型的基础特征维度
• DEPTHS: 各阶段的Transformer块数量，这里配置为[2,2,18,2]，意味着有4个阶段，第三阶段特别深
• NUM_HEADS: 各阶段的多头注意力头数，随着网络深度增加而增加
• WINDOW_SIZE: 12x12的局部窗口大小，这是Swin-Transformer的核心创新之一

专家混合(MoE)特定配置

MLP_FC2_BIAS: False
INIT_STD: 0.005
MOE_BLOCKS: [ [ -1 ], [ -1 ], [ 1, 3, 5, 7, 9, 11, 13, 15, 17 ], [ 1 ] ]
NUM_LOCAL_EXPERTS: 1
TOP_VALUE: 1
CAPACITY_FACTOR: 1.25
IS_GSHARD_LOSS: False
MOE_DROP: 0.1
AUX_LOSS_WEIGHT: 0.01

这些参数控制着MoE机制的行为：

• MOE_BLOCKS: 指定哪些块使用MoE机制。-1表示不使用，数字表示使用。这里第三阶段的大部分块(奇数编号)使用了MoE
• NUM_LOCAL_EXPERTS: 每个MoE层有1个本地专家
• TOP_VALUE: 每个token选择1个专家
• CAPACITY_FACTOR: 1.25的容量因子，控制专家负载平衡
• MOE_DROP: 0.1的MoE层dropout率
• AUX_LOSS_WEIGHT: 0.01的辅助损失权重，用于平衡专家利用率

训练策略配置

TRAIN:
  EPOCHS: 90
  WARMUP_EPOCHS: 10
  WEIGHT_DECAY: 0.1
  BASE_LR: 1.25e-4 # 4096 batch-size
  WARMUP_LR: 1.25e-7
  MIN_LR: 1.25e-6
  CLIP_GRAD: 3.0

• EPOCHS: 总共训练90个epoch
• WARMUP_EPOCHS: 前10个epoch用于学习率预热
• BASE_LR: 基础学习率1.25e-4，针对4096的总batch size
• WARMUP_LR: 预热阶段起始学习率1.25e-7
• MIN_LR: 最小学习率1.25e-6
• CLIP_GRAD: 梯度裁剪阈值为3.0，防止梯度爆炸

测试配置

TEST:
  SHUFFLE: True

• SHUFFLE: 测试时打乱数据顺序，通常用于更可靠的评估

技术要点总结

1. 模型规模: 这是一个"Small"规模的Swin-MoE模型，适合中等规模的计算资源
2. 专家分布: 专家主要集中在网络的第三阶段(深层)，这与人类视觉系统深层处理更复杂特征的特性相符
3. 训练策略: 采用了学习率预热、梯度裁剪等稳定训练的技术
4. MoE设计: 使用相对保守的MoE配置(1个专家，TOP1选择)，平衡了模型容量和训练稳定性

这种配置特别适合从大规模数据(如ImageNet-22K)中学习通用视觉特征，同时保持合理的计算成本。MoE机制的引入可以在不显著增加计算量的情况下提升模型容量。