通用视觉预训练大模型巡礼系列（一）UFO大模型

牛股行天下烁

前言：
近来，以ChatGPT为代表的自然语言对话大模型，以及以Stable Diffusion为代表的多模态文图生成大模型展现出惊艳的效果。而面向视觉领域的通用视觉模型相对来说似乎略显冷清。事实上，从2021年商汤联合上海人工智能实验室发布“书生”通用视觉体系以及在2022年世界人工智能大会上发布的“书生2.0”版本、2021年华为发布盘古CV大模型、2022年百度发布UFO通用视觉大模型、智源研究院开源EVA视觉大模型等等可以看出，学术界和工业界一直持续关注预训练大模型技术在视觉领域的发展和应用落地。
结合自己的研究兴趣，笔者一直以来比较关注预训练大模型技术在视觉领域的技术进展。人工智能技术在面向视觉的工业应用落地场景非常广阔。尽管在人脸、车辆识别等场景中，利用传统AI技术已经能够解决单个场景的业务需求。但是这种“手工作坊式”的开发模式，导致模型开发成本随着场景增长而提高。大量长尾场景中数据获取代价高昂，AI技术落地困难。随着预训练大模型技术的发展，这种AI落地困境可能得到解决。
通用视觉预训练大模型就是要解决AI技术在视觉任务场景中的落地难题，为AI模型降低对场景数据依赖，支持下游多任务通过小样本微调达到高性能。尽管通用预训练视觉大模型要解决的问题是一致的，但是各家视觉大模型的技术路线不尽相同。其中有很多值得研究的地方，希望这些研究也会给通用视觉预训练大模型技术的发展、乃至未来多模态大模型中视觉模态的融入方法带来启发。
本次主要分享百度UFO视觉大模型相关技术。主要研究内容是基于ECCV 2022的论文《UFO: Unified Feature Optimization》[1]和百度UFO模型的开源代码：https://github.com/PaddlePaddle/VIMER/blob/main/UFO/


UFO大模型整体概述：
针对预训练大模型落地所面临的问题，百度提出统一特征表示优化技术（UFO：Unified Feature Optimization），在充分利用大数据和大模型的同时，兼顾大模型落地成本及部署效率。百度UFO大模型是行业最大 170 亿参数视觉多任务模型，覆盖人脸、人体、车辆、商品、食物细粒度分类等 20+ CV 基础任务，单模型 28 个公开测试集效果 SOTA。


UFO在28个公开测试集上达到SOTA(分类任务)
百度UFO大模型的整体构建、部署可以看成三个过程。

• 多数据集训练一个通用网络（百度称之为AllInOne）。
  
• 通用网络通过多任务训练扩展成为由共享网络和专家网络构成的170亿参数的稀疏大模型，称之为超网络（OneForAll）。
  
• 从完整参数的稀疏大模型中抽取子网络用于下游任务。
  

UFO模型的整体构建过程
模型结构：


UFO超网模型结构
传统的多路径大模型的参数共享方式主要有软参数共享和硬参数共享两种。软参数共享方式允许每个任务有单独的模型和参数，但是强制每个模型可以通过正则化访问其他模型内部的信息[2,3]。硬参数共享则是共享部分网络结构[4,5]。UFO综合了两种参数共享方式，基于 Vision Transformer 结构设计了多任务多路径超网络。与谷歌 Switch Transformer 以图片为粒度选择路径不同，UFO以任务为粒度进行路径选择，这样当超网络训练好以后，可以根据不同任务独立抽取对应的子网络进行部署，而不用部署整个大模型。UFO的超网中不同的路径除了可以选择不同 FFN 单元，Attention 模块和 FFN 模块内部也支持弹性伸缩，实现网络的搜索空间扩展。超网络分为多路径 FFN 超网和与可伸缩 Attention 超网两部分。首先针对多路径 FFN 超网模块，每个任务都有两种不同的路径选择，即选择共享 FFN（FFN-shared）或者专属 FFN（FFN-task），当选定好 FFN 以后，还可根据放缩系数弹性选择FFN中参数规模；因此FFN超网络中共有（T * ratio）^L 种不同的 FFN 路径，其中 T 为 task 的数量，L 为网络的层数, ratio 为放缩系数的数量。而对于 self-attention 超网，每个子网络可以选择不同的 Head 数量 QKV 矩阵参数量。
模型训练与子网参数抽取策略：


