AAAI | 2025 | The Dynamic Duo of Collaborative Masking and Target for Advanced
Masked Autoencoder Learning

  • 论文:https://doi.org/10.1609/aaai.v39i18.34145
  • 代码:NONE
  • 会议:AAAI
  • 时间:2025

创新点

  • ⾸次让学⽣模型参与 “选遮挡块” 和 “定重建⽬标”,不是⽼师模型的单⽅⾯指挥
  • ⽤ “线性聚合” 的⽅法融合师⽣注意⼒,既⾼效⼜好⽤,不⽤复杂计算
  • 两阶段训练循序渐进,先学基础再练协作,模型学得更稳、更好

本文贡献

  • 本文团队提出一种简单而有效的掩码自编码器,可实现协作掩码与目标,用于增强基于MIM的视觉预训练
  • 本文团队提出一种新颖的协同掩蔽机制,通过教师网络学生网络注意力图进行线性聚合,以实现强大的引导效果
  • 大量实验结果全面证明,他们的CMT-MAE在下游任务中展现出超越以往基线模型的前沿优势(如下图所示)

图1:本研究提出的CMT-MAE方法与MAE及DINO在预训练ViT-B/16模型上的性能对比(横坐标是不同下游任务,纵坐标是性能分数)。对比得出:在所有下游任务中,本文团队的方法均显著优于既往基线方法

摘要及引言

  • 掩码⾃编码器是当前主流的视觉学习⼯具,它的思路很简单:故意挡住图像的⼀部分(如 75%),让模型根据没挡住的部分补出挡住的内容,从⽽学会图像的核⼼特征
  • 但之前的 MAE 有个⼤问题:要么随机遮图像块,要么只通过“⽼师模型”(⽐如 CLIP 这种提前学好的⼤模型)提供的信息来掩码,完全忽略了 “学⽣模型”(正在训练的模型)在训练中慢慢积累的知识
  • 本文团队发现其实学生模型的积累及反馈能反过来帮⽼师选更该掩码的地⽅、定更合理的重建⽬标,让学习效果更好,故这篇论⽂就是要解决 “师⽣协作” 的问题,让⽼师和学⽣⼀起工作,提升模型的学习及重建能力
  • 本文提出的模型在预训练阶段分为两个阶段:
    1.第一阶段:教师转换器编码器(即CLIP)从输入图像的最后一个注意力层提取注意力图,以指导掩码。学生编码器从未掩码的图像块中生成特征,这些特征与掩码标记拼接后输入解码器,用于恢复掩码图像块的教师特征
    2.第二阶段:应用学生动量编码器生成学生引导注意力图,并将其与教师引导注意力图线性聚合,生成具有协作比例的协作注意力图用于协作掩码。随后,掩码标记与学生编码器生成的未掩码图像块特征拼接后输入解码器
    3.最后,两个预测头线性应用到掩码图像块的教师和学生特征重建中,用于协作目标

预备知识

  • 教师模型:提前训练好、知识储备丰富的模型(⽐如 CLIP),⽤来指导学⽣模型学习
  • 学⽣模型:正在训练的模型,通过和教师模型协作,逐步提升⾃⾝能⼒
  • 协作掩码:将⽼师和学⽣所得的信息结合在一起从而决定该遮挡哪些图像块,而不是⽼师独断或随机选择
  • 协作⽬标:学⽣重建遮挡部分时,不仅要参考⽼师的特征,还要结合学生学到的特征,双重⽬标让重建更精准
  • 遮罩图像建模:通过对图像的部分区域进行遮盖,再让模型根据未被遮盖的图像信息,重建出被遮盖的部分从而使模型能自主学到图像的局部细节、全局结构等深层特征的过程

方法

核心思路

  • 综合教师模型(CLIP)和学⽣模型的能力,⼀起选遮挡的图像块(协作掩码),⼀起定重建的⽬标(协作⽬标),让学⽣在训练中既能学⽼师的标准答案,⼜能发挥⾃⼰的学习能力,最终学到更实⽤的图像特征

