POA 一次性预训练是什么？一次训练同时得到多种尺寸的模型

一次性预训练（Pre-training Once，POA）是一种一次训练就能得到多种尺寸模型的自监督训练框架，由蚂蚁集团在 2024 年提出，论文标题是《POA：一次预训练得到所有尺寸的模型》，发表于 ECCV 2024。

问题：不同场景需要不同大小的模型

实际部署 AI 模型时，不同的场景对模型大小的要求很不一样：云端服务器算力充足，可以用大模型追求最好的效果；手机、摄像头等端侧设备算力和内存有限，只能用小模型。通常的做法是，为每一种尺寸单独训练一个模型，或者先训练一个大模型，再通过蒸馏、剪枝等方法压缩出小模型。无论哪种方式，每多一种尺寸，就要多付出一份训练成本。自监督预训练本身已经非常耗费计算资源，为每种尺寸都重复一遍，代价很高。

POA 怎么做

POA 采用三分支的自监督训练框架。它基于自蒸馏的思路，在常见的“教师-学生”两个分支之外，又引入了一个创新的弹性学生分支。在每一个预训练步骤中，弹性学生分支会从完整的学生网络中随机采样出一个子网络，子网络的宽度和深度都可以变化，比如只用一部分层、每层只用一部分通道，然后让这个子网络也参与训练，向教师网络学习。

因为每一步采样的子网络大小都不一样，经过整个预训练过程，从完整网络中“切”出的各种大小的子网络都得到了充分的训练。预训练结束后，就可以直接从这个模型中提取出不同尺寸的子网络使用，不需要额外的训练。论文称，一次预训练就能生成数百种不同尺寸的模型。研究者还发现，弹性学生分支的引入不仅没有拖累完整网络，反而通过子网络之间的相互促进，提升了整体的表示学习效果。

效果

论文在基于视觉 Transformer 和卷积网络的模型上进行了验证，结果显示，POA 一次预训练得到的各种尺寸的模型，在图像分类、目标检测、语义分割等下游任务上，表现与单独训练的同尺寸模型相当甚至更好，达到了当时最好的水平。这大大降低了为多种部署场景准备模型的成本，适用于需要在多种硬件上部署同类模型的企业。