在5月20日举办的WAVE SUMMIT 2022深度学习开发者峰会上，百度发布了行业最大视觉多任务文心VIMER-UFO 2.0大模型，模型参数量达到170亿，单模型28项公开数据集SOTA，基于飞桨Task MoE架构，根据任务的不同自动选择激活最优的区域，从而实现100倍参数压缩，同时支持下游任务快速扩展。
针对以上问题，百度文心大模型提出统一特征表示优化技术(UFO:Unified Feature Optimization)，在充分利用大数据和大模型的同时，兼顾大模型落地成本及部署效率。本次发布的文心VIMER-UFO 2.0大模型技术方案的主要内容包括：
All in One：发布行业最大的 170 亿参数视觉多任务模型，覆盖人脸、人体、车辆、商品、食物等细粒度分类的20+CV基础任务，单模型28项公开测试集SOTA。
One for All：首创针对视觉多任务的超网络预训练方案，支持各类任务、各类硬件的灵活部署，解决大模型应用时模型参数量大，推理性能差等问题。
Task MoE: 飞桨多任务超网络分布式训练架构，支持训练任务动态扩展，特定任务任意切分，保证多任务之间信息有效借鉴，负载均衡，高效协同。
1
All-in-One
功能更强大、更通用的视觉模型
之前主流的视觉模型生产流程，通常采用单任务 “train from scratch” 方案。每个任务都从零开始训练，各个任务之间也无法相互借鉴。由于单任务数据不足带来偏置问题，实际效果过分依赖任务数据分布，场景泛化效果往往不佳。近两年随着预训练技术的蓬勃发展，基于海量数据训练获得的预训练模型展现出强大的泛化能力，在下游任务通过少量数据finetune即可获得良好的效果。但美中不足的是，基于预训练+下游任务finetune的模型生产流程，需要针对各个任务分别训练模型，存在较大的研发资源消耗。
百度文心大模型提出的All in One训练方案，通过多任务协同学习方式训练得到一个功能强大的通用模型，可被直接应用于处理多个任务。在通过跨任务的信息提升单个任务效果的同时，也免去了下游任务finetune过程，可被广泛应用于各类多任务AI系统。以智慧城市场景为例，All in One可以用单模型实现多个任务的SOTA识别效果，同时多任务模型可获得显著优于单任务模型的效果，证明了多任务之间信息借鉴机制的有效性。
基于All in One训练方案，文心VIMER-UFO 2.0大模型参数量达到了170亿，位居行业首位。单模型在不进行下游finetune的情况下，在28项公开数据集上取得了 SOTA的结果。
2
One-for-All
灵活可伸缩的弹性部署方案
在大模型落地时，受限于硬件算力、存储等的限制，往往无法直接部署。针对该问题，文心VIMER-UFO One for All方案通过引入“超网络“进行解决。超网络由众多稀疏的子网络构成，每个子网络是超网络中的一条路径，将不同参数量、不同任务功能、不同精度的模型训练整合为一个超网络的训练。训练完成的大模型即可针对不同的任务和设备低成本生成相应的可即插即用的小模型，实现One for All Tasks和One for All Chips的能力。
文心VIMER-UFO 2.0基于Vision Transformer结构设计了多任务多路径超网络。与谷歌Switch Transformer以图像为粒度选择路径不同，其以任务为粒度进行路径选择，超网络训练完成后，可根据不同任务独立抽取对应的子网络进行部署，无需部署整个大模型。文心VIMER-UFO 2.0的超网中不同的路径除了可以选择不同FFN单元，Attention模块和FFN模块内部也支持弹性伸缩，实现网络的搜索空间扩展，为硬件部署提供更多可选的子网络，并提升精度。
文心VIMER-UFO 2.0大模型基于飞桨的Task MoE架构构建多任务超网络。其主要由Attention模块和FFN模块构成，不同任务子模型会共享Attention模块。在FFN模块中 ，每个任务都有两种不同的路径选择，即选择共享 FFN（FFN-shared）或者专属FFN(FFN-taskX)；在应用时，文心VIMER-UFO 2.0提出的Path Routing方法会根据任务的不同自动选择激活最优的区域（类似于人类的大脑，人类的大脑经过数百万年的进化，形成了分区的结构，不同区域负责特定功能，同时又是相互协作的一个整体）。在实际部署时，既可以部署整个大模型实现多任务处理能力，还可以从大模型中直接抽取单个和少数几个任务模型部署，可大幅降低模型的参数量，同时不同任务之间可自由组合，显著提升了AI应用的研发效率。
One For All Chips
基于文心VIMER-UFO 2.0的Stochastic Architecture Slimming方法，大模型在抽取的任务子网基础上进一步构建了芯片超网络。对于Attention模块，每个芯片子网可以选择不同的 Head 数量、 QKV 矩阵参数量。对于FFN模块，可根据放缩系数弹性选择FFN中参数规模。在应用落地时，针对不同平台存储容量和算力不同，可以抽取不同深度和宽度的子网络进行部署，进一步压缩模型的参数和计算量。
多任务大模型直接部署
单任务抽取子网络部署
针对只需要个别处理能力的AI应用，可根据任务需求从文心VIMER-UFO 2.0大模型中抽取部分参数，得到针对特定任务的模型进行部署，大幅减少模型的参数量。例如文心VIMER-UFO 2.0具备170亿参数规模，而抽取的单任务模型只包含6亿参数。基于单任务模型抽取的芯片级模型参数量可进一步降低到1亿规模，模型参数量压缩超过100倍。并且不同任务之间可自由组合，大大提升了AI服务的开发和部署效率。
新任务快速扩展
针对未覆盖的新任务，文心VIMER-UFO 2.0支持在只更新部分参数的情况下，仅使用少量数据finetune，实现任务的快速扩展。文心VIMER-UFO 2.0的超网络中有一个Shared的分支（Attention 与 FFN-Shared），该分支在文心VIMER-UFO 2.0大模型的训练过程中使用全部任务数据进行优化，因此具备了强大的任务泛化性，对于不支持的新任务，只需要抽取该分支的参数使用少量数据进行finetune，便可在新任务上达到良好的效果。同时由于只需要更新部分参数，下游finetune的成本大大降低，解决了目前主流大模型落地应用的难题。新任务扩展结果（不使用外部数据）如下图所示。
子网络异构蒸馏
为了更好地支持在移动和边缘设备上进行部署，文心VIMER-UFO 2.0还可抽取子网络模型进行模型蒸馏。结合百度研发的异构蒸馏技术，可实现Transformer架构模型到CNN架构模型的高效知识迁移，模型参数量从亿级别进一步降低到百万级别。
3
飞桨Task MoE分布式训练架构
如此大的参数规模和任务数，给模型的训练带来了巨大的挑战。文心VIMER-UFO 2.0大模型采用稀疏门控混合专家设计，仅参数存储就需要68G，给训练时的模型存储带来了压力；该模型在前向反向时所有计算节点间会进行同步等待的All-to-All通信，使得通信负担明显加大；此外，该模型的多任务数目是动态的，且多个任务之间样本严重不均衡，使得计算节点之间的同步等待较长，影响并发效率。
针对这些挑战，飞桨提出了Task MoE分布式训练架构，不仅实现多级并行存储稀疏参数，还支持硬件拓扑感知通信，使得层次化All-to-All通信效率提升20%。同时飞桨还创新性地提出了基于Task的负载均衡机制，支持任务数量的动态扩展、特定任务的任意切分以及多个任务在不同的专家下的并发训练，同等实验环境下训练性能比PyTorch提升66%。同时，该方案保障多任务之间信息借鉴机制的有效性，使得VIMER-UFO 2.0模型精度大幅提升。此外，在推理阶段，基于飞桨Task MoE架构构建的多任务多路径的超网络，可支持任务粒度的路径选择，方便灵活部署。
更多详情请查看：
SE-MoE: A Scalable and Efficient Mixture-of-Experts Distributed Training and Inference System（https://arxiv.org/abs/2205.10034）
4
结语
百度文心大模型提出统一特征表示优化技术(UFO：Unified Feature Optimization)的初衷，在于为大模型应用落地探索解决方案，本次发布的文心VIMER-UFO 2.0大模型可被广泛应用于智慧城市、无人驾驶、工业生产等各类多任务AI系统，支持多种应用模式配合，兼顾效率和效果，近期亦会在百度零门槛AI开发平台EasyDL中开放，敬请期待。
点击阅读原文获得链接
了解更多UFO技术细节
https://github.com/PaddlePaddle/VIMER/tree/main/UFO/OneForAll
相关阅读
关注【飞桨PaddlePaddle】公众号
获取更多技术内容~
