GHM论文理解及实现

原论文：Gradient Harmonized Single-stage Detector

本文主要基于tf.keras讨论分类部分，论文也提出了适用于检测的方法。 实验表明具有一定效果，可以尝试，感觉比focal loss要好用。

概述

本文针对问题：one-stage的目标检测算法一直存在的正负/难易样本（样本梯度）失衡问题。

在one-stage算法中，负样本的数量要远远大于正样本，而且大多数负样本是简单样本（well-classified）；当然，大多正样本也算是简单样本。简单样本的小梯度，通过样本量的增加，量变引起质变，主导了模型的训练过程。focal loss中通过引入$\alpha$大大降低简单样本的分类损失，以平衡正负/难易样本，但是设计的损失函数$L$引入两个超参(OHEM只学习困难/loss大的一部分样本，训练有效性降低）：

$$L = - \sum_{i}^{c} ( \alpha * | p_{i} - p_{i}^{*} |^{\gamma} * p_{i}^{*} * \ln(p_{i}))$$

def focal_loss(y_truth, y_pred, _):
    epsilon = keras.backend.epsilon()
    y_pred = keras.backend.clip(y_pred, epsilon, 1.0 - epsilon)
    cross_entropy = -y_truth * keras.backend.log(y_pred)
    # 这里没有区分正负样本的alpha
    weight = alpha * keras.backend.pow(keras.backend.abs(y_truth - y_pred), gamma)
    loss = weight * cross_entropy
    loss = keras.backend.sum(loss, axis=1)
    return loss

本论文认为，模型训练过程中，信息/知识的传导靠的是梯度；只不过梯度的大小，表象刚好是样本的难易/正负。也就是说，正负/难易样本的失衡，其实是梯度的失衡！所以想让模型学得更好，应该从更本质的梯度入手。故提出gradient harmonizing mechanism (GHM)，拟解决模型中的梯度失衡问题。

我的理解

梯度调和机制（gradient harmonized mechanism, GHM）其实就是将不同梯度对损失函数的影响，进行基于密度的平衡。（如果为了简单了解，我们暂且可以把密度理解为样本数量）

调和机制（harmonized mechanism）

这里把harmonized mechanism翻译为调和机制，其实是调和有取倒数的意思。

基于类别平衡的损失（class-balanced loss）

可以从样本的类别不平衡来理解：如果有10个类$A$的样本，2个类$B$的样本，那么基于类别平衡的损失为：

$$L = \frac{12}{10} * \sum_{i \in A}(L_{i}) + \frac{12}{2} * \sum_{i \in B}(L_{i})$$

其中，10可以理解为类别$A$的密度，2可以理解为类别$B$的密度。

推广到多个类别可得：

$$L = \sum_{i}^{N}( \frac{N}{M_{i}} * L_{i} )$$

其中，$N$是样本总数，$M_{i}$是样本$i$所属的类别的密度，$L_{i}$是样本$i$的损失。以上如果取平均损失函数，则需要乘上$\frac{1}{N}$。

基于梯度平衡的损失（gradient-balanced loss）

这里，我们将基于类别的平衡推广到基于梯度的平衡：

$$L = \sum_{i}^{N}( \frac{N}{GD_{g_{i}}} * L_{g_{i}} )$$

其中，$N$是样本总数，$g_{i}$是样本$i$的梯度，$GD_{g_{i}}$是梯度$g_{i}$的密度，$L_{g_{i}}=L_{i}$是样本$i$的损失。

也就是说，用样本梯度的密度取倒数，乘上样本损失，便可以平衡不同梯度区域的损失。

为什么需要调和

基于类别的平衡，在实际的数据分析场景中，也要case-by-case分析是否适用。那么，基于梯度的平衡是否科学呢？作者给出了下图的解释（图中还包含了cross-entropy和focal-loss的梯度加和方式）：

最左边的图是样本的梯度分布，梯度小的表示已经被模型学习到的了（容易样本），梯度大的表示模型很难学到（困难样本）。这两种样本，梯度密度都比较大，主导了整个模型的训练方向。而作者认为，学到的了可以不用学了，没有学到的可能是异常样本，也不用学了；我们这时候应该提升模型，学习中间那段梯度密度较小、还有信息可以学习的样本。

求解出整个样本集梯度的概率密度的调和曲线（1/概率密度，中间那张图），来调和原始的梯度（也就是与左图相乘），得到最终的模型回传梯度（最右边的图）。

和focal loss和OHEM的比较来看，改进源于更本质的角度——模型的梯度（当然，这里的梯度并不是整个模型的所有参数$W$构成的梯度，而是简化为最后一层sigmoid回传的梯度；同时也不是指梯度向量，而是梯度向量的L1范数）。

以上，便是论文的核心理解，接下来掰一掰公式和实现。

Gradient Harmonizing Mechanism —— GHM-C Loss

原论文主要针对的是sigmoid二分类情况，我泛化为softmax来分析。

回过头来看一下GHM的损失函数为：

$$L = \frac{1}{K} * \sum_{i}^{K}( \frac{N}{GD_{g_{i}}} * L_{g_{i}} )$$

其中，$K$是batch size，最重要的部分在于求解梯度密度$GD_{g_{i}}$上，

论文通过两个机制来近似这个梯度密度：

1. 将梯度取值区间(0, 1)切割为多个bin，统计不同bin的梯度数量R，作为梯度密度（论文中梯度密度调和参数是$\beta_{i}=\frac{N}{GD_{g_{i}}},GD_{g_{i}}=\frac{R_{i}}{\epsilon},\epsilon=\frac{1}{bin}$）；
2. 用逐batch的指数加权移动平均(EMA)来近似总样本下的梯度密度$S_{i}^{t}=\alpha S_{i}^{t-1}+(1-\alpha)R_{i}^{t}$。

cross-entropy的损失函数是：

$$L = -\sum_{i}^{c}(p_{i}^{*} * \ln(p_{i}))$$

其中，$p_{i}^{*}$指真实类别概率，预测类别概率为$p_{i}=\frac{e^{z_{i}}}{\sum_{j}e^{z_{j}}}$； 代入计算梯度得 （与论文中针对sigmoid的推导结果是一致的）：

$$\frac{\partial L}{\partial x_{i}} = p_{i} - p_{i}^{*}$$

综上，可得求解GHM-C Loss：

1. $|\frac{\partial L}{\partial x_{i}}| = |p_{i} - p_{i}^{*}|$；
2. 统计$R(g)$（$g=|p-p^{*}|$ 所在梯度区间$[(i-1)*\epsilon, i*\epsilon]$的样本数）；
3. 指数加权移动平均计算$S(g)$；
4. 计算梯度密度 $GD(g) = \frac{S(g)}{\epsilon} = S(g) * M$；
5. 计算损失 $L = \sum_{i}(\frac{I(x_{i})}{GD(g_i)})$。

def _categorical_ghm_loss(bins=30, momentum=0.75):
    """ 返回多分类 GHM 损失函数：
            把每个区间上的梯度做平均，也就是说把梯度拉平，回推到公式上等价于把loss做平均
    Formula:
        loss = sum(crossentropy_loss(p_i,p*_i) / GD(g_i))
        GD(g) = S_ind(g) / delta = S_ind(g) * M
        S_ind(g) = momentum * S_ind(g) + (1 - momentum) * R_ind(g)
        R_ind(g)是 g=|p-p*| 所在梯度区间[(i-1)delta, i*delta]的样本数
        M = 1/delta，这个是个常数，理论上去掉只有步长影响
    Parameters: （论文默认）
        bins -- 区间个数，default 30
        momentum -- 使用移动平均来求区间内样本数，动量部分系数，论文说不敏感
    """
    # 区间边界
    edges = np.array([i/bins for i in range(bins + 1)])
    edges = np.expand_dims(np.expand_dims(edges, axis=-1), axis=-1)
    acc_sum = 0
    if momentum > 0:
        acc_sum = tf.zeros(shape=(bins,), dtype=tf.float32)

    def ghm_class_loss(y_truth, y_pred, valid_mask):
        epsilon = keras.backend.epsilon()
        y_pred = keras.backend.clip(y_pred, epsilon, 1.0 - epsilon)
        # 0. 计算本次mini-batch的梯度分布：R_ind(g)
        gradient = keras.backend.abs(y_truth - y_pred)
        # 获取概率最大的类别下标，将该类别的梯度做为该标签的梯度代表
        # 没有这部分就是每个类别的梯度都参与到GHM，实验表明没有这部分会更好些
        # truth_indices_1 = keras.backend.expand_dims(keras.backend.argmax(y_truth, axis=1))
        # truth_indices_0 = keras.backend.expand_dims(keras.backend.arange(start=0,
        #                                                                  stop=tf.shape(y_pred)[0],
        #                                                                  step=1, dtype='int64'))
        # truth_indices = keras.backend.concatenate([truth_indices_0, truth_indices_1])
        # main_gradient = tf.gather_nd(gradient, truth_indices)
        # gradient = tf.tile(tf.expand_dims(main_gradient, axis=-1), [1, y_pred.shape[1]])

        # 求解各个梯度所在的区间，并落到对应区间内进行密度计数
        grads_bin = tf.logical_and(tf.greater_equal(gradient, edges[:-1, :, :]), tf.less(gradient, edges[1:, :, :]))
        valid_bin = tf.boolean_mask(grads_bin, valid_mask, name='valid_gradient', axis=1)
        valid_bin = tf.reduce_sum(tf.cast(valid_bin, dtype=tf.float32), axis=(1, 2))
        # 2. 更新指数移动平均后的梯度分布：S_ind(g)
        nonlocal acc_sum
        acc_sum = tf.add(momentum * acc_sum, (1 - momentum) * valid_bin, name='update_bin_number')
        # sample_num = tf.reduce_sum(acc_sum)  # 是否乘以总数，乘上效果反而变差了
        # 3. 计算本次mini-batch不同loss对应的梯度密度：GD(g)
        position = tf.slice(tf.where(grads_bin), [0, 1], [-1, 2])
        value = tf.gather_nd(acc_sum, tf.slice(tf.where(grads_bin), [0, 0], [-1, 1]))  # * bins
        grad_density = tf.sparse.SparseTensor(indices=position, values=value,
                                              dense_shape=tf.shape(gradient, out_type=tf.int64))
        grad_density = tf.sparse.to_dense(grad_density, validate_indices=False)
        grad_density = grad_density * tf.expand_dims(valid_mask, -1) + (1 - tf.expand_dims(valid_mask, -1))

        # 4. 计算本次mini-batch不同样本的损失：loss
        cross_entropy = -y_truth * keras.backend.log(y_pred)
        # loss = cross_entropy / grad_density * sample_num
        loss = cross_entropy / grad_density
        loss = keras.backend.sum(loss, axis=1)
        """
        # 调试用，打印tensor
        print_op = tf.print('acc_sum: ', acc_sum, '\n',
                            'grad_density: ', grad_density, '\n',
                            'cross_entropy: ', cross_entropy, '\n',
                            'loss:', loss, '\n',
                            '\n',
                            '=================================================\n',
                            summarize=100)
        with tf.control_dependencies([print_op]):
            return tf.identity(loss)
        """
        return loss
    return ghm_class_loss

Gradient Harmonizing Mechanism —— GHM-R Loss

一般的回归损失是：

$$SL_{1}(d) = \left\{\begin{matrix} \frac{d^{2}}{2\delta} & if |d| <= \delta \\ |d| - \frac{\delta}{2} & otherwise \end{matrix}\right.$$

这样得到的梯度为：

$$\frac{\partial SL_{1}}{\partial d} = \left\{\begin{matrix} \frac{d}{\delta} & if \; |d| <= \delta \\ sgn(d) & otherwise \end{matrix}\right. \Rightarrow \vee |d| > \delta, |\frac{\partial SL_{1}}{\partial d}| = 1$$

由于大部分为梯度为1，没法计算梯度密度。故改进回归损失为：

$$ASL_{1}(d) = \sqrt{(d^{2} + \mu_{2})} - \mu \Rightarrow \frac{\partial ASL_{1}(d)}{\partial d} = \frac{d}{\sqrt{(d^{2} + \mu_{2})}}$$

然后按照 GHM-C Loss 的步骤计算 GHM-R Loss 即可（按经验$u=0.02$）。