超越BN层的归一化方法FRN

BN层作为深度神经网络的必选组件在各种网络结构中得到了广泛的应用。然而训练输入的batchsize过小时，尤其是没有足够计算资源可用，不得不使用小batchsize来训练模型时，你会发现使用同样的方法和训练策略和别人有足够计算资源训练出来模型的表现差的不是一点半点。这一点我在自己设计（抄别人的网络。。。使用别人的训练策略）网络结构跑实验的时候也有发现在batch较小（显存不够办法呀。。。。）的情况下bn层没啥效果。还有一个就是评价和训练的结果相差较远。所以google设计一个非常简单但效果非常好的不依赖batch的归一化方法FRN。

FRN

输入特征图$[B,W,H,C]$四个维度中FRN仅涉及其中的$WH$两个维度，这样就避免了Batch和Channel对归一化结果的影响，即无论batch大小如何都不会影响FRN的表现。

因此在归一化时只需关注上图中红色部分即可，其它通道都使用相同的计算方式。设红色部分为向量$x$ 则FRN归一化方法为：

• 先求向量$x$的均方$v^2$：
  $$v^2=\sum _i\frac{x_i^2}{N}(N\in W\times H)$$

• 然后正则化$x$:
  $$\hat{x}=\frac{x}{\sqrt{v^2+\epsilon }}$$
  其中$\epsilon$为一个极小的数防止分母为0。

在上述计算之后使用仿射变换抵消因normal带来尺度变化的影响。仿射变换公式为：
$$y=r\hat{x}+\beta$$
其中$r$和$\beta$都为可学习的参数。

TLU

然而由于FRN缺乏对均值中心的处理，因此在激活函数计算结果之后相对于0会产生一个绝对误差。为了解决这个问题，作者使用一个带有可学习截断值$\tau$增强型ReLU。
$$\mathbf{z}=\mathbf{m}\mathbf{a}\mathbf{x}(\mathbf{y},\tau )$$
因为$z=max(y,\tau )=max(y-\tau ,0)+\tau =ReLU(y-\tau )+\tau$所以理论上来说TLU激活函数的作用应该和在使用ReLU之前和之后带有偏置项的函数作用相同（ the effect of TLU activation is the same as having a shared bias before and after ReLU），但是实验结果表明TLU形式的激活似乎更好，作者把这归咎于TLU形式的激活函数更容易优化。
以下是tensorflow的代码实现：

def FRNLayer(x, tau, beta, gamma, eps=1e-6):
    # x: Input tensor of shape [BxHxWxC].
    # alpha, beta, gamma: Variables of shape [1, 1, 1, C].
    # eps: A scalar constant or learnable variable.
    # Compute the mean norm of activations per channel.
    nu2 = tf.reduce_mean(tf.square(x), axis=[1, 2], keepdims=True)
    # Perform FRN.
    x = x * tf.rsqrt(nu2 + tf.abs(eps))
    # Return after applying the Offset-ReLU non-linearity.
    return tf.maximum(gamma * x + beta, tau)

实验对比

最后就是使用不同激活函数在网络模型上的效果对比啦：

从实验结果来说FRN层还是非常优秀的，这对计算资源不够只能使用小batch跑实验的一些同学来说真的是一个不错的技术。