Abstract

我们提出了一种新的基于频率的自监督学习(SSL)方法,显著提高了预训练的效果。以往在这一方向上的工作会屏蔽输入图像中的预定义频率,并采用重建损失来进行模型预训练。虽然取得了一些有希望的结果,但这种实现方式在我们的论文中指出了两个基本局限。首先,使用预定义频率忽略了图像频率响应的可变性。其次,通过频率过滤后的图像进行预训练后,所得模型在微调时需要相对更多的数据来适应自然图像。为了克服这些缺点,我们提出了基于傅里叶变换压缩的自知识蒸馏 FOurier transform compression with seLf-Knowledge distillation(FOLK)方法,整合了两个专门的理念。首先,受到图像压缩的启发,我们根据图像的频率响应自适应地选择屏蔽频率,从而为预训练创建更合适的自监督任务。其次,我们采用了一个由知识蒸馏驱动的双分支框架,使模型能够同时接收过滤后的图像和原始图像作为输入,大大减轻了下游任务的负担。我们的实验结果证明了FOLK在多个下游任务中(包括图像分类、小样本学习和语义分割)达到了与许多最新自监督学习方法竞争的性能。

1. Introduction

近年来,自监督学习(SSL)在视觉预训练方面受到了极大的关注。这种关注源于其能够从大量现成的无标签图像中提取有意义的视觉表征,而不需要昂贵的人工标注[Ben-Shaul et al., 2024, Su et al., 2024, Almalki and Latecki, 2024]。这一进展的关键在于通过不同的前置任务建立的多种预训练方法,包括多视图对比学习[Oord et al., 2018, Chen et al., 2020b, Tian et al., 2020b, He et al., 2020]、掩码图像建模(MIM)[Bao et al., 2022, He et al., 2022a, Xie et al., 2022, Monsefi et al., 2024a, Oquab et al., 2024]、掩码频率建模(MFM)[Xie et al., 2023, Liu et al., 2023, Zheng et al., 2024],以及自监督知识蒸馏(KD)[Kakogeorgiou et al., 2022, Zhou et al., 2022, Chen et al., 2020c, Chen and He, 2021, Caron et al., 2021]。在最近流行的掩码图像建模(MIM)方法中,一个关键策略是对图像的部分区域进行掩码,然后让模型重建这些隐藏部分或为其生成特征表示[Bao et al., 2022, Xie et al., 2022, He et al., 2022a, Yi et al., 2023]。通过这个过程,模型被鼓励学习捕捉未掩码和掩码图像部分之间内在结构的鲁棒特征表示,从而增强其对图像语义的理解。

不同于空间域的掩码(MIM),掩码频率建模 Masked Frequency Modeling(MFM) [Xie et al., 2023] 引入了一种自监督方法,对输入图像的频率成分进行掩码。由于图像的高层语义和低层细节可以被分离到不同的频率成分中[Oppenheim and Lim, 1981, Piotrowski and Campbell, 1982, Navard and Yilmaz, 2024],频率域为揭示图像的潜在模式提供了一个更便利的途径。MFM的前置任务是从频率过滤的图像中预测被掩码的频率成分(见第3.1节)。

它有两个主要优点:

  • 首先,它可以帮助避免在分析空间域中的原始像素值时遇到的问题,如空间冗余[Wang et al., 2020, Chen et al., 2023];
  • 其次,与MIM中常用的基于patch 的掩码不同,MFM适用于Vision Transformer(ViT)和卷积神经网络(CNN)模型。

然而,尽管MFM [Xie et al., 2023]展示了有希望的结果,但它有两个基本局限性。首先,MFM使用恒定滤波器,通过预定义的超参数半径来控制频谱中的掩码。这忽视了个别图像特有的内在结构,导致重建任务难度较低(见图1b和图1d)。其次,在前置阶段,MFM仅向模型展示频率掩码的图像,而没有机制适当地暴露原始图像的原始信息。这可能限制了预训练模型对正常图像分布的理解,进一步阻碍了MFM的训练效率和模型效果,尤其在小样本学习场景中表现不佳(见第4.2.2节)。

我们的动机是通过解决上述MFM的局限性,来实现更有效的视觉预训练。为此,我们提出了一个新的框架,整合了傅里叶变换压缩与自知识蒸馏,称为FOLK。与MFM类似,但与MIM不同,FOLK对图像的频率响应进行掩码,并支持ViT和CNN架构。此外,它从两个主要方面解决了MFM的问题。