初步研究

  • 问题设定与符号说明:对于一个维度为3 × H × W的图像和P的补丁分辨率,我们的目标是学习一个掩码自编码器框架,该框架包含 f e ( ⋅ ) f_e(\cdot) fe()编码器和解码器 f d ( ⋅ ) f_d(\cdot) fd(),利用未被掩码的图像补丁来恢复被掩码的补丁
  • 重访掩码自编码器:首先在所有图像块上应用随机掩码集M,随后通过编码器从未被掩码的图像块中提取特征。最终,未被掩码的嵌入向量与掩码标记在解码器中拼接,以恢复被掩码图像块的原始像素每个图像的原始掩码损失通过目标像素 p i p_i pi与预测归一化像素 p ^ i \hat{\mathbf{p}}_i p^i之间的均方损失计算得出:
    L mae = 1 ∣ M ∣ ∑ i ∈ M ∥ p i − p ^ i ∥ 2 2 \mathcal{L}_{\text{mae}} = \frac{1}{|M|} \sum_{i \in M} \| \mathbf{p}_i - \hat{\mathbf{p}}_i \|_2^2 Lmae=M1iMpip^i22

其中|M|表示掩膜集M中的掩膜块总数

  • 重访教师引导的蒙版自编码器:一些学者通过教师Transformer编码器提取注意力图A 生成掩码集M ,并将表征作为重建目标。针对每张图像的教师引导掩码损失,通过目标 f i t f^t_i fit与预测归一化特征 f i t f^t_i fit f ^ i t \hat{\mathbf{f}}_i^t f^it之间的均方损失计算得出,其公式为:
    L teacher-mae = 1 ∣ M t ∣ ∑ i ∈ M t ∥ f i t − f ^ i t ∥ 2 2 \mathcal{L}_{\text{teacher-mae}} = \frac{1}{|M^t|} \sum_{i \in M^t} \| \mathbf{f}_i^t - \hat{\mathbf{f}}_i^t \|_2^2 Lteacher-mae=Mt1iMtfitf^it22

其中| M t M_t Mt|表示教师引导掩蔽集| M t M_t Mt|中的掩蔽块总数

  • 然而,这种训练目标将对协同掩码和目标提出主要挑战:这些方法并未明确将自训练(即学生)模型纳入预训练过程中与教师网络的合作机制。此外,教师网络与学生网络存在知识水平差异
  • 为应对这一挑战,本文团队提出了一种简单而有效的掩码自编码器,该模型能够实现协作式掩码和目标生成,即CMT-MAE

协作式掩蔽

  • 为实现教师与学生共同指导的协作掩码生成,本文团队提出了一种两阶段训练范式。该方案通过线性聚合两个模型的注意力图,构建了简洁的协作掩码机制,从而增强生成掩码的效能
  • 在第一阶段,本文团队采用教师端Transformer编码器,通过输入图像提取 A t A^t At从末尾注意力层用于指导掩码。学生编码器从未被掩码的图像块中生成特征,这些特征与被掩码的标记拼接后输入解码器,以恢复被掩码图像块 f i t f_i^t fit的教师特征
  • 在第二阶段,本文团队采用学生动量编码器,将输入图像转化为学生引导的注意力图 A s A^s As,并使用指数移动平均法对在线编码器进行更新,其系数为m。最终将学生 A s A^s As和教师 A t A^t At注意力图线性聚合,生成形式如下的协作注意力图A :
    A c = a ∗ A s + ( 1 − a ) ∗ A t \mathbf{A}^c = \boldsymbol{a} * \mathbf{A}^s + (1 - \boldsymbol{a}) * \mathbf{A}^t Ac=aAs+(1a)At

其中α表示协作比。协作图 A c A^c Ac用于生成掩膜集 M c M^c Mc,将输入图像分割为掩膜和非掩膜两部分,用于第二阶段的训练

在这里插入图片描述

