首页 星云 工具 资源 星选 资讯 热门工具
:

PDF转图片 完全免费 小红书视频下载 无水印 抖音视频下载 无水印 数字星空

POA:已开源,蚂蚁集团提出同时预训练多种尺寸网络的自监督范式 | ECCV 2024

编程知识
2024年08月30日 12:34

论文提出一种新颖的POA自监督学习范式,通过弹性分支设计允许同时对多种尺寸的模型进行预训练。POA可以直接从预训练teacher生成不同尺寸的模型,并且这些模型可以直接用于下游任务而无需额外的预训练。这个优势显著提高了部署灵活性,并有助于预训练的模型在各种视觉任务中取得SOTA结果。

来源:晓飞的算法工程笔记 公众号

论文: POA: Pre-training Once for Models of All Sizes

Abstract


  大规模自监督预训练为一个基础模型处理多种不同的视觉任务铺平了道路。大多数预训练方法在一次训练中训练一个特定大小的单一模型。然而,在现实世界的场景中,由于各种计算或存储限制,需要大量的工作来开发一系列不同大小的模型进行部署。因此,在这项研究中,我们提出了一种新颖的三分支自监督训练框架,称为POAPre-training Once for All),来解决上述问题。我们的方法在现代自蒸馏范式中引入了一种创新的弹性student分支。在每个预训练步骤中,我们随机从原始student中抽样一个子网络来形成弹性student,并以自蒸馏的方式训练所有分支。一旦预训练完成,POA允许提取不同大小的预训练模型用于下游任务。值得注意的是,弹性student促进了多个不同大小模型的同时预训练,同时也作为各种大小模型的额外集合,增强了表示学习。大量实验证明了我们的POA的有效性和优势,包括k最近邻、线性探测评估以及多个下游任务的评估。它使用ViTSwin TransformerResNet骨干网络实现了最先进的性能,并通过一次预训练会话生成了大约一百个不同大小的模型。代码可在以下链接找到:https😕/github.com/Qichuzyy/POA

Introduction


  通过自监督学习在大型模型中学习可泛化的视觉表示,近年来在各种视觉任务上取得了卓越的性能。然而,当部署到现实世界的应用程序时,大型模型必须根据计算、存储、功耗等各种资源限制进行调整。例如,一个良好设计的人工智能产品通常包括一套为不同场景量身定制的模型,比如Gemini NanoProUltra。对于一个大型预训练模型,将其部署到具有不同资源约束的多个应用场景的常见解决方案包括额外的权重修剪、知识蒸馏,甚至从头开始重新训练一个小网络,这些都需要大量的开发工作。因此,这个问题引发了一个关键问题:是否可能进行一次预训练以同时生成多个具有不同大小的模型,每个模型都提供足够好的表示。

  为了解决这一挑战,论文引入了一种名为POAPre-training Once for All)的新型自监督学习范式。POA建立在流行的teacher-student自蒸馏框架之上,具有一个额外的创新性弹性student分支。弹性student分支通过参数共享嵌入了一系列子网络,这是基于观察到对于现代网络结构来说,较小尺寸的模型是较大尺寸模型的子网络。此外,该分支的参数与原始的或完整的studennt共享。在每个预训练步骤中,从完整student中随机抽样一部分参数,形成相应的弹性studennt。原始完整student和弹性student都被训练以模拟teacher网络的输出。teacher本身通过对student参数的指数移动平均(EMA)不断优化,包括采样的弹性student。弹性student有助于在不同参数子集上进行有效和高效的预训练,从而成功地从预训练teacher中提取出高性能子网络,用于后续的下游场景。它还作为一种训练正则化形式,通过强制teacher和各种子网络之间的输出匹配来促进稳定的训练过程。