首先,我们寻求一种改进的掩码方案,不再采用恒定滤波器,而是考虑每个图像独特的频率响应,以提高预训练效率和模型效果。受到MIM方法(如AttMask [Kakogeorgiou et al., 2022])中基于注意力的掩码的启发,在这些方法中,图像的最关键部分对学生模型是隐藏的,以创建更具挑战性的前置任务。FOLK利用傅里叶变换压缩[Pratt et al., 1969]的概念来掩码图像频率响应中最关键的部分。我们设计了两种滤波器,Com滤波器及其对应的RCom滤波器,用于保留(或去除)频谱中最高系数值(见图1)。与MFM中使用的恒定滤波器相比,这些Com和RCom滤波器自适应地掩码携带每个图像本质(或更精细细节)的频率,从而在训练样本中创建更大的变化。因此,模型面临的是一个更具挑战性的前置任务,强化其对每个图像独特的宏观和微观视觉线索的理解。

其次,我们考虑如何在预训练期间适当地向模型暴露自然图像信息,以增强微调效率。为此,FOLK采用了一种使用自监督教师-学生设计的知识蒸馏策略。将原始图像输入教师模型,并将频率掩码后的图像输入学生模型,学生模型不仅学习重建被掩码的频率(如同MFM所做的),还学习从同一图像的频率掩码视角重建原始图像的表示(由教师模型生成)。这种多任务的教师-学生方法允许模型在掩码图像和原始图像领域都有所感知,从而增强训练的稳定性和预训练模型在应用于下游任务时的效果,正如我们的实验结果所示(见第4.2节)。

总结起来,贡献有以下三点:

  • 在掩码频率建模中引入了一种新的掩码技术,使用Com和RCom滤波器,提出了一种更具挑战性和意义的前置任务,提高了自监督学习(SSL)的效率。
  • 提出了FOLK框架,这是一种创新的多任务自监督方法,通过自知识蒸馏,使模型在预训练阶段能够感知频率掩码图像和原始图像。
  • 通过广泛的实验,我们证明了FOLK的有效性。实验结果表明,在图像分类、小样本学习和语义分割等多种下游任务中,FOLK的表现与许多最先进的MIM和MFM技术相当或更好。

Related Work就简单了解下就行

3. Method

3.1 Preliminary and Background

在视觉模型的自监督学习领域,MFM [Xie et al., 2023] 引入了一种新颖的方法,与传统的空间域掩码策略有所不同。通过利用频率域,它包含了高频细节和低频元素,MFM 基于频率成分的掩码和被掩码频率的预测来进行学习过程。具体而言,给定一个单通道图像 x ∈ R H × W x \in \mathbb{R}^{H \times W} xRH×W,其频率表示通过二维快速傅里叶变换(FFT) F ( x ) \mathcal{F}(x) F(x) 获得:

F ( x ) ( u , v ) = ∑ h = 0 H − 1 ∑ w = 0 W − 1 x ( h , w ) e − i 2 π ( u h H + v w W ) , (1) \mathcal{F}(x)(u,v) = \sum_{h=0}^{H-1}\sum_{w=0}^{W-1} x(h, w) e^{-i2\pi \left( \frac{uh}{H} + \frac{vw}{W} \right)}, \tag{1} F(x)(u,v)=h=0H1w=0W1x(h,w)ei2π(Huh+Wvw),(1)

其中 x ( h , w ) x(h, w) x(h,w) 表示图像在空间坐标 ( h , w ) (h, w) (h,w) 处的像素值,而 F ( x ) ( u , v ) \mathcal{F}(x)(u, v) F(x)(u,v) 表示频谱坐标 ( u , v ) (u, v) (u,v) 处的复数频率值。这里, e e e 是欧拉数, i i i 是虚数单位。

为了从频谱中掩码一些频率,并让模型重建这些缺失的频率,首先通过以下方式获得频率掩码图像 x ~ \tilde{x} x~

x ~ = F − 1 ( F ( x ) ⊙ M ) , (2) \tilde{x} = \mathcal{F}^{-1}(\mathcal{F}(x) \odot M), \tag{2} x~=F1(F(x)M),(2)

其中 M ∈ { 0 , 1 } H × W M \in \{0, 1\}^{H \times W} M{0,1}H×W 是一个掩码,0表示相应的频率被掩码,1表示保留频率。 F − 1 \mathcal{F}^{-1} F1 表示逆傅里叶变换操作, ⊙ \odot 表示逐元素乘法。MFM 的学习目标是最小化重建的频率成分和原始频率成分之间的差异,可以写成:

