ICLR2023-适用于密集预测任务的通用小样本学习(VTM)
1. 创新点:
目前的小样本学习方法通常只针对一组有限的任务,比如分类、语义分割。本文提出了可以用于任意密集预测任务的小样本学习器——VTM,它通过 图像和标签的Patch级嵌入令牌的非参数匹配 概括了所有任务。
密集预测任务是计算机视觉领域的一类重要任务,如语义分割、深度估计、边缘检测和关键点检测等。对于这类任务,手动标注像素级标签面临着难以承受的巨额成本。因此,如何从少量的标注数据中学习并作出准确预测,即小样本学习,是该领域备受关注的课题。
2. 问题建模
首先对于任务T,可以表示为:
即将任意的三通道任务转化为
C
T
C^T
CT 个单通道子任务,对于每一个子任务都使用公式2中共享的模型F独立建模。
即对任意一个任务T,根据查询图像
X
q
X^q
Xq和支持集
S
T
S_T
ST 产生预测
Y
q
Y^q
Yq 。
将公式2推广到一般函数形式,使他产生具有统一框架的适应任意任务的结构化标签,本文选择在Patch上操作的非参数方法,其中 查询标签是通过支持标签的加权组合获得的。
设
X
=
{
x
j
}
j
≤
M
X = \{xj\}_{j≤M}
X={xj}j≤M 表示大小为
M
=
h
×
w
M = h × w
M=h×w 的Patch网格上的图像(或标签),其中
x
j
x_j
xj 是第
j
j
j 个Patch。给定查询图像
X
q
=
{
x
j
q
}
j
≤
M
X^q = \{x^q_j\}_{j≤M}
Xq={xjq}j≤M 和任务
T
T
T的支持集
{
(
X
i
,
Y
i
)
}
i
≤
N
=
{
(
x
k
i
,
y
i
k
)
}
k
≤
M
,
i
≤
N
\{(X^i, Y^i)\}_{i≤N} = \{(x^i_k, y^k_i)\}_{k≤M, i≤N}
{(Xi,Yi)}i≤N={(xki,yik)}k≤M,i≤N ,
将所有Patch投影到嵌入空间,并预测 查询图像 的每一个Patch的标签
Y
q
=
{
y
j
q
}
j
≤
M
Y_q = \{y^q_j\}_{j≤M}
Yq={yjq}j≤M。
其中, f T ( x ) = f ( x ; θ , θ T ) ∈ R d f_T(x) = f(x; θ, θ_T) ∈ R^d fT(x)=f(x;θ,θT)∈Rd 表示图像编码器, g ( y ) = g ( y ; φ ) ∈ R d g(y) = g(y; φ) ∈ R^d g(y)=g(y;φ)∈Rd 表示标签编码器, σ : R d × R d → [ 0 , 1 ] σ: R^d × R^d → [0, 1] σ:Rd×Rd→[0,1]表示定义在图像补丁嵌入空间上的相似性函数。然后,可以将每个预测的标签嵌入映射到一个标签补丁 ˆ y j q = h ( g ( y j q ) ) ˆy_j^q = h(g(y_j^q)) ˆyjq=h(g(yjq)),其中引入了一个标签解码器 h ≈ g − 1 h ≈ g^{-1} h≈g−1。
3. VTM架构
图1:VTM的整体架构。该模型是一个分层编码器-解码器,有四个主要组件:图像编码器 f T f_T fT、标签编码器 g g g、标签解码器 h h h和匹配模块。
给定一个新任务的少量标记示例,VTM 首先会根据给定的示例以及示例的标签调整其对相似性的理解,从示例图像块中锁定与待预测图像块相似的图像块,通过组合它们的标签来预测未见过的图像块的标签。
图1表明该模型具有分层编码器-解码器架构,该架构在多个层次中实现了公式3的Patch级非参数匹配,其中包含四个组件:图像编码器fT、标签编码器g、标签解码器h和匹配模块。给定查询图像和支持集,图像编码器首先独立提取每个查询和支持图像的Patch级嵌入(标记)。标签编码器类似地提取每个支持标签的标记。给定每个层次的标记,匹配模块执行公式 3的非参数匹配以推断查询标签的标记,最终由标签解码器推断出查询图像的标签。
图像编码器:
作者使用Vision Transformer作为图像编码器。该方法在共享权重的同时独立地对图像进行查询和支持,从而在多个层次上对图像patch进行标记化表示。作者在四个中间ViT块中提取令牌以形成分层特征。为了帮助学习广泛任务的一般表示,本文初始化了预训练的BEiT (Bao等人,2021)的参数,这是自我监督的,因此对特定任务的偏见较小。
如之前所述,作者将图像编码器设计为具有两组参数θ和θT,其中θ在所有任务中共享,θT特定于每个任务T。 在通过θT设计适应机制的众多候选中,我们发现 偏差调谐(Cai et al, 2020;Zaken等人,2022) 在经验上提供了最佳的效率和性能。为此,我们在元训练和元测试中为每个任务使用单独的偏差集,同时共享所有其他参数。
标签编码器
标签编码器采用与图像编码器相同的ViT架构,并提取支持标签的令牌嵌入。与图像编码器类似,标签编码器独立地应用于每个支持标签,主要区别在于它一次只看到一个通道,因为我们将每个通道视为独立的任务。然后从与图像编码器匹配的多个层次中提取标签令牌。与图像编码器相反,标签编码器的所有参数都是从头开始训练并跨任务共享的。
匹配模块
作者在每个层次上实现匹配模块作为 多头注意力层(Vaswani等人,2017)。在图像和标签编码器的每个层次上,作者首先从图像和标签编码器的中间层获取查询图像
X
q
X^q
Xq的标记,表示为
{
q
j
}
j
≤
M
\{q_j\}_{j≤M}
{qj}j≤M,以及支持集
{
(
X
i
,
Y
i
)
}
i
≤
N
\{(X_i,Y_i)\}_{i≤N}
{(Xi,Yi)}i≤N的标记,表示为
{
(
k
k
i
,
v
k
i
)
}
k
≤
M
,
i
≤
N
\{(k^i_k, v^i_k)\}_{k≤M,i≤N}
{(kki,vki)}k≤M,i≤N。然后将这些标记堆叠成行矩阵,
q
∈
R
M
×
d
,
k
、
v
∈
R
N
M
×
d
q ∈ R^{M×d},k、v ∈ R^{N M×d}
q∈RM×d,k、v∈RNM×d。然后,查询标签在该层次的标记被推断为多头注意力层的输出 ,如下所示:
其中 H H H是头的数量, d H d_H dH是头大小, w h Q 、 w h K 、 w h V ∈ R d × d H , w O ∈ R H d H × d w_h^Q、w_h^K、w^V_h ∈ R^{d×d_H},w^O ∈ R^{Hd_H×d} whQ、whK、whV∈Rd×dH,wO∈RHdH×d。
值得注意的是,公式5中的每个注意力头实现了方程式3中的非参数匹配的直觉。这是因为每个查询标签标记是根据查询和支持图像标记q、k之间的相似性,以支持标签标记v的加权组合来推断的。在这里,方程式3的相似性函数σ被实现为缩放的点积注意力。由于每个头涉及不同的可训练投影矩阵 w h Q 、 w h K 、 w h V w_h^Q、w_h^K、w^V_h whQ、whK、whV,方程式4中的多头注意力层能够学习具有不同相似性函数的匹配算法的多个分支(头)。
标签解码器
标签解码器 h h h接收在多个层次推断的查询标签令牌,并将它们组合以预测原始分辨率的查询标签。作者采用了 密集预测transformer的多尺度解码器架构(Ranftl等人,2021),因为它可以与ViT编码器和多级令牌无缝协作。在解码器的每个层次结构中,推断的查询标签令牌首先在空间上连接到一个恒定大小的特征映射 ( M — > h × w ) (M —>h × w) (M—>h×w)。然后,对每个特征图应用不同步长的(转置的)卷积层,产生分辨率不断增加的特征金字塔。多尺度特征通过卷积块逐步上采样和融合,然后使用卷积头进行最终预测。
与标签编码器类似,标签解码器的所有参数都是从头开始训练并跨任务共享的。这允许解码器元学习一种从预测的查询标签令牌解码结构化标签的通用策略。根据之前提到的的通道分割,解码器的输出是单通道,这允许它应用于任意数量通道的任务。
4. 模型训练
在训练任务
T
t
r
a
i
n
T_{train}
Ttrain的带标签数据集
D
t
r
a
i
n
D_{train}
Dtrain上按照标准的逐个迭代的元学习协议来训练我们的模型。在每个任务
T
T
T的一个迭代中,我们从
D
t
r
a
i
n
D_{train}
Dtrain中采样两个带标签集合
S
T
S_T
ST和
Q
T
Q_T
QT。然后我们使用
S
T
S_T
ST作为支持集来训练模型来预测
Q
T
Q_T
QT中的标签。我们重复使用
D
t
r
a
i
n
D_{train}
Dtrain中的各种密集预测任务的迭代,以使模型可以学习有关少样本学习的一般知识。记
F
(
X
q
;
S
T
)
F(X^q; S_T)
F(Xq;ST)表示模型在使用支持集
S
T
S_T
ST上的
X
q
X^q
Xq进行预测。然后模型
(
f
T
,
g
,
h
,
σ
)
(f_T,g,h,σ)
(fT,g,h,σ)通过以下端到端的学习目标进行训练:
其中L是损失函数。对于语义分割任务使用交叉熵损失,对其他任务使用L1损失。请注意,目标函数(eq6)并没有明确强制执行令牌空间中的匹配方程(eq3),允许一些预测的知识由标签解码器h处理,因为作者发现在初始实验中引入显式的令牌重构损失会降低性能。在训练过程中,由于我们有固定数量的训练任务 T t r a i n T_{train} Ttrain,我们为每个训练任务保留并训练图像编码器 f T f_T fT的不同偏置参数集(假定进行了通道分割)。
在对
D
t
r
a
i
n
D_{train}
Dtrain进行训练后,模型在给定支持集
S
T
t
e
s
t
S_{T_{test}}
STtest 的新任务
T
t
e
s
t
T_{test}
Ttest上进行少样本评估。我们首先通过使用支持集
S
T
t
e
s
t
S_{T_{test}}
STtest对图像编码器
f
T
f_T
fT的偏置参数进行微调来对模型进行自适应。为此,我们通过随机将支持集划分为子支持集
˜
S
˜S
˜S 和子查询集
˜
Q
˜Q
˜Q 来模拟逐个迭代的元学习,使得
S
T
t
e
s
t
=
˜
S
U
˜
Q
S_{T_{test}} =˜S U ˜Q
STtest=˜SU˜Q 。然后通过以下目标函数对图像编码器
f
T
f_T
fT的偏置参数
θ
T
θ_T
θT进行微调:
其中 θ T θ_T θT表示图像编码器fT的偏置参数。要微调的参数比例很小,因此模型可以避免在小的支持集 S T t e s t S_{T_{test}} STtest上过拟合。微调后,通过使用支持集 S T t e s t S_{T_{test}} STtest对未见查询图像的标签进行预测来评估模型。
5. 实验
数据集
作者构建了Taskonomy数据集(Zamir等,2018)的一个变种,以模拟对未见密集预测任务的少样本学习。Taskonomy包含具有各种注释的室内图像,作者选择了十个具有不同语义和输出维度的密集预测任务:语义分割(SS),表面法线(SN),欧几里得距离(ED),Z缓冲深度(ZD),纹理边缘(TE),遮挡边缘(OE),2D关键点(K2),3D关键点(K3),着色(RS)和主曲率(PC)。将这十个任务分为5个部分,每个部分中有两个任务用于少样本评估(Ttest),剩下的八个任务用于训练(Ttrain)。为了在具有新语义的任务上进行评估,作者仔细构建了划分,使得训练和测试任务彼此之间足够不同,例如通过将边缘任务(TE,OE)组合在一起作为测试任务。此外,作者还将一些单通道任务(ED,TE,OE)处理为多通道以增加任务的多样性,并将所有标签标准化为[0, 1]。
Baseline
将(VTM)与两类学习方法进行比较
- 完全监督Baseline:在训练过程中,完全监督Baseline可以完全访问测试任务Ttest的监督信息,因此可以作为少样本性能的上限。作者考虑了两种最先进的Baseline方法,即 通用密集预测任务的监督学习和多任务学习方法-DPT(Ranftl等,2021)和InvPT(Ye&Xu,2022) ,其中DPT分别独立训练每个单一任务,InvPT联合训练所有任务。
- 少样本学习Baseline:在训练过程中,少样本学习Baseline没有访问测试任务Ttest的权限,只在测试时给出少量标记图像。由于没有先前针对通用密集预测任务开发的少样本方法,作者将最先进的少样本分割方法调整到他们的设置中。作者选择了三种方法,即 DGPNet(Johnander等,2021),HSNet(Min等,2021)和VAT(Hong等,2022) ,它们的架构要么本质上是任务不可知的(DGPNet),要么可以简单扩展(HSNet,VAT)以处理密集预测任务的通用标签空间。
6. 结果
VTM在很大程度上优于所有的FSL Baseline,并且在许多任务上与监督Baseline竞争。
7. 总结
VTM 的本质是一个元学习方法。其训练由多个 episode 组成,每个 episode 模拟一个小样本学习问题。VTM 训练运用到了元训练数据集 D_train,其中包含多种有标签的密集预测任务示例。每个训练 episode 都会模拟数据集中特定任务 T_train 的小样本学习场景,目标是在给定支持集的条件下,为查询图像产生正确的标签。通过多个小样本学习的经验,模型能够学习到通用的知识,以便快速、灵活地适应新的任务。在测试时,模型需要在训练数据集 D_train 中未包含的任意任务 T_test 上进行小样本学习。
总结来说,尽管 VTM 的底层思路非常简单,但它具有统一的体系结构,可用于任意密集预测任务,因为匹配算法本质上包含所有任务和标签结构(例如,连续或离散)。此外,VTM 仅引入了少量的任务特定参数,就能具备抗过拟合性与灵活性。未来研究员们希望进一步探究预训练过程中的任务类型、数据量、以及数据分布对模型泛化性能的影响,从而帮助我们构建一个真正普适的小样本学习器。 [ 1 ] ^{[1]} [1]
参考文献
https://blog.csdn.net/amusi1994/article/details/130460588
https://arxiv.org/abs/2303.14969