POA代表了第一个能够同时训练多个不同尺寸模型的自监督学习方法,每个模型在不需要进一步预训练的情况下,都能获得适用于不同资源约束的高质量表示。图1显示了通过POA预训练的ViT-L模型提取的143个子网络的k最近邻(k-NN)评估结果。通过选择不同的弹性宽度和深度,预训练teacher模型可以根据可用计算资源定制的适用于下游应用的合适模型,生成足够数量的候选子网络以供选择。值得注意的是,由于在同视图蒸馏上进行了精心设计,每个子网络都经过了良好训练,并表现出优越性能。特别是,ViT-SViT-BViT-L模型创造了新的基准,与那些由现有方法预训练的模型相比取得了SOTA结果。

  为了严格评估方法的有效性,使用三种广泛使用的骨干架构,即ViTSwin TransformerResNet,进行了大量实验。每个骨干架构都在ImageNet-1K数据集上进行了预训练,并使用k-NN和线性探测分类评估,以及在下游密集预测任务进行评估,如目标检测和语义分割。POA在单次预训练会话中跨多种模型尺寸实现了最先进的准确性。

  本文的技术贡献总结如下:

  1. POA是第一个将无监督表示学习和一次性模型生成集成到单个预训练会话中的预训练范式,解决了社区很少探讨的一次性预训练挑战。这对实际部署非常重要,因为实际部署通常需要一套模型。
  2. 提出了一个新颖而优雅的组件,称为弹性studentElastic Student),具有一系列弹性算子,可以使POA与包括ViTSwin TransformerResNet在内的流行骨干结构兼容,具备生成各种大小模型的能力。此外,还作为模型集成来平滑训练过程并改善学到的表示。
  3. 通过对k-NN、线性探测和下游密集任务评估的彻底评估,在多个指标上展现出优于现有最先进预训练方法的性能。此外,将POA与自监督蒸馏(SEED)进行了比较,SEED是一种专为自监督学习设计的知识蒸馏方法,进一步验证了POA的有效性。

POA Self-supervised Learning Framework


  论文的主要目标是通过单次自监督预训练会话来预训练多种规模的模型,受到自蒸馏技术最新进展的启发,提出了一个名为POA的新型SSLSelf-supervised Learning)框架。POA架构如图2所示,包括一个teacher模型、一个完整的student模型、一个弹性student模型以及对应的头部。teacher模型使用student模型的指数移动平均(EMA)进行更新。弹性student模型是完整student模型的派生版本,其主干网络和头部参数是共享的。

  在两个方面利用蒸馏技术:完整student和弹性student都是通过使用同一图像不同视图的teacher模型进行蒸馏,而弹性student还通过使用相同视图的完整student进行学习。交叉视图蒸馏作为一种表示学习形式,如所介绍的那样。值得注意的是,除了仅使用完整student进行常规EMA更新外,弹性student在每个预训练步骤中还提供一个随机抽样的子网络,参与teacher模型的EMA优化。这个过程实际上模拟了多个子网络的集成,这在监督学习领域也被证明是有益的。同视图蒸馏是完整student和弹性student之间的标准知识蒸馏,提升了弹性student的质量。

Design of Elastic Student

  弹性student是一个子网络,其参数是从完整student中提取的。在transformer主干网络的背景下,宽度指的是标记的维度,而在卷积主干网络中,宽度表示通道数。深度则定义为transformer或卷积网络中基本块的数量。给定宽度和深度的值,会产生一定的网络结构。为简单起见,论文将重点放介绍ViT的弹性设计。

ViT的基本块主要由多头自注意力(MSA)模块和多层感知器(MLP)模块组成。在每个模块之前应用层归一化(LN),并在每个模块后使用残差连接。如图3的左侧所示,弹性块是指在ViT原始基本块中调整宽度后堆叠的弹性MSAMLPLN。在论文的方法中,弹性student分支是通过在每个训练迭代中组装特定数量的这些弹性块来构建的。

  • Elastic MSA

  一个原始或完整的MSA模块由三个主要组件组成,即输入投影层,包含注意力和连接的操作符,以及输出投影层。将投影层定义为( \(w^{\ast}, b^{\ast}\) ),其中 \(w^{\ast}\) 表示线性转换权重, \(b^{\ast}\) 表示相应的偏置, \(\ast\) 表示层的名称。如图3的右侧所示,给定一个标记维度 \(D_{max}=N_h \cdot D_h\) ,其中 \(N_h\) 是注意力头的数量, \(D_h\) 是头部维度,具有长度 \(T\) 的输入序列 \(z \in \mathbb{R}^{T \times D_{max}}\) 最初被投影以形成查询 \(Q \in \mathbb{R}^{T \times D_h}\) 、键 \(K \in \mathbb{R}^{T \times D_h}\) 和值 \(V \in \mathbb{R}^{T \times D_h}\) 。为了生成弹性MSA,定义了M+1个弹性宽度,包括 \(D_{max}\) ,间隔为 \(D_h\)