UFO超网与子网抽取
下面结合论文分析UFO的超网训练策略和下游任务子网络抽取方法。
UFO超网是基于Visual Transformer (ViT)。子网会从超网中选择部分模块、继承对应参数。传统基于transformer的网络结构搜索（NAS）方法主要是三种搜索方向：弹性深度、弹性注意力头、弹性前馈网络（FFN）比例。UFO引入了一种弹性FFN路径的网络结构搜索方向。UFO中的子网最终可以弹性地从超网中抽取FFN路径、FFN权重、注意力权重或者是是否ViT的完整子模块。
现有的其他超网训练策略中，因为没有没有FFN路径参与，超网模型只能一步一步训练。UFO中加入了FFN路径来减轻共享注意力之间严重的参数竞争，并且可以实现端到端的训练。UFO中可以为每个下游任务的FFN网络选择提供三种选项：共享FFN，任务专属FFN或者两者都选择。其中对下游任务的概率分布采样使用了Gumbel-softmax重参数化手段。
在多任务联合训练时，UFO根据每个对应任务数据集的大小，按照比例采样部分数据，形成一个batch，输入到超网中的主干（backbone）进行训练。主干网络获得的特征被分隔送入到任务专属的Head网络中。任务专属网络计算后可以得到每个任务的误差。最后将这些误差求和后，一起送回共享主干网络中，误差值在网络中反传，这样就完成了一个后向梯度传播的过程。
如何根据下游任务的应用特性从超网中抽取对应任务最优子网络是一个复杂的问题。在子网络的参数选择过程中，不仅考虑子网络在目标任务上的性能，还考虑了抽取的子网在所有任务上的平均性能。因为UFO超网模型的完整参数量达到了170亿，对每种网络结构都进行搜索计算是不现实的。研究者通过GP-NAS[6]方法从超网中先生成一个包含多个网络结构的子集。然后训练了一个“子网任务性能预测器”来预测子集中的网络结构的下游任务性能，最终从中选择一个作为最优子网结构。
开源代码分析&实验：
UFO模型的开源代码结构如图：


UFO开源代码结构
如果想进行下游任务测试，主要的功能代码都在OneForAll路径下。


UFO下游任务相关代码
目前UFO的开源版本只支持分类任务的微调。笔者也对UFO模型进行了分类任务的初步测试。相同小样本数据量情况下，UFO的性能确实可以超过一般常用的预训练分类网络Resnet 50。
总结：
UFO通用视觉预训练大模型是国内较早开源的面向视觉的大模型。通过上文的分析，UFO大模型的主要技术路线是MoE（Mix of Expert）[7]这种稀疏大模型技术。UFO创新性地提出了以任务为粒度进行大模型的子网抽取方法。目前，在百度开发者大会上，UFO模型已经跟很多具体需求方联合发布了多个行业大模型，具有一定的应用落地能力[8]。不过开源版本功能较少（只支持分类任务微调），且大模型的实验成本比较高（算力要求高）。
参考文献
[1]XI T, SUN Y, YU D, et al. UFO: Unified Feature Optimization[Z]. XI T, SUN Y, YU D, et al. UFO: Unified Feature Optimization[Z]. arXiv, 2022(2022–07–21). DOI:10.48550/arXiv.2207.10341.
[2]GAO Y, BAI H, JIE Z, et al. MTL-NAS: Task-Agnostic Neural Architecture Search Towards General-Purpose Multi-Task Learning[C]//2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)GAO Y, BAI H, JIE Z, et al. MTL-NAS: Task-Agnostic Neural Architecture Search Towards General-Purpose Multi-Task Learning[C]//2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). . DOI:10.1109/CVPR42600.2020.01156.
[3]LIU P, QIU X, HUANG X. Recurrent Neural Network for Text Classification with Multi-Task Learning[C]//Proceedings of the Twenty-Fifth International Joint Conference on Artificial Intelligence. New York, New York, USA: AAAI Press, 2016: 2873–2879.LIU P, QIU X, HUANG X. Recurrent Neural Network for Text Classification with Multi-Task Learning[C]//Proceedings of the Twenty-Fifth International Joint Conference on Artificial Intelligence. New York, New York, USA: AAAI Press, 2016: 2873–2879.
[4]DUONG L, COHN T, BIRD S, et al. Low Resource Dependency Parsing: Cross-Lingual Parameter Sharing in a Neural Network Parser[C]//Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 2: Short Papers). Beijing, China: Association for Computational Linguistics, 2015: 845–850. DOI:10.3115/v1/P15-2139.DUONG L, COHN T, BIRD S, et al. Low Resource Dependency Parsing: Cross-Lingual Parameter Sharing in a Neural Network Parser[C]//Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 2: Short Papers). Beijing, China: Association for Computational Linguistics, 2015: 845–850. DOI:10.3115/v1/P15-2139.
[5]MISRA I, SHRIVASTAVA A, GUPTA A, et al. Cross-Stitch Networks for Multi-Task MISRA I, SHRIVASTAVA A, GUPTA A, et al. Cross-Stitch Networks for Multi-Task Learning[C]//2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR). . DOI:10.1109/CVPR.2016.433.
[6]LI Z, XI T, DENG J, et al. GP-NAS: Gaussian Process Based Neural Architecture Search[C]//2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)LI Z, XI T, DENG J, et al. GP-NAS: Gaussian Process Based Neural Architecture Search[C]//2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). . DOI:10.1109/CVPR42600.2020.01195.
[7]KUDUGUNTA S, HUANG Y, BAPNA A, et al. Beyond Distillation: Task-Level Mixture-of-Experts for Efficient Inference[Z]. KUDUGUNTA S, HUANG Y, BAPNA A, et al. Beyond Distillation: Task-Level Mixture-of-Experts for Efficient Inference[Z]. arXiv, 2021(2021–09–24). DOI:10.48550/arXiv.2110.03742.
[8]https://wenxin.baidu.com/wenxin/news/2. 文心大模型-产业级知识增强大模型[EB].