跳到主要内容
返回时间线
arXiv来源发表:

一个Transformer块反复用:深度坐标调度专家库,B/16追平DeiT III而参数少约七成

导语

视觉Transformer的整个深度堆叠被压缩成一个反复执行的Transformer块,每个循环深度用归一化深度坐标把一个小专家库软合并成一个FFN,B/16在相近推理FLOPs下追平DeiT III而存储参数少约七成。

Source-provided article image: One Block, Multiple Depths: Recurrent Vision Transformers with Depth-Programmed Experts
arxiv.org

正文

视觉编码器的深度堆叠可以换成一个反复执行的Transformer块,深度相关的计算由每个循环深度上不同的FFN权重提供。 此前视觉Transformer按固定深度堆叠,每层有各自独立的参数;权重共享方案要么只做压缩,要么保留多个不同块。 在ImageNet-1k监督训练下,reViT-B/16达到83.0% top-1,DeiT III为82.8%,两者推理FLOPs接近,而reViT存储参数为23.6M对86.6M。

每个循环深度的FFN由一个小专家库的凸组合生成,组合系数只由该深度的归一化坐标决定,因此同一组系数对所有图像和所有token相同。 此前的专家混合在token层面做路由,每个token走不同的专家子集,专家参数在执行前并不合并。 在相同循环骨干、任务损失和基础训练配方下比较七种MoE机制,一个稠密FFN的名义算力预算下,深度编程合并优于所测试的token分发与输出混合方案。

一个训练好的检查点可以在多个推理深度上运行,只需在归一化坐标区间上重新采样。 固定深度模型改变深度需要重新训练或另存检查点。 在蒸馏设置下,深度条件变体在ADE20k、ImageNet线性探针和NYUv2上随深度增加而改善,而仅用特征条件的对照在ImageNet和ADE20k上几乎持平。

接下来

需要按固定深度部署的读者可以把合并后的FFN预先算好,写成常规稠密图,去掉在线路由与合并,代价是存储展开为每个深度一个FFN。想换推理深度的读者可以直接在归一化坐标区间上重新采样同一个检查点,无需重训。想扩大存储容量而不增加每步稠密FFN计算的读者可以增大专家库,S/16从4个专家到8个专家在ImageNet上提升2.0个点。

当前延迟测量使用直接PyTorch实现,合并权重用通用FP32归约重建,没有跨输入缓存或自定义内核,因此反映的是当前实现而非优化后的部署表现。动态形式在batch 1下总运行内存降低约一半,但带来延迟开销,batch 64时延迟开销升至约39%且总内存相近。蒸馏比较中reViT只匹配教师最后一层特征,而Raptor使用中间层特征,因此表2属于参考对比而非受控消融。专家删除分析显示,删除最早被使用的专家造成的对齐损失最大,因此删除敏感度本身无法把专家角色与其在循环中的位置分开。

来源