Softmax，Label Smoothing原理与其NVCaffe实现

背景

对于多分类的任务，最后通常会连接一个全连接层作为分类器，输出每个类别预测出的confidence scores。这些scores会被softmax归一化为一个概率分布，表示每一个类别的预测概率。假设共有K个类别，对类i，网络输出的confidence score为$z_i$，那么softmax得到的输出为$q_i=softmax(z_i)$，其中softmax的计算定义为：
$$q_i=\frac{exp(z_i)}{{\sum }_{j=1}^{K}exp(z_j)}$$

显然，有$q_i>0$和${\sum }_{i=1}^K{q}_i=1$，这是一个有效的概率分布。

假设图片的真实标签为$y$，我们通常会认为真实的概率分布应该为$p_i=1ifi=y$否则$p_i=0$。在训练时，就会最小化预测概率分布与真实概率分布的交叉熵：
$$l(p,q)=-\sum _{i=1}^K{p}_{i}log{q}_i$$

为了最小化交叉熵，最优的预测概率分布应为$z_y^{\ast }=inf$，且$z_i^{\ast }\to 0(i̸=y)$，也就是说，神经网络的输出会被鼓励尽量去增加正确标签与错误标签的输出值之差，这样就容易出现过拟合问题。

Label smoothing

Label smoothing最早在$Inceptionv2$中被提出，作为一种正则化手段，可以提高网络的top-1和top-5各0.2%

