论文笔记《Self-distillation with Batch Knowledge Ensembling Improves ImageNet Classification》
Self-distillation with Batch Knowledge Ensembling Improves ImageNet Classification
- 2021.5.13
- Project Page: https://geyixiao.com/projects/bake
- https://arxiv.org/abs/2104.13298
Introduction
- 主要目标在于给batch内的每一个作为anchor的图片通过传播同一个batch内其他样本的知识生成细化的软标签refined soft label，所传播的知识为batch内样本之间的相似度
- 该方法基于这样的假设：外观相似的样本应有更一致的类别预测
- 实际中对于每个样本，将batch内别的样本的预测结果或通过加权传播的方式形成软目标
- 知识的传播是经过迭代处理的，直到其收敛
- 折是第一个没有用多个网络或额外的分支来生成 ensembled soft target的自蒸馏方法
- 所提出的batch knowledge ensembling使用样本之间的知识生成惊喜的蒸馏目标
Method
-
具体的实现方法其实通过伪代码可以非常清晰的理解，这边还是介绍一下具体过程，整体框架图如下
-
对于我们要传播的label知识，首先是要对batch内所有的样本进行相似性计算，生成相似性矩阵\(A\in \mathbb{R}^{N\times N}\)，相似性计算后，去除对角线删去自己与自己的相似性，然后进行一个归一化，对于每一个样本的相似性向量和各元素和为1，记作\(\hat{A}\)
-
将原先所预测出每个样本所对应的logit记作\(P^{\tau}\)，然后将上一步计算出的相似性矩阵与之相乘，相当于利用相似性作一个加权\(\hat{P^{\tau}}=\hat{A}P^{\tau}\)
-
对于相似性传播过来的label，我们也要进行一个加权，相当于是得到了我们想要的细化的logit：
-
这样的知识传播需要进行数次，直到收敛，这时候公式中t表示第t次传播与迭代
- 当我们t趋于无穷大时，我们对此求极限，相当于是一个等比级数的极限，证明也很简单，值得注意的是上式中第一项极限为0
- 所以最后我们的知识加权传播模块最终可以表示成如下，该公式在下面的代码中也有所体现，算是真正得到了我们想要的吸取了batch内所有其他样本后的logit，值得注意的是对于每一个样本的refined logit和刚好为1，所以可以直接用：
- 最后就对于原本的logit和我们细化的logit之间做一次KL loss，加上一定的权重后和原本的CE loss成为本自蒸馏项目全部的loss
- 值得注意的是，在这个工作中有非常重要的一点，因为logit的细化蒸馏主要依赖于相似性，在一个batch内如果没有相似的样本其实是本方法是无效的，所以我们引入了一个对每类都采样的机制，对于batch大小为\(N\)内有一张图片后随机选取同类的\(M\)张图放入同一个batch中，组成新的batch，这时新的batch大小为\(N\times (M+1)\)
# w: ensembling weight
# t: temperature
# r: loss weight
for (x, gt_labels) in loader:
# features: N×D, logits: N×K 分别是embedding特征和logit
f, logits = net.forward(x)
# classification loss with ground-truth labels
loss = CrossEntropyLoss(logits, gt labels)
# produce soft targets
f = normalize(f)
# 计算batch内各样本之间的相似度并去除中间的自己与自己，进行一个softmax变成0-1之间，得到公式中的A
A = softmax(mm(f, f.t())-eye(N)*1e-9) # row-wise normalization of affinity matrix with zero diagonal
# 最后求过极限之后得到的公式 得到soft_target
soft_targets = mm((1-w)·inv(eye(N)-w·A),softmax(logits/t)) # approximate inference for propagation and ensembling
soft_targets = soft_targets.detach() # no gradient
# distillation loss with soft targets 两个target之间的KL loss
loss += KLDivLoss(log_softmax(logits/t), soft_targets)*t^2*r
# SGD update
loss.backward()
update(net.params)
Experiments & Result
-
做了很多实验来证明其有效性，首先给出了训练的细节，如\(N=256,M=1,lr = base\_lr×batch\_size/256\)
-
首先是不同架构下与原baseline之间的差距和别的label regularzation方法和别的self-distillation之间的区别，常规实验对比
-
和别的ensembel distillation方法之间的对比
-
和别的label refinery方法之间的对比
-
Transfer learning下游任务上目标检测结果
-
鲁棒性测量实验结果
-
每类数据采样方法的实验，这个实验很重要，因为它证明了BAKE方法效果好的原因还是在于knowledge ensemble而不是采样方法，因为可以看到在正常情况下采用这种采样方法反而会使效果下降，可能是因为这导致了同一个batch内多样性下降，而且也并不是同一batch内相同样本越多越好
-
小数据集上的实验结果
Conclusion
- 一种全新的batch knowledge ensemble方法，为自蒸馏生成了refined soft target，不过这也是建立在一定的采样方法基础之上的，虽然该方法还挺有意思的，但受限于这个条件显得就没有那么厉害了，因为蒸馏中利用batch之内样本的相似性来作文章真的挺多了，但这个工作是用来生成新的logit，所以我个人感觉还是挺有意思的，而且这篇文章的算法过程描述的非常清楚了很容易就懂。但目前还并不知道这篇文章中了没有，其实其对比的自蒸馏方法还是相对来说比较少的，不知道最后结果如何，感谢作者的工作给我带来的启发。