图2:所提出的协作掩码与目标掩码自编码器(CMT-MAE)框架示意图

  • 第一阶段
    教师(CLIP 模型):先看输入图像,从最后一个注意力层里提炼出注意力图 A t A_t At,告诉学生该关注图像哪些关键部分,以此指导掩码生成
    学生(编码器):照着教师给的 “重点提示”,先分析没被遮住的图像块,提炼出这些块的特征;再把这些特征和 “遮罩标记”(代表被遮住的部分)拼在一起,交给 “解码器”,目标是学着恢复被遮住部分的 “教师级特征”(教师之前提炼的关键信息),相当于学生先模仿教师的思路学基础
  • 第二阶段
    学生升级:学生多了个 “动量编码器”,自己也能看图像生成学生注意力图 A s A_s As,不再只靠教师指导
    协作掩码:把教师的 A t A_t At和学生的 A s A_s As按一定比例(协作比例 α)融合,形成 协作注意力图 A c A_c Ac,按这个图来遮图像块 —— 既参考教师的成熟经验,又结合学生自己学到的新发现,遮得更合理
    协作目标:解码器不仅要恢复教师的特征,还要恢复学生自己学到的特征(学生特征);最后用两个预测头分别评判恢复效果,再按 α 比例算总误差(协作损失),让学生既不脱离教师指导,又能发挥自己的学习能力

协作目标

  • 除了协作掩码机制,本文团队在第二阶段引入协作目标作为训练目标,通过动态选择教师和学生模型的表征进行同步优化
  • 基于协作注意力图A指导的掩码集M,本文团队将被掩码标记的标记学生编码器未被掩码标记区域的特征进行拼接,输入到解码器中
  • 最终,本文团队通过两个预测头对被掩码标记区域的教师特征 f i t f^t_i fit和学生特征 f i s f^s_i fis进行线性重构,从而生成协作目标
  • 本文团队的模型在第二阶段的总体目标是以端到端方式实现优化,具体表现为:
    L cmt-mae = 1 ∣ M c ∣ ∑ i ∈ M c α ∗ ∥ f i s − f ^ i s ∥ 2 2 + ( 1 − α ) ∗ ∥ f i t − f ^ i t ∥ 2 2 \mathcal{L}_{\text{cmt-mae}} = \frac{1}{|M^c|} \sum_{i \in M^c} \alpha * \| \mathbf{f}_i^s - \hat{\mathbf{f}}_i^s \|_2^2 + (1 - \alpha) * \| \mathbf{f}_i^t - \hat{\mathbf{f}}_i^t \|_2^2 Lcmt-mae=Mc1iMcαfisf^is22+(1α)fitf^it22

其中| M c M^c Mc|表示协作掩蔽集 M c M^c Mc中掩蔽补丁的总数。其中 f i s f^s_i fis f ^ i s \hat{\mathbf{f}}_i^s f^is分别表示学生特征的目标值和预测值, f i s f^s_i fis f ^ i s \hat{\mathbf{f}}_i^s f^is表示教师特征。α为协作比例。需要说明的是,α这一协作比例用于计算教师和学生目标的协作损失

实验

实验设计

  • 数据集:使用ImageNet-1K进行图像分类,使用MS-COCO进行目标检测,并采用ADE20K数据集进行语义分割;采用MaskR-CNN作为检测器,并采用基于ImageNet-1K预训练ViT-Base的UPerNet方法进行评估
  • 评估指标说明:采用线性探测和微调的分类准确率作为评估指标,使用 A P b AP^b APb A P m AP^m APm作为边界框和实例掩码的评估指标,并且在ImageNet-1K数据集上进行预训练,采用AdamW优化器,学习率设为1.5e-4,衰减率为0.05,批量大小为4096

与既往研究的比较

图像分类

在这里插入图片描述

表格⾏是不同模型,列是预训练数据集、预训练轮次,以及ViT-B/16、ViT-L/16模型的 “线性探测”“微调” Top-1 准确率

实验结果:

  • ViT-B/16 模型:CMT-MAE的线性探测准确率79.8%、微调准确率85.7%,⽐原版 MAE(68.0%、83.6%)分别⾼11.8%、2.1%,⽐强基线DINO(78.2%、82.8%)也更⾼
  • ViT-L/16 模型:CMT-MAE的微调准确率 87.2%,同样是所有模型⾥最⾼的

结论:

  • 证明不管是基础版还是⼤型版ViT模型,CMT-MAE的预训练⽅法都能让模型学到更好的分类特征

目标检测与实例分割&语义分割

在这里插入图片描述

表格⾏是不同模型,列是 AP 盒(⽬标检测准确率)、AP 掩码(实例分割准确率)、mIoU(语义分割准确率)

