Gradient Centralization: 一行代码加速训练并提升泛化能力

加入极市专业CV交流群，与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流！
同时提供每月大咖直播分享、真实项目需求对接、干货资讯汇总，行业技术交流。关注 极市平台 公众号，回复 加群，立刻申请入群~

来源｜晓飞的算法工程笔记

梯度中心化GC对权值梯度进行零均值化，能够使得网络的训练更加稳定，并且能提高网络的泛化能力，算法思路简单，论文的理论分析十分充分，能够很好地解释GC的作用原理。

论文: Gradient Centralization: A New Optimization Technique for Deep Neural Networks

论文地址：https://arxiv.org/abs/2004.01461
论文代码：https://github.com/Yonghongwei/Gradient-Centralization

Introduction

优化器(Optimizer)对于深度神经网络在大型数据集上的训练是十分重要的，如SGD和SGDM，优化器的目标有两个：加速训练过程和提高模型的泛化能力。目前，很多工作研究如何提高如SGD等优化器的性能，如克服训练中的梯度消失和梯度爆炸问题，有效的trick有权值初始化、激活函数、梯度裁剪以及自适应学习率等。而一些工作则从统计的角度对权值和特征值进行标准化来让训练更稳定，比如特征图标准化方法BN以及权值标准化方法WN。。

与在权值和特征值进行标准化方法不同，论文提出作用于权值梯度的高性能网络优化算法梯度中心化(GC, gradient centralization)，能够加速网络训练，提高泛化能力以及兼容模型fine-tune。如图a所示，GC的思想很简单，零均值化梯度向量，能够轻松地嵌入各种优化器中。论文主要贡献如下：

提出新的通用网络优化方法，梯度中心化(GC)，不仅能平滑和加速训练过程，还能提高模型的泛化能力。
分析了GC的理论属性，表明GC能够约束损失函数，标准化权值空间和特征值空间，提升模型的泛化能力。另外，约束的损失函数有更好的Lipschitzness(抗扰动能力，函数斜率恒定小于一个Lipschitze常数)，让训练更稳定、更高效。

Gradient Centralization

Motivation

BN和WS使用Z-score标准化分别操作于特征值和权重，实际是间接地对权值的梯度进行约束，从而提高优化时损失函数的Lipschitz属性。受此启发，论文直接对梯度操作，首先尝试了Z-score标准化，但实验发现并没有提升训练的稳定性。之后，尝试计算梯度向量的均值，对梯度向量进行零均值化，实验发现能够有效地提高损失函数的Lipschitz属性，使网络训练更稳定、更具泛化能力，得到梯度中心化(GC)算法。

Notations

定义一些基础符号，使用统一表示全连接层的权值矩阵和卷积层的权值张量，为权值矩阵的第列，为目标函数，和为对和的梯度，与的大小一样。定义为输入特征图，则为输出特征图，为位单位向量(unit vector)，为单位矩阵(identity matrix)。

Formulation of GC

对于卷积层或全连接层的权值向量，通过反向传播得到其梯度，然后如图b所示计算其均值，GC操作定义如下：

也可以将公式1转换为矩阵形式：

由单位矩阵以及单位向量形成矩阵构成，分别负责保留原值以及求均值。

Embedding of GC to SGDM/Adam

GC能够简单地嵌入当前的主流网络优化算法中，如SGDM和Adam，直接使用零均值化的梯度进行权值的更新。

算法1和算法2分别展示了将GC嵌入到SGDM和Adam中，基本上不需要对原优化器算法进行修改，仅需加入一行梯度零均值化计算即可，大约仅需0.6sec。

Properties of GC

下面从理论的角度分析GC为何能提高模型的泛化能力以及加速训练。

Improving Generalization Performance

GC有一个很重要的优点是提高模型的泛化能力，主要得益于权值空间正则化和特征值空间正则化。

Weight space regularization

首先介绍的物理意义，经过推算可以得到：

即可以看作映射矩阵，将映射到空间向量中法向量为的超平面，为映射梯度。

以SGD优化为例，权值梯度的映射能够将权值空间约束在一个超平面或黎曼流形(Riemannian manifold)中，如图2所示，梯度首先映射到的超平面中，然后跟随映射梯度的方向进行更新。从可以得到，目标函数实际变为：

这是一个权值空间的约束优化问题，正则化的解空间，降低了过拟合的可能性(过拟合通常是学习了复杂的权值来适应训练数据)，能够提升网络的泛化能力，特别是当训练样本较少的情况下。
WS对权值进行的约束，当初始权值不满足约束时，会直接修改权值来满足约束条件。假设进行fine-tune训练，WS则会完全丢弃预训练模型的优势，而GC可以适应任何初始权值。