Label;smoothing将认为的真实概率分布改变为
$$p_i=\{\begin{array}{ll}1-ϵ& \text{if i = y}\\ ϵ/(K-1)& \text{otherwise}\end{array}$$
这里的$ϵ$代表一个非常小的常数。

现在，最优的预测概率分布变成了
$$z_i^{\ast }=\{\begin{array}{ll}log((K-1)(1-ϵ)/ϵ+\alpha )& \text{if i = y}\\ \alpha & \text{otherwise}\end{array}$$
其中$\alpha$为一个任意实数，这样，神经网络正确标签与错误标签的输出值之差将会被鼓励输出在一个固定值之内，帮助神经网络更好地泛化。

如果把真实分布概率改为一个性能更优的CNN的输出概率，那么又有点类似知识蒸馏了，当然知识蒸馏不是把student的输出概率趋于teacher，而是求两个网络输出的logits的KL散度，具体的就不再多说了。

前向传播

在这里我们以NVCaffe中的softmax_loss_layer.cpp文件为例，介绍Label smoothing的前向传播与反向传播，当然Caffe的实现也是类似的。
在Caffe与NVCaffe中，softmax_loss_layer包含了求softmax与求loss两部分。

首先定义两个成员变量inner_num_与outer_num_

outer_num_ = bottom[0]->count(0, softmax_axis_);
inner_num_ = bottom[0]->count(softmax_axis_ + 1);

假设$blob$的维度为$(N,C,H,W)$，softmax_axis_ =1
那么outer_num_求的就是$N$，inner_num_求得就是特征图的大小$H\ast W$
outer_num_用于在softmax_axis_之外遍历，inner_num_用于在softmax_axis_之内遍历，这样既可以用于图像分类求loss，也可以用于图像分割在特征图上求loss

在Forward_cpu中，首先得到softmax计算出的预测概率分布

softmax_layer_->Forward(softmax_bottom_vec_, softmax_top_vec_);
const Ftype* prob_data = prob_->template cpu_data<Ftype>();

再定义一个常量，表示当前softmax_axis_上的通道数，也即分类任务的类别个数

int num_class = bottom[0]->shape(softmax_axis_) - 1;  // defined in softmax_loss_layer.cpp, equals prob_->count() / outer_num_ / inner_num_
const int channels = dim / spatial_dim;  // defined in softmax_loss_layer.cu, spatial_dim is inner_num_

根据softmax公式与Label smoothing定义的真实概率分布，有
$$\begin{array}{rl}l(p,q)& =-\sum _{i=1}^K{p}_{i}log{q}_i\\ & =-\frac{ϵ}{K-1}\sum _{i=1,i̸=y}^{K}log{q}_i-(1-ϵ)log{q}_y\end{array}$$
接下来，只需要遍历softmax_axis_中的每一个通道，判断该通道代表的值是否为正确的标签，选择不同的权重去乘以log的结果

for (int i = 0; i < outer_num_; ++i) {
  for (int j = 0; j < inner_num_; j++) {
    const int label_value = static_cast<int>(label[i * inner_num_ + j]);
    if (has_ignore_label_ && label_value == ignore_label_) {
      continue;
    }
    DCHECK_GE(label_value, 0);
    DCHECK_LT(label_value, prob_->shape(softmax_axis_));
    if (label_smooth_ > 0.F) {
      for (int c = 0; c <= num_class; c++) {
        Ftype coeff = (c == label_value) ? (1.F - label_smooth_) : (label_smooth_ / float(num_class));
        loss -= coeff * log(std::max(prob_data[i * dim + c * inner_num_ + j],
          min_dtype<Ftype>()));
      }
    } else {
      // min_dtype<Ftype>() get the minimum value of Ftype, 
      // BVLC Caffe use Dtype(FLT_MIN) instead, where Dtype(FLT_MIN) = 1.17549435E-38F
      // ref : https://blog.csdn.net/jkfdqjjy/article/details/52268565?locationNum=14
      loss -= log(std::max(prob_data[i * dim + label_value * inner_num_ + j],
          min_dtype<Ftype>()));  
    }
    ++count;
  }
}

反向传播

首先，我们要继续推导前向传播的损失函数，把softmax的公式带入(2)式，得到
$$\begin{array}{rl}l(p,q)& =-\frac{ϵ}{K-1}\sum _{i=1,i̸=y}^{K}log{q}_i-(1-ϵ)log{q}_y\\ & =-\frac{ϵ}{K-1}\sum _{i=1,i̸=y}^{K}log\frac{e^{z_i}}{{\sum }_j{e}^{z_j}}-(1-ϵ)log\frac{e^{z_y}}{{\sum }_j{e}^{z_j}}\\ & =-\frac{ϵ}{K-1}\sum _{i=1,i̸=y}^K(z_i-log\sum _j{e}^{z_j})-(1-ϵ)(z_y-log\sum _j{e}^{z_j})\\ & =-\frac{ϵ}{K-1}\sum _{i=1,i̸=y}^K{z}_i-(1-ϵ)z_y+log\sum _j{e}^{z_j}\end{array}$$
对其求偏导数，当$i̸=y$时，有
$$\begin{array}{rl}\frac{\partial l}{\partial z_i}& =-\frac{ϵ}{K-1}+\frac{\partial log{\sum }_j{e}^{z_j}}{\partial z_i}\\ & =-\frac{ϵ}{K-1}+\frac{e^{z_i}}{{\sum }_j{e}^{z_j}}\\ & =softmax(z_i)-\frac{ϵ}{K-1}\end{array}$$
当i=y时，有
$$\begin{array}{rl}\frac{\partial l}{\partial z_y}& =-(1-ϵ)+\frac{\partial log{\sum }_j{e}^{z_j}}{\partial z_y}\\ & =-(1-ϵ)+\frac{e^{z_y}}{{\sum }_j{e}^{z_j}}\\ & =softmax(z_y)-(1-ϵ)\end{array}$$
因此，梯度的计算与前向时类似，遍历所有通道并把计算出的预测概率分布值减去其真实概率分布即可，部分代码如下：

for (int c = 0; c <= num_class; c++) {
  int index = i * dim + c * inner_num_ + j;
  float coeff = (c == label_value) ? (1.F - label_smooth_) : (label_smooth_ / num_class);
  bottom_diff[index] -= coeff;
}
++count;

当然，当$ϵ=0$时，得到的梯度公式与没有Label smoothing时的梯度公式是一样的。