实验结果:

  • ⽬标检测(AP 盒):CMT-MAE 达到 52.8%,⽐原版 MAE(1600 轮训练,50.3%)⾼ 2.5 个百分点
  • 实例分割(AP 掩码):CMT-MAE45.7%,⽐ MAE(44.9%)⾼ 0.8 个百分点
  • 语义分割(mIoU):CMT-MAE52.9%,⽐ MAE(46.1%)⾼ 6.8 个百分点,甚⾄超过了⽤ 4 亿图像 - ⽂本对预训练的 MVP 模型

结论:

  • 说明 CMT-MAE 学到的特征不仅能分类,还能精准定位物体、分割像素,实⽤性更强

视频目标分割

在这里插入图片描述

表格⾏是不同模型,列是模型 backbone(主⼲⽹络,ViT-B/16、ViT-L/16),以及 ( J (J (J& F ) m F)_m F)m J m J_m Jm F m F_m Fm 三个评估指标

实验结果:

  • ViT-B/16 模型:CMT-MAE 的 ( J (J (J& F ) m F)_m F)m=57.6%,⽐ MAE(51.0%)⾼ 6.6 个百分点,⽐I-JEPA(56.2%)也更⾼
  • ViT-L/16 模型:CMT-MAE 的 ( J (J (J& F ) m F)_m F)m=60.5%,同样是最佳,⽐MAE(53.4%)⾼ 7.1 个百分点

结论:

  • 证明 CMT-MAE 的特征能很好地适应视频的 “时空变化”持续跟踪并分割⽬标
    在这里插入图片描述

每组有四⾏:第⼀⾏是视频原始帧(⽐如动物、⼈物场景),第⼆⾏是 “真实掩膜”(⼈⼯标注的物体轮廓,作为标准答案),第三⾏是 MAE 的预测结果,第四⾏是 CMT-MAE 的预测结果

实验结果:

  • MAE 的预测结果:轮廓不精准,⽐如物体边缘有缺⼝、和背景混淆(⽐如动物的⽑发部分没分割⼲
    净)
  • CMT-MAE 的预测结果:轮廓更贴合真实掩膜,边缘更平滑、细节更完整(⽐如能精准分割出物体
    的⼩部件、复杂轮廓)

结论:

  • 说明 CMT-MAE 学到的特征能更好地捕捉视频中物体的 “持续特征”分割质量更⾼

实验分析

协作掩码与协作目标

在这里插入图片描述

表格⾏是 “CM 是否使⽤”“CT 是否使⽤” 的四种组合(都不⽤、只⽤ CM、只⽤ CT、都⽤),列是线性探测、微调、AP 盒等所有下游任务指标

实验结果:

  • 都不⽤(基线 MAE):所有指标最低(⽐如线性探测 68.0%、 ( J (J (J& F ) m F)_m F)m 51.0%)
  • 只⽤ CM:指标明显提升(线性探测 73.5%、 ( J (J (J& F ) m F)_m F)m 53.8%)
  • 只⽤ CT:指标也提升(线性探测 74.2%、 ( J (J (J& F ) m F)_m F)m 54.6%)
  • 都⽤(CMT-MAE):所有指标达到最⾼(线性探测 79.8%、 ( J (J (J& F ) m F)_m F)m 57.6%)

结论:

  • 证明 CM 和 CT 是 “1+1>2” 的组合,两者共同作⽤才是 CMT-MAE 性能提升的关键

协作比例的权衡考量

在这里插入图片描述

表格⾏是不同协作⽐例(0%、10%、30%、…、100%),列是所有下游任务指标

实验结果:

  • 协作⽐例 α=30%时,所有指标都达到最佳(⽐如线性探测79.8%、AP盒52.8%、 ( J (J (J& F ) m F)_m F)m 57.6%)
  • α=0%(全听⽼师):指标⼀般(线性探测77.1%);α=100%(全听学⽣):指标最差(线性探测77.6%)

结论:

  • 说明师⽣协作不能偏废,让学⽣占30%的话语权、⽼师占70%,既能发挥⽼师的经验优势,⼜能利⽤学⽣的实时学习成果,效果最优

结论

  • 在本研究中,本文团队提出了一种简单却高效的掩码自编码器CMT-MAE,该模型能够同时实现协同掩码和目标的双重目标
  • 本文团队通过教师模型和学生模型的注意力机制进行线性聚合,创新性地设计了协同掩码机制。进一步地,他们将两者的表征作为解码器的协同目标用于重构任务
Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