L M F M = ∥ ( F ( x ) − F ( g θ ( x ~ ) ) ) ⊙ ( 1 − M ) ∥ 2 , (3) \mathcal{L}_{MFM} = \| (\mathcal{F}(x) - \mathcal{F}(g_{\theta}(\tilde{x}))) \odot (\mathbb{1} - M) \|_2, \tag{3} LMFM=(F(x)F(gθ(x~)))(1M)2,(3)

其中 F ( x ) \mathcal{F}(x) F(x) 是原始图像的频谱, F ( g θ ( x ~ ) ) \mathcal{F}(g_{\theta}(\tilde{x})) F(gθ(x~)) 是使用神经网络模型 g θ g_{\theta} gθ (由参数 θ \theta θ 表示)重建的频谱。 1 − M \mathbb{1} - M 1M 表示只考虑频谱中被掩码的区域用于计算损失。

3.2 FOLK Framework

在介绍我们提出的方法之前,我们先重申一下 MFM 方法 [Xie et al., 2023] 的局限性。首先,注意在公式 (3) 中,MFM 的损失依赖于应用于频谱的滤波器 M。然而,MFM 中使用的低/高通滤波器非常简单,使用固定半径的圆形区域(见图 1)。这可能会降低频率重建任务的难度,从而阻碍模型学习。MFM 的另一个限制是模型在预训练阶段仅看到频率掩码图像,即公式 (3) 中的 g θ ( x ~ ) g_{\theta}(\tilde{x}) gθ(x~)。因此,预训练的模型可能对自然图像相对不熟悉,在微调时需要更多的数据才能有效适应(见第 4.2.2 节)。

为了克服这些局限性并实现有效的掩码频率建模,我们提出了 FOLK 框架。我们的关键思想包括创建基于频率的智能滤波器 Com 和 RCom,以及基于师生设计的自蒸馏策略。

在这里插入图片描述

3.2.1 Informed Filters

成功的视觉预训练在很大程度上依赖于向模型呈现合适且具有挑战性的预任务。AttMask [Kakogeorgiou et al., 2022] 通过掩盖图像中最受关注的区域,比随机掩码创建了更有效的训练方案。然而,对于 MFM 方法,这个差距依然存在,因为目前仅探索了恒定的掩码/滤波器 [Xie et al., 2023],这可能在预训练中呈现了较少的挑战性任务。

为弥补这一差距,我们引入了两种类型的滤波器,Com 和 RCom,用于智能掩码。受傅里叶图像压缩技术 [Pratt et al., 1969] 的启发,其中携带大部分图像视觉信息的最显著频率(具有最高幅度的频率)被保留,而其他频率被丢弃以实现高效存储或带宽使用,我们的 Com 滤波器有选择地保留这些显著频率并丢弃其余部分。这种方法强调了图像的主要语义,并要求重建对应于边缘等细节的次要频率。相反,RCom 滤波器去除显著频率,并需要从次要频率保留的细节中重建它们。通过在预训练期间应用这两种滤波器,模型被有效地训练以理解宏观和微观视觉线索,从而增强其在下游任务中的泛化能力和有效性。

Com 和 RCom 滤波器的生成如图 2 所示。首先将输入图像转换为灰度图像,以确保为所有三个 RGB 通道创建一个通用滤波器,因为分别对每个通道生成和应用滤波器可能导致不自然和损坏的视觉信息。图像通过 2D FFT 转换为频谱,然后根据一个阈值(在我们的实验中,均匀采样自一组值 [0.005, 0.01, 0.05]),确定具有最高幅度的频率分量。随后创建智能滤波器以保留(Com)或掩盖(Rcom)这些频率。两种滤波器以相等的概率(即 50%)随机选择并应用于频谱。通过对过滤后的频谱应用逆 FFT,我们恢复了一个频率掩码图像,该图像将在预训练期间作为模型的输入。需要注意的是,Com 和 RCom 滤波器是基于每个图像唯一生成的(附录 C 中提供了示例),这在训练样本中引入了更大的变化,并且与使用恒定滤波器相比,呈现了更困难的训练任务。除了将我们的方法与 MFM [Xie et al., 2023] 中使用的低/高通滤波器进行比较外,我们还进行了一组使用随机频率滤波器的消融研究,以证明我们提出的智能滤波器的有效性(见附录 B.3)。

2

3.2.2 Making Backbone Familiar with Natural Images