\[\begin{equation} D_i = (N_h - i) \cdot D_h,\quad\forall i \in \{0, 1, ..., M\},\quad M < N_h. \label{eq:elastic_width} \end{equation} \]

  对于每个弹性宽度 \(D_i\) ,从完整MSA中的相应输入投影层( \(w^{a1}\) , \(b^{a1}\) )中提取生成每个头部的 \(Q\)\(K\)\(V\) 的权重 \(w^{a1}_i \in \mathbb{R}^{D_h \times D_i}\) 和偏置 \(b^{a1}_i \in \mathbb{R}^{D_h}\) ,如 \(w^{a1}_i = w^{a1}[:, :D_i]\cdot\alpha_i\)\(b^{a1}_i = b^{a1}\) 。这里, \(\alpha_i\) 表示用于应对输入维度减少的缩放因子,计算公式为 \(\alpha_i = D_{max}/D_i\) 。随着宽度的减小,弹性MSA中的注意力头数量自然减少到 \(N_h - i\) 。类似地,对于输出投影层( \(w^{a2}\) , \(b^{a2}\) ),权重 \(w^{a2}_i \in \mathbb{R}^{D_i \times D_i}\) 和偏置 \(b^{a2}_i \in \mathbb{R}^{D_i}\) 被提取为:

\[\begin{equation} w^{a2}_i = w^{a2}[:D_i, :D_i]\cdot\alpha_i ~~~~~ b^{a2}_i = b^{a2}[:D_i]. \label{eq:w_b_extract} \end{equation} \]

  • Elastic MLP

ViT块中的原始或完整MLP模块包含两个投影层。第一层( \(w^{m1}, b^{m1}\) )将嵌入维度扩展了 \(s\) 倍,通常在ViT结构中设置为4。然后,第二层( \(w^{m2}, b^{m2}\) )将其投影回原始维度。弹性MLP的两个层的参数以类似于公式2描述的方式提取,如下所示:

\[\begin{equation} \begin{aligned} & w^{m1}_i = w^{m1}[:D_i \cdot s, :D_i]\cdot\alpha_i ~~~~~ b^{m1}_i = b^{m1}[:D_i \cdot s] \\ & w^{m2}_i = w^{m2}[:D_i, :D_i \cdot s]\cdot\alpha_i ~~~~~ b^{m2}_i = b^{m2}[:D_i]. \end{aligned} \label{eq:mlp_extract} \end{equation} \]

  • Elastic LN

  对于弹性LN,直接使用原始LN内部参数的前 \(D_i\) 个元素,类似于公式2中的偏置提取。

  • Elastic depth

  要从包含 \(L_{max}\) 个块的完整ViT中创建一个包含 \(L_i\) 个弹性块的子网络,引入了一组N+1个弹性深度,定义为 \(L_i = L_{max} - i,~~\forall i \in \{0, 1, ..., N\},~~N < L_{max}\) 。对于特定深度 \(L_i\) ,根据块ID在等间隔上选择相应的块。激活深度 \(L_i\) 的每个块ID \(BID^{L_i}_j\) 可以表示为:

\[\begin{equation} BID^{L_i}_j = \left\lfloor \frac{(L_{max} - 1) \cdot j}{L_i - 1} \right\rfloor,\quad \forall j \in \{0, 1, ..., L_i - 1\}. \label{eq:block_id_extraction} \end{equation} \]

  因此,通过结合弹性宽度和深度,可以生成总共 \((N+1)\cdot(M+1)\) 个不同的子网络。例如,通过将弹性宽度设置为384,弹性深度设置为12,可以直接从如ViT-L的完整网络中提取一个ViT-S。在预训练的每次迭代中,随机选择其中一个子网络作为弹性student分支。

Distillation between Views

POA根据其三个分支执行蒸馏。给定输入图像 \(x\) 的一对全局增强视图,表示为 \(x_a\)\(x_b\)teacher编码器 \(E_{T}\) 使用 \(x_a\) 作为输入提取特征 \(Z_a = E_{T}(x_a)\) 。同时, \(x_b\) 被输入到完整student编码器 \(E_{IS}\) 和弹性student编码器 \(E_{ES}\) 中,分别产生特征 \(Z_{b1} = E_{IS}(x_b)\)\(Z_{b2} = E_{ES}(x_b)\) 。从teacher编码器输出的特征 \(Z_a\) 经过teacher头部 \(H_T\) 处理,然后使用Sinkhorn-KnoppSK)算法进行居中处理,并使用温度缩放softmax进行归一化,生成概率 \(p_a\) ,如下所示:

