论文笔记《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内所有的样本进行相似性计算，生成相似性矩阵(Ain mathbb{R}^{N imes N})，相似性计算后，去除对角线删去自己与自己的相似性，然后进行一个归一化，对于每一个样本的相似性向量和各元素和为1，记作(hat{A})
将原先所预测出每个样本所对应的logit记作(P^{ au})，然后将上一步计算出的相似性矩阵与之相乘，相当于利用相似性作一个加权(hat{P^{ au}}=hat{A}P^{ au})
对于相似性传播过来的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 imes (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 reﬁnery方法之间的对比
Transfer learning下游任务上目标检测结果
鲁棒性测量实验结果
每类数据采样方法的实验，这个实验很重要，因为它证明了BAKE方法效果好的原因还是在于knowledge ensemble而不是采样方法，因为可以看到在正常情况下采用这种采样方法反而会使效果下降，可能是因为这导致了同一个batch内多样性下降，而且也并不是同一batch内相同样本越多越好
小数据集上的实验结果

Conclusion

一种全新的batch knowledge ensemble方法，为自蒸馏生成了refined soft target，不过这也是建立在一定的采样方法基础之上的，虽然该方法还挺有意思的，但受限于这个条件显得就没有那么厉害了，因为蒸馏中利用batch之内样本的相似性来作文章真的挺多了，但这个工作是用来生成新的logit，所以我个人感觉还是挺有意思的，而且这篇文章的算法过程描述的非常清楚了很容易就懂。但目前还并不知道这篇文章中了没有，其实其对比的自蒸馏方法还是相对来说比较少的，不知道最后结果如何，感谢作者的工作给我带来的启发。