为了进一步提高在频域掩码建模中的训练效率和模型的鲁棒性,我们在我们提出的方法中引入了自蒸馏设计 [Grill et al., 2020; Caron et al., 2021; Tarvainen 和 Valpola, 2017]。原始的 MFM [Xie et al., 2023] 方法仅要求模型从频率掩码视图中重建缺失的频率分量。这样的方法可能忽略了原始图像空间的数据分布,因为模型在预训练阶段只看到频率掩码的图像,导致在微调过程中需要更多的数据来适应自然图像。为了解决这个问题,我们提出通过一种自蒸馏技术将原始图像信息正确地注入训练过程中,这种技术在 BYOL [Grill et al., 2020] 和 DINO [Caron et al., 2021] 中得到了应用。

FOLK 框架如图 3 所示。需要注意的是,FOLK 不需要像 BEiT [Bao et al., 2022] 那样的离线标记器预处理的额外训练阶段。FOLK 从输入图像 x x x 生成两个视图 u u u v v v,通过不同的变换创建多样的视角。这些变换包括随机裁剪、颜色抖动等,遵循 DINO [Caron et al., 2021] 的方法。与 DINO 或 ATTMask [Kakogeorgiou et al., 2022] 不同,我们不使用局部视图的概念,以保持框架的高效性。对视图 u u u(或 v v v)应用 2D FFT 后,按该视图(详见第 3.2.1 节和图 2)生成独特的 Com 和 RCom 滤波器。然后随机选择一个滤波器应用于频谱,保留的频率分量经过逆 FFT 处理恢复为频率掩码视图 u ~ \tilde{u} u~(或 v ~ \tilde{v} v~),见方程 2。

学生模型接收这个频率掩码视图 u ~ \tilde{u} u~(或 v ~ \tilde{v} v~)并对两个目标进行预测:重建滤波器丢弃的缺失频率,以及重建由教师模型生成的另一原始视图 v v v(或 u u u)的特征表示。学生模型附加了两个不同的头,每个头实现一个预测任务。MFM 头 h θ ^ \hat{h_\theta} hθ^ 用于重建缺失的频率,其实现方式是原始 MFM [Xie et al., 2023] 设计之后的单线性层。学生头 h θ h_\theta hθ(和教师头 h φ h_\varphi hφ)旨在重建另一原始视图的特征表示,采用三层多层感知器(MLP)设计。此外,学生头(和教师头)之后跟随一个缩放的 softmax 函数,将输出转化为概率分布,以计算蒸馏损失(如下所述)。

具体来说:
P s ( x ) ( i ) = e x p ( f θ ( x ) ( i ) / τ s ) ∑ k = 1 K e x p ( f θ ( x ) ( k ) / τ s ) (4) P_s(x)^{(i)} = \frac{exp(f_{\theta}(x)^{(i)} / \tau_s)}{\sum_{k=1}^K exp(f_{\theta}(x)^{(k)} / \tau_s)} \tag{4} Ps(x)(i)=k=1Kexp(fθ(x)(k)/τs)exp(fθ(x)(i)/τs)(4)
其中, f θ ( x ) = h θ ( g θ ( x ) ) f_\theta(x) = h_\theta(g_\theta(x)) fθ(x)=hθ(gθ(x)) K K K h h h 的输出维度, τ s > 0 \tau_s > 0 τs>0 是温度参数。教师模型的情况类似,使用 P t P_t Pt τ t \tau_t τt。缩放的 softmax 的使用允许输出分布的锐化(特别是对于教师模型),以避免模型崩溃 [Caron et al., 2021]。此外,我们遵循 DINO 中提出的中心化方法,进一步避免模型崩溃并减少对大批量数据的依赖。头部的详细描述见附录 A.3。

3

图 3:提出的 FOLK 框架。输入图像的两个视图(u 和 v)通过图 2 中介绍的有信息的滤波过程进行处理。此过程生成两个频率掩码视图( u ~ \tilde{u} u~ v ~ \tilde{v} v~),作为学生模型的输入。学生模型的任务是从掩码视图 u ~ \tilde{u} u~(或 v ~ \tilde{v} v~)中重建缺失的频率,并且从教师模型生成的另一原始视图 v(或 u)中重建特征表示,使用的是掩码视图 u ~ \tilde{u} u~(或 v ~ \tilde{v} v~)。学生模型 g θ g_\theta gθ 和教师模型 g φ g_\varphi gφ 具有相同的架构,但参数不同,唯一不同的是学生模型有一个额外的 MFM 头 h θ ^ \hat{h_\theta} hθ^。只有学生模型(及其两个头)通过反向传播进行更新,而教师模型的参数则通过对应学生参数的指数移动平均(EMA)进行定期更新。