\[\begin{equation} \begin{aligned} l_a = SK(H_T(Z_a)),~l_a \in \mathbb{R}^P ~~\quad p^i_a = \frac{\exp(l^i_a / \tau)}{\sum^{P-1}_{k=0}\exp(l^k_a / \tau)},~\forall i \in \{0, ..., P-1\}, \end{aligned} \label{eq:prob_teacher} \end{equation} \]

  其中 \(P\) 是原型(logits?)的数量, \(\tau > 0\) 是温度参数。类似地,通过使用student头部 \(H_{IS}\)\(H_{ES}\) 处理输出来计算完整和弹性student编码器的概率 \(p^i_{b1}\)\(p^i_{b2}\) 。然后,这些输出通过一个针对student量身定制的温度参数 \(\tau'\) 的温度缩放softmax函数进行处理。值得注意的是, \(H_{IS}\)\(H_{ES}\) 共享相同的参数,只是 \(H_{ES}\) 的第一个投影层进行公式2的相应调整,以便对齐相应的维度。为简单起见,省略了 \(p^i_{b1}\)\(p^i_{b2}\) 的显式表达式,因为它们遵循与公式5类似的计算方式。对于完整student分支,使用跨视图数据从teacher进行蒸馏如下:

\[\begin{equation} \mathcal{L}^g_{IS} = -p_a \log(p_{b1}). \label{eq:distill_is} \end{equation} \]

  弹性student分支在POA框架中发挥着至关重要的作用。为了确保这一分支的充分训练,采用了从teacher和完整student分支进行的双重蒸馏。第一次蒸馏涉及到teacher模型,利用跨视图数据来引导表示学习。第二次是与完整student模型进行的蒸馏过程,使用同视图数据。这种同视图蒸馏负责将完整student学到的表示转移到弹性student分支。这种双重蒸馏过程的损失函数制定如下

\[\begin{equation} \begin{aligned} \mathcal{L}^g_{ES1} = - p_a \log(p_{b2}), \quad \mathcal{L}^g_{ES2} = - p_{b1} \log(p_{b2}). \label{eq:distill_es} \end{aligned} \end{equation} \]

  请注意,在这两个损失函数中,对所有原型求和,以计算相应概率分布之间的交叉熵损失。

Overall Loss of POA

  根据SSL方法,采用多裁剪策略从单个图像中创建各种失真视图。除了之前提到的两个全局视图外,还生成 \(v\) 个分辨率较低的局部视图 \(x_{l_1}, x_{l_2}, ..., x_{l_v}\) 。这些局部视图由两个student共同处理,以促进局部到全局的对应关系。完整和弹性student的局部蒸馏损失计算如下:

\[\begin{equation} \mathcal{L}^{l}_{IS} = - \frac{1}{v}\sum^{v}_{i=1} p_a \log(p_{l_i1}), \label{eq:local_distill_is_loss} \end{equation} \]

\[\begin{equation} \mathcal{L}^{l}_{ES1} = - \frac{1}{v}\sum^{v}_{i=1}p_a \log(p_{l_i2}), \quad \mathcal{L}^{l}_{ES2} = - \frac{1}{v}\sum^{v}_{i=1}p_{l_i1} \log(p_{l_i2}), \label{eq:local_distill_es_loss} \end{equation} \]

  其中, \(p_{l_{i1}}\)\(p_{l_{i2}}\) 分别是完整和弹性student分支对于局部视图 \(l_i\) 产生的概率。完整和弹性student的总蒸馏损失通过将它们与因子 \(\lambda\) 相加来计算:

\[\begin{equation} \begin{aligned} \mathcal{L_S} &= \lambda(\mathcal{L}^g_{IS} + \mathcal{L}^{l}_{IS}) + (1-\lambda)((\mathcal{L}^g_{ES1} + \mathcal{L}^{l}_{ES1}) + (\mathcal{L}^g_{ES2} + \mathcal{L}^{l}_{ES2})) \\ &= \lambda\mathcal{L}_{IS} + (1-\lambda)(\mathcal{L}_{ES1} + \mathcal{L}_{ES2}). \end{aligned} \label{eq:total_distill_loss} \end{equation} \]

  为了确保弹性student的每个子网络都得到充分的训练,在主干网络之后引入了多个投影头(MPH)。每个投影头具有完全相同的结构,只是原型数量不同。对于每个投影头,根据公式10计算完整和弹性student的蒸馏损失 \(\mathcal{L_S}_i\) 。最终,在具有 \(H\) 个投影头的POA框架中,整体损失函数被表述为: \(\mathcal{L} = \frac{1}{H} \sum^H_{i=1}\mathcal{L_S}_i\)

Experiments




如果本文对你有帮助,麻烦点个赞或在看呗~
更多内容请关注 微信公众号【晓飞的算法工程笔记】

work-life balance.

From:https://www.cnblogs.com/VincentLee/p/18388586
本文地址: http://www.shuzixingkong.net/article/1581
0评论
提交 加载更多评论
其他文章 投屏协议
AirPlay协议是苹果开发、广泛应用于iPhone、iPad和Mac设备,可以通过WiFi将iPhone、iPad等iOS设备上的图片、音频、视频通过无线的方式传输到支持AirPlay 设备。即移动终端显示什么电视大屏就显示什么。随着AirPlay协议逐步普及,国内越来越多网络机顶盒,智能电视都集
服务重启了,如何保证线程池中的数据不丢失?
大家好,我是苏三,又跟大家见面了。 前言 最近有位小伙伴在我的技术群里,问了我一个问题:服务down机了,线程池中如何保证不丢失数据? 这个问题挺有意思的,今天通过这篇文章,拿出来跟大家一起探讨一下。 1 什么是线程池? 之前没有线程池的时候,我们在代码中,创建一个线程有两种方式: 继承Thread
服务重启了,如何保证线程池中的数据不丢失? 服务重启了,如何保证线程池中的数据不丢失? 服务重启了,如何保证线程池中的数据不丢失?
工作 6 年,@Transactional 注解用的一塌糊涂
接手新项目一言难尽,别的不说单单就一个 @Transactional 注解用的一塌糊涂,五花八门的用法,很大部分还失效无法回滚。 有意识的在涉及事务相关方法上加@Transactional注解,是个好习惯。不过,很多同学只是下意识地添加这个注解,一旦功能正常运行,很少有人会深入验证异常情况下事务是否
工作 6 年,@Transactional 注解用的一塌糊涂
好多kafka难题啊,看看其中的化解之道
前面已经分享过几篇面试了,这是一篇关于更加面向项目和技术的面经详解,第一次遇见问那么多kafka的问题,看看这个粉丝是怎么回答的。
[WPF]数据绑定时为何会出现StringFormat失效
在数据绑定过程中,我们经常会使用StringFormat对要显示的数据进行格式化,以便获得更为直观的展示效果,但在某些情况下格式化操作并未生效,例如 Button的 Content属性以及ToolTip属性绑定数据进行StringFormat时是无效的。首先回顾一下StringFormat的基本用法
[WPF]数据绑定时为何会出现StringFormat失效 [WPF]数据绑定时为何会出现StringFormat失效
折腾 Quickwit,Rust 编写的分布式搜索引擎 - 可观测性之分布式追踪
概述 分布式追踪是一种跟踪应用程序请求流经不同服务(如前端、后端、数据库等)的过程。它是一个强大的工具,可以帮助您了解应用程序的工作原理并调试性能问题。 Quickwit 是一个用于索引和搜索非结构化数据的云原生引擎,这使其非常适合用作追踪数据的后端。 此外,Quickwit 本地支持 OpenTe
折腾 Quickwit,Rust 编写的分布式搜索引擎 - 可观测性之分布式追踪 折腾 Quickwit,Rust 编写的分布式搜索引擎 - 可观测性之分布式追踪 折腾 Quickwit,Rust 编写的分布式搜索引擎 - 可观测性之分布式追踪
从代码到产品,我的IT职业成长之路
每个人的职业生涯都是一段充满转折和挑战的旅程,当然每一次职业转型都是一次重新定义自己的机会,从2015年开始,当时我刚踏入IT行业,成为一名Java开发者,后来随着时间的推移,我的职业方向逐渐转向了前端开发者,埋头于代码的世界。最终在2018年找到了属于自己的职业定位——产品经理。一路走来,我不断扩
从代码到产品,我的IT职业成长之路
基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型
过去的几个月,我们目睹了使用基于 transformer 模型作为扩散模型的主干网络来进行高分辨率文生图 (text-to-image,T2I) 的趋势。和一开始的许多扩散模型普遍使用 UNet 架构不同,这些模型使用 transformer 架构作为扩散过程的主模型。由于 transformer
基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型 基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型 基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型