标题:基于生成模型的特征知识蒸馏用于动作识别
源文链接:Generative Model-Based Feature Knowledge Distillation for Action Recognition| Proceedings of the AAAI Conference on Artificial Intelligencehttps://ojs.aaai.org/index.php/AAAI/article/view/29473
源码链接: aaai-24/Generative-based-KD (github.com)https://github.com/aaai-24/Generative-based-KD
发表:AAAI-2024
摘要
知识蒸馏(KD),一种在计算机视觉中广泛使用的技术,已经成为提高小型神经网络性能的事实标准。然而,在视频任务中,基于KD的主流方法主要集中在设计损失函数和融合跨模态信息上。这忽视了时空特征语义,导致模型压缩方面的进展有限。为了弥补这一差距,本文介绍了一种创新的基于生成模型的知识蒸馏框架,用于训练轻量级的学生模型。特别地,该框架分为两个步骤:首先是特征表示阶段,其中训练一个基于生成模型的注意力模块来表示特征语义;随后是基于生成的特征蒸馏阶段,包括生成蒸馏和注意力蒸馏,目的是使用生成模型来传递基于注意力的特征语义。通过在多种流行数据集上的综合实验,我们证明了所提方法在视频动作识别任务中的显著效果。此外,我们还在更复杂的视频动作检测任务的背景下验证了所提框架的有效性。我们的代码已公开可用,地址为:https://github.com/aaai-24/Generative-based-KD。
引言
近年来,各种深度学习技术在智能视频分析领域取得了显著的成功(Foo等人,2023;Yang等人,2022a)。特别地,动作识别作为智能视频分析中的一项关键任务,涉及将动作实例分类到相应的标签中。最近,动作识别的性能得到了实质性的提升(Sun等人,2022)。直观地看,一个更大的模型通常对应着更好的性能。这一观点促使许多研究人员设计复杂的骨干网络来捕捉以视频为中心的特征语义(例如,C3D(Xu、Das和Saenko,2017),I3D(Carreira和Zisserman,2017),S3D(Xie等人,2018))。这些努力在动作识别领域取得了令人称赞的结果。然而,部署一个更大的骨干网络引入了极高的资源和内存约束,使得它在资源有限的边缘设备上变得不切实际。为了解决这个问题,Hinton等人提出了知识蒸馏(KD)作为一种解决方案,它促进了从重型模型(教师模型)到轻量级模型(学生模型)的已学知识的转移。
根据模型压缩的基本目标,知识蒸馏(KD)方法主要包括两种研究范式:基于logits的方法和基于特征的方法。基于logits的KD(Hinton, Vinyals, and Dean 2015; Li et al. 2023)涉及将大型教师模型压缩为更紧凑的学生模型,这通过传递来自教师模型输出的软标签(即“暗知识”)来实现。自2015年Romero等人引入该方法以来,比较评估始终验证了在各种任务(Zhao et al. 2022)中基于特征的KD方法的优越性(Xu et al. 2020; Zhang and Ma 2020; Yang et al. 2022c)。因此,学术界的关注明显转向了从模型中间层提取复杂特征中的知识。
然而,大多数基于特征的KD方法(Zhao et al. 2022; Zhang and Ma 2020; Xu et al. 2020)主要集中在设计各种损失函数以匹配特征图,忽视了变分特征语义的内在重要性。实际上,特别是在视频的背景下,特征本身包含强烈的语义信息(Quader et al. 2020),这主要来自于时间和空间的变化。这些特征语义在通过从教师模型转移特征变化来提高学生模型精度方面起着关键作用。遗憾的是,当前的方法往往忽视了这些语义,导致在视频任务中仅有微不足道的改进。因此,当前的KD方法主要适用于与图像相关的任务(Zhao et al. 2022; Yang et al. 2022b; Lin et al. 2022)。在视频动作识别领域,KD方法的主要焦点在于跨模态蒸馏,旨在提高模型的准确性(Liu et al. 2021; Thoker and Gall 2019; Dai, Das, and Bremond 2021)。
因此,在基于知识蒸馏(KD)的方法中,特别是在视频动作分析的背景下,针对压缩3D-CNN模型的研究领域出现了一个明显的空白。为了解决这些问题,我们提出了一种基于生成模型的特征知识蒸馏框架,用于视频动作识别。该框架能够在中间层之间有效地将特征语义从重型模型(教师模型)转移到轻量级模型(学生模型)。具体来说,我们提出的方法包含两个组件。首先,我们设计了一个特征表示模块,该模块利用基于生成的注意力模型在3D-CNN架构中获取特征语义。其次,我们构建了一个基于生成的KD模块,该模块包括生成蒸馏和注意力蒸馏两个过程。这个模块旨在从教师模型中提炼基于注意力的特征信息,如图1所示。生成蒸馏组件被定制为通过匹配重构特征来优化学生模型的特征表示模块,从而促进基于生成的注意力模型的学习。随后,注意力蒸馏过程被协调以通过匹配注意力图来转移基于注意力的特征语义,这以不变的基于注意力的特征分布为条件。
图1:我们提出的框架包括两个阶段:阶段1增加了一个注意力模块来表示特征语义;阶段2构建了一个基于产生式的KD模块来从教师模型中提取特征知识。
据我们所知,我们是首次考虑特征语义的时间变化并研究知识蒸馏(KD)中的生成模型机制的。我们设计了一个基于生成模型的知识蒸馏框架,有效提高了3D卷积神经网络(3D-CNN)模型的压缩性能。我们的主要贡献总结如下:
• 为了提取时空特征,我们设计了一个新颖的基于生成模型的注意力模块,以在3D-CNN架构内表示特征语义。
• 我们构建了一个新框架,引入了利用生成模型来提炼基于注意力的特征的新概念。特别是,我们的知识蒸馏框架是首个在视频上压缩3D-CNN的框架,它使用生成模型来传递时空信息。
• 通过广泛的实验,我们的方法在两个著名的动作识别数据集上展示了在各种网络架构中显著的性能提升。此外,我们还将我们的框架扩展到更复杂的任务——动作检测,这也提供了相当的性能提升。
相关工作
动作识别
大多数之前关于动作识别的卷积神经网络(CNN)架构可以分为两类:3D CNNs(Carreira 和 Zisserman 2017; Xu, Das, 和 Saenko 2017; Ji 等人 2012; Hara, Kataoka, 和 Satoh 2017)和部分 3D CNNs(Tran 等人 2018; Qiu, Yao, 和 Mei 2017; Xie 等人 2018)。3D CNNs 最初在(Carreira 和 Zisserman 2017; Tran 等人 2015)中被提出,它们将视频视为帧的堆叠,通过3D卷积核来学习动作的时空特征。此外,Carreira 和 Zisserman 提出了I3D(Carreira 和 Zisserman 2017),它通过基于3D卷积融合RGB和光流来捕获空间和时间信息。部分 3D CNNs 在(Tran 等人 2018; Qiu, Yao, 和 Mei 2017)中被提出,它们用深度可分离卷积替代了3D卷积以减少资源成本。同时,Xie 等人用2D卷积替代了3D卷积以降低计算复杂度。然而,3D CNNs消耗大量资源以达到高准确率,这使得它们不适合在资源受限的设备上部署。虽然部分 3D CNNs 的使用减轻了资源需求,但与功能完备的3D CNNs相比,它们仍然存在一定的准确率差距。
生成模型
近年来,生成模型(Goodfellow 等人 2014;Kingma 和 Welling 2014)通过与深度学习的结合得到了迅速发展。生成对抗网络(GAN,Goodfellow 等人 2014)最大化生成变量子集和识别网络输出之间的近似真实数据分布信息。然而,GAN 隐式地学习分布并缺乏样本多样性。变分自编码器(VAE,Kingma 和 Welling 2014)通过优化数据边缘似然性的变分下界来近似真实分布。然而,VAE 不适合建模多模态输出的分布(Sohn, Lee, 和 Yan 2015)。为了将注意力转移到指导知识蒸馏上,我们使用条件变分自编码器(Sohn, Lee, 和 Yan 2015)来建模基于注意力值的特征语义分布。
知识蒸馏
知识蒸馏最初由(Hinton, Vinyals, 和 Dean 2015)提出,它通过教师产生的软标签来转移输出概率分布。进一步地,Romero 等人提出了从倒数第二层蒸馏特征表示,称为特征知识蒸馏。Xu 等人提出了特征归一化知识蒸馏,以减少标签噪声的影响。此外,Zagoruyko 和 Komodakis 提出了注意力知识蒸馏,它试图匹配注意力图来传递特征知识。同时,还有将基于 logits 的方法与基于特征的方法相结合的努力(Zhao 等人 2020;Shen 等人 2019),所有这些方法的目标都是提高整体模型性能。然而,这些方法主要集中在对齐特征和注意力图上,因此可能忽略了潜在的特征语义。
方法
在这一部分中,我们首先介绍了特征KD在动作识别任务中的原理。然后,我们给出了KD框架的框架,并介绍了两个子模块。
定义
特征蒸馏 知识蒸馏的理念是训练一个紧凑的学生模型以近似一个复杂教师模型的能力。具体来说,假设我们有一个预训练的教师模型T和一个未训练的学生模型S,它们在这项工作中都被参数化为神经网络。我们用F ∈ RT×C×HW来表示3D-CNN的输出特征图,其中T、C和HW分别代表时间、通道和空间维度。为了更好地说明,我们用FT和FS分别表示教师模型和学生模型某一层的特征图。
对于特征知识蒸馏(KD),学生模型和教师模型之间的蒸馏距离是通过两个特征图来计算的:
其中f(·)是一个明确的映射函数,n是时间维度。学生模型被鼓励最小化目标函数LKD以模仿教师模型。然而,由于这种方式忽略了特征的语义,学生在动作识别中无法从教师那里学习到时序依赖关系。
特征表示 对于动作识别任务,我们首先定义了一个注意力模块来表示特征的语义,这同时也能有效提高识别性能。我们学习了一个生成模型来通过利用特征F和动作类别C来优化注意力λ。为了简化问题,我们使用贝叶斯定理来转换优化目标:
在最后一步中,我们丢弃了常数项log p(F, C, λ)并将λ设为均匀分布。
正如等式(2)所示,当我们优化注意力值时,我们不仅使用特征语义和注意力来提高分类性能(第一项),还确保基于注意力的特征分布与原始特征分布一致(第二项)。
特征知识蒸馏
这项工作采用了特征知识蒸馏来帮助学生模型学习特征表示。我们从训练好的教师模型中提取出以注意力图形式表达的语义,并要求学生模型去模仿它。我们的知识蒸馏框架的流程如图2所示,包含两个步骤(即生成蒸馏和注意力蒸馏)。给定一个视频特征F,在第一步中,教师模型利用一个预训练的生成模型CVAE(条件变分自编码器)来传递特征和注意力的知识,以确保分类的准确性。在第二步中,教师会产生一个注意力图λ,它代表了教师模型的特征语义。
图2:我们知识蒸馏(KD)框架的流程。我们的KD模型通过两个阶段交替进行训练。在第一阶段,生成蒸馏通过自重构损失LCVAE和生成蒸馏损失LKD−gen进行训练,以匹配重构特征。在第二阶段,注意力蒸馏通过表示损失Lrecon、注意力蒸馏损失LKD−att和分类损失Lclf进行更新,以提炼基于注意力的特征语义。
生成蒸馏:为了重构特征语义,我们引入了一个生成蒸馏模块,通过从教师模型转移基于注意力的特征来实现。此外,条件变分自编码器(CVAE)通过学习特征图和注意力图之间的对应关系来有效地重构特征语义。具体来说,如图2(a)所示,我们冻结了学生模型和注意力模块,利用教师模型的CVAE来更新CVAE模块。为了更新CVAE,我们构建了一个基于动作的生成问题:
其中zt是潜在变量,ϕ表示可学习的模型,pϕ(zt|λt)表示先验模型,而pϕ(Ft|λt, zt)是生成模型的后验模型,即解码过程。值得注意的是,潜在变量zt是从学习到的先验分布中采样的,这被设置为特征重构的过程。
在训练过程中,我们的目标是让学生模型学习教师模型重构的特征,这旨在拟合教师模型的分布并提高学生模型的重构质量。具体来说,我们提出了生成损失LKD−gen来优化特征重构:
其中S和T分别表示学生模型和教师模型,F表示特征图,A是注意力图,f(·)是基于注意力的重构特征函数。值得注意的是,我们在教师模型和学生模型的CVAE模块中设置了相同的输入,即学生模型的特征图和注意力图。相同的输入是为了让学生专注于教师基于注意力的特征重构,而不受不同模型提取的特征和注意力图差异的影响。在CVAE训练过程中,如图2(a)所示,我们的目标是最小化证据下界(ELBO)损失Lcvae:
其中zt是潜在变量,α=0.1是一个经验超参数。
具体来说,如图3(a)所示,首先,我们使用编码器通过特征Ft和注意力λt生成中间变量zt,函数fz如下:
其中ffci是线性模型,θi是一个1×1×1的3D卷积,旨在通过减少通道数并添加可学习模型参数来为fz连接λi。fz是中间变量z的均值和方差函数,用于生成潜在变量Z。然后,我们使用解码器生成重构的特征图Fb,它利用潜在变量Z和注意力λ,如下:
其中θ′i是一个1×n×n(n是核大小)的3D反卷积,旨在通过增加通道数来生成重构的特征图。我们使用Lcvae来衡量和F之间的差异。
注意力蒸馏 如公式2所示,生成蒸馏(Generative Distillation)主要优化第一项logp(C∣F,λ),而本节主要关注第二项logp(λ∣F)。注意力(Attention)代表了视频的特征分布,因此学生模型可以通过学习注意力图来学习教师模型的表示能力(Quader 等人 2020;Zagoruyko 和 Komodakis 2017)。
具体来说,如图2(b)所示,我们冻结了条件变分自编码器(CVAE),通过利用注意力蒸馏来更新学生模型和注意力模块。
对于视频特征F,我们使用注意力A来表示特征的语义,如下所示:
其中,GN表示组归一化(Group Normalization)。
进一步,我们计算基于注意力的特征F′,如下所示:
其中,θ(⋅)表示3D转置卷积(3D transposed convolution),σ是缩放因子,定义为:
如图2(b)所示,为了蒸馏注意力,我们提出了注意力蒸馏损失LKD−att,定义为:
其中,AT和AS分别表示教师模型和学生模型的注意力图。
在训练过程中,对于动作识别,我们鼓励提高动作分类的能力。这相当于最小化以下损失:
其中,yc(x)和pc(x)分别表示真实标签和预测结果的动作概率分布。
同时,如图3(b)所示,我们使用解码器(Decoder)来生成重构的特征图,其目标是在特征表示过程中保持基于注意力的特征的分布。这相当于最小化重构损失Lrecon:
其中,zt(l)是由生成模型的编码器中的zt和λt生成的。特别地,按照(Shi等人 2020)的做法,在最后一步,我们设定L为1。
训练与推理 在我们提出的框架中,注意力模块需要基于注意力的特征作为输入。因此,如图2所示,我们考虑交替训练注意力模块和主干模块。如图2(a)所示,在第一阶段,我们冻结学生主干模型并激活CVAE模块,以通过从主干模型提取的特征和生成的注意力来优化生成蒸馏过程。第一阶段的整体损失函数如下:
其中,β是一个经验超参数,我们设定β=0.01。
如图2(b)所示,在第二阶段,我们冻结CVAE模块并激活学生主干模型,以通过基于注意力的重构特征、注意力分布和预测的动作概率来优化注意力蒸馏过程,如下:
其中,γ是一个经验超参数,我们设定γ=0.1。
值得注意的是,在推理过程中,我们只利用带有注意力模块的学生模型来预测动作标签,不包括生成模型。
实验
结论
在本文中,我们提出了一种新颖的基于生成模型的知识蒸馏框架,通过注意力表示来传递视频特征语义。我们的框架主要包括两个步骤:特征表示和基于生成的知识蒸馏(特征蒸馏和注意力蒸馏)。通过利用注意力机制,我们捕获特征语义并通过生成建模实现知识蒸馏。在3D卷积神经网络(3D-CNNs)蒸馏领域,我们的方法相较于现有方法在视频动作识别领域取得了显著的性能提升。结果证明了我们的方法在通过基于生成模型的特征蒸馏来增强知识蒸馏性能方面的有效性。因此,本研究为专注于3D模型蒸馏的创新知识蒸馏框架奠定了基础。