FOLK框架通过其教师-学生设计有效地实现了未掩盖的原始图像信息的揭示。在预训练阶段,教师模型看到的是自然的图像,这些图像与在微调阶段遇到的图像更为一致,从而提高了微调效率,特别是在少量样本学习的场景中。另一方面,学生模型只观察掩码视图,但通过教师模型的指导,使用以下蒸馏损失进行训练。学生模型和教师模型 g θ g_\theta gθ g ϕ g_\phi gϕ 以及它们的头部 h θ h_\theta hθ h ϕ h_\phi hϕ 具有相同的架构和初始化,但在训练过程中它们的参数是不同的。只有学生模型 g θ g_\theta gθ 及其两个头部 h θ h_\theta hθ h θ ^ \hat{h_\theta} hθ^ 会通过损失反向传播进行更新,而教师模型的参数则通过指数移动平均(EMA)周期性地更新。

为了让信息较少的学生模型模仿信息更多的教师模型,蒸馏损失被用来强制学生模型学习教师模型感知的原始视图。对于单个输入图像 x x x,该损失可以写作:

L dis = − [ P t ( u ) log ⁡ ( P s ( v ~ ) ) + P t ( v ) log ⁡ ( P s ( u ~ ) ) ] (5) \mathcal{L}_{\text{dis}} = - \left[ P_t(u) \log \left( P_s(\tilde{v}) \right) + P_t(v) \log \left( P_s(\tilde{u}) \right) \right] \tag{5} Ldis=[Pt(u)log(Ps(v~))+Pt(v)log(Ps(u~))](5)

其中 u u u v v v x x x 的两个不同视图, u ~ \tilde{u} u~ v ~ \tilde{v} v~ 是对应的频率掩码视图。 P s P_s Ps P t P_t Pt 分别是学生模型和教师模型的输出概率分布,正如方程(4)所示。教师模型的 EMA 更新规则为: ϕ ← λ ϕ + ( 1 − λ ) θ \phi \leftarrow \lambda \phi + (1 - \lambda) \theta ϕλϕ+(1λ)θ,确保教师模型的知识逐渐整合到学生模型中。

与 MFM [Xie et al., 2023] 相似,FOLK框架与 ViT 和基于 CNN 的架构都兼容。如图3所示,当使用基于 ViT 的模型作为学生(和教师)时,最终编码器层的补丁令牌会输入到 MFM 头部,而类别令牌 [CLS] 则传递到学生(和教师)头部。相反,当使用基于 CNN 的模型时,框架会使用来自 CNN 编码器的最终特征图作为 MFM 头部的输入,平均池化后的特征图会作为学生(和教师)头部的输入。我们在附录B.1中提供了使用CNN模型(如ResNet-50 [He et al., 2016])的实验和结果。因此,FOLK提供了一种在不同架构范式之间一致的处理方法。

3.2.3 Comprehensive Loss Calculation

最终,通过结合两个主要的损失组件,可以得到一个综合损失:

L tot = α ⋅ L dis + L MFM (6) \mathcal{L}_{\text{tot}} = \alpha \cdot \mathcal{L}_{\text{dis}} + \mathcal{L}_{\text{MFM}} \tag{6} Ltot=αLdis+LMFM(6)

其中超参数 α \alpha α 控制着两个损失项之间的权重,在我们的实验中,除非另有说明, α \alpha α 被设置为1。需要注意的是,对于单个输入图像, L MFM \mathcal{L}_{\text{MFM}} LMFM 会对两个不同视图的两个项取平均。这个综合损失有助于同时进行两个任务的模型学习:一个是掩码频率的重建(通过 L MFM \mathcal{L}_{\text{MFM}} LMFM),另一个是通过自蒸馏进行的原始图像特征重建(通过 L dis \mathcal{L}_{\text{dis}} Ldis)。此外,关于超参数 α \alpha α 不同选择的消融实验会在第4.2.4节中提供。

4. Experiments

目前还没有公开代码,不知道之后会不会公开。

Logo

GitCode 天启AI是一款由 GitCode 团队打造的智能助手,基于先进的LLM(大语言模型)与多智能体 Agent 技术构建,致力于为用户提供高效、智能、多模态的创作与开发支持。它不仅支持自然语言对话,还具备处理文件、生成 PPT、撰写分析报告、开发 Web 应用等多项能力,真正做到“一句话,让 Al帮你完成复杂任务”。

更多推荐