梯度下降优化算法综述与PyTorch实现源码剖析

现代的机器学习系统均利用大量的数据，利用梯度下降算法或者相关的变体进行训练。传统上，最早出现的优化算法是SGD，之后又陆续出现了AdaGrad、RMSprop、ADAM等变体，那么这些算法之间又有哪些区别和联系呢？本文试图对比的介绍目前常用的基于一阶梯度的优化算法，并给出它们的(PyTorch)实现。

SGD

算法描述

随机梯度下降法(Stochastic Gradient Descent，SGD)是对传统的梯度下降算法(Gradient Descent,GD)进行的一种改进。在应用GD时，我们需要对整个训练集进行一次反向传播计算梯度后再进行参数更新，对系统的计算能力和内存的需求较高，而SGD在计算梯度更新参数时刚好相反，每次只使用整个训练集中的一个样本，因此具有更快地计算速度和较少的内存占用。同时，因为每次只使用一个样本更新参数，使得参数更新更加频繁，更新的参数间具有更高的方差，损失函数会朝不同的方向有较大的波动，这有助于发现新的极值点，避免优化器陷入一个局部极值点。但是也由于这种频繁的震荡，出现了一种折中的方法，即小批量(mini-batch)梯度下降法，每次只取训练集中一个batch的样本进行梯度的计算与参数更新，一般batch的大小为4的倍数。原始SGD的更新法则如下：θ=θ−η⋅∇θJ(θ)(1)(1)θ=θ−η⋅∇θJ(θ)

传统SGD面临的问题

传统的SGD在训练的过程中主要存在以下几个问题：

1. 很难选择一个合适的学习速率，太小的学习速率导致算法收敛很慢，而太大的学习速率会导致在极值点附近震荡甚至错过，因此需要经过多次尝试。
2. Learning rate schedules往往实现定义一个学习速率衰减表，比如每过多少step对学习速率进行decay，但是这些策略往往没法按照某个数据集的具体参数特性进行定制。
3. 对于比较稀疏的数据，不同的特征出现的频率差别很大，如果所有的参数均使用一个相同的学习速率进行更新，这样做是不合理的。对于出现频率的特征，我们应该使用一个较大的学习速率。
4. 深度神经网络之所以难以训练，并不是因为容易陷入局部最小值，而是在学习的过程中陷入到鞍点(saddle point)，此时往各个方向的梯度几乎均为0。如果以二维平面为例，y=x3y=x3中x=0处即为一个鞍点。对于传统的SGD而言，一旦优化的过程中进入鞍点，就很难再离开这一位置。

Momentum

针对以上提到的第四点问题，可以通过增加动量(Momentum)的SGD进行缓解，加速优化函数的收敛。vt=γvt−1−η⋅∇θJ(θ)θ=θ+vt(2)(2)vt=γvt−1−η⋅∇θJ(θ)θ=θ+vt所谓的添加动量项，即在一定程度上保留上一次梯度更新的方向，γ,ηγ,η分别用来控制上次梯度方向和本次梯度方向对最终更新方向的贡献程度，其中γ∈(0,1]γ∈(0,1]在开始阶段常常被设置为0.5，当学习趋向稳定后，逐渐增加到0.9甚至更高。 可以把待优化的目标函数想象成一座山，在山顶将一个小球推下，小球在山坡上滚动的位置即系统的loss值，在往下滚动的过程中小球的动量不断增加，由于动量的存在，当小球滚动到山坡中较为平坦的地带时，小球将更容易越过这片地带继续往下滚而不是陷在这一区域停滞不前，并最终到达山谷。

图1 左：原始SGD 右：SGD+Momentum

Nesterov Accelerated Gradient

Its better to correct a mistake after you have made it!

目前我们有了一个带有动量的小球，但是这个小球在滚动的过程中总是随着山势的变化滚动，因此其行进的路径极不稳定。因此我们希望有一个更加“聪明”的小球，它不但拥有动量，而且能够知道自己将要去哪，这样当前面出现上坡小球能够进行减速。比如说，当接近坡底时，小球应该提前减速避免错过坡底。vt=γvt−1−η∇θJ(θ+γvt−1)θ=θ+vt(3)(3)vt=γvt−1−η∇θJ(θ+γvt−1)θ=θ+vt具体的实现也非常的直接，就是将传统的Momentum方法对θθ计算梯度变为对θ+γvt−1θ+γvt−1求梯度，这一项可以看做对小球下一步将会往哪运动的一个粗略估计。也就是说，我们的小球有了一定的对未来的“预测”能力。就像本节开头说的，如果我们知道了小球之后会犯什么错误，那么是否更容易更正错误呢？下图上半部分是传统Momentum求下一次梯度更新方向，下半部分则是使用NAG求下一次更新方向的方法。

图2 Momentum与NAG更新的区别

当然，在具体实现时，直接计算θ+γvt−1θ+γvt−1项的梯度比较麻烦，希望更新参数时计算能和传统的SGD或者Momentum方法类似，因此需要对上式的计算步骤做一些改进。

v_prev = v  #备份vt-1项
v = mu*v - lr * g  #这一步和传统的Momentum计算一样
p += -mu*v_prev + (1+mu)*v  #更新时真实的p应该为p-mu*v_prev，更新后为p-mu*v_prev+v，但是为了方便计算加上上次动量项的梯度，这里的p直接保存为p-mu*v_prev+v+mu*v,也就是p(小球)的“未来位置”。

PyTorch实现

Momentum/NAG的实现和原始论文中的实现有些许的不用，具体的，在PyTorch实现中按照如下的公式更新梯度,其中ηη为learning rate,gg为θθ的梯度。目前尚不清楚为什么要做出这样的改变？vt=γvt−1+gθ=θ−η⋅vt(4)(4)vt=γvt−1+gθ=θ−η⋅vt具体代码如下： > class torch.optim.SGD(params, lr=required, momentum=0, dampening=0, weight_decay=0, nesterov=False)

def step(self, closure=None):
    """Performs a single optimization step.

    Arguments:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        loss = closure()

    for group in self.param_groups:
        weight_decay = group['weight_decay']
        momentum = group['momentum']
        dampening = group['dampening']
        nesterov = group['nesterov']

        for p in group['params']:
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            if momentum != 0:   #动量项添加
                param_state = self.state[p]
                if 'momentum_buffer' not in param_state:
                    buf = param_state['momentum_buffer'] = d_p.clone()
                else:
                    buf = param_state['momentum_buffer']
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:    #如果使用NAG，则为t+1步先保存可能到达的“位置”
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf

            p.data.add_(-group['lr'], d_p)

    return loss

AdaGrad

算法描述

AdaGrad为的是解决传统的SGD对所有参数使用相同的学习速率的问题(即1.2节中提到的第三点问题)。它使用参数的历史梯度累计和去归一化该参数对应的学习速率。具体的，对于经常出现的参数，那么其梯度累积和较大，归一化的学习速率就较小。而对于不常见的参数，往往包含更多关于特征的信息，累积和较小，归一化后的学习速率较大，也即是学习算法应该更加关注这些罕见的特征的出现。Gt,ii=Gt−1,ii+g2t,iθt+1,i=θt,i−η√Gt,ii+ϵ⋅gt,i(5)(5)Gt,ii=Gt−1,ii+gt,i2θt+1,i=θt,i−ηGt,ii+ϵ⋅gt,i当然，通过观察式(5)，我们也发现AdaGrad在学习速率的调整上存在过于激进的问题，随着时间的累积，Gt,iiGt,ii这一项会越来越大，导致归一化的学习速率越来越小，这有可能导致优化函数在收敛之前就停止更新。

PyTorch实现

class torch.optim.Adagrad(params, lr=0.01, lr_decay=0, weight_decay=0)

def step(self, closure=None):
        """Performs a single optimization step.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = None
        if closure is not None:
            loss = closure()

        for group in self.param_groups:
            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad.data
                state = self.state[p]

                state['step'] += 1

                if group['weight_decay'] != 0:
                    if p.grad.data.is_sparse:
                        raise RuntimeError("weight_decay option is not compatible with sparse gradients ")
                    grad = grad.add(group['weight_decay'], p.data)

                clr = group['lr'] / (1 + (state['step'] - 1) * group['lr_decay'])

                if p.grad.data.is_sparse:
                    grad = grad.coalesce()  # the update is non-linear so indices must be unique
                    grad_indices = grad._indices()
                    grad_values = grad._values()
                    size = torch.Size([x for x in grad.size()])

                    def make_sparse(values):
                        constructor = type(p.grad.data)
                        if grad_indices.dim() == 0 or values.dim() == 0:
                            return constructor()
                        return constructor(grad_indices, values, size)
                    state['sum'].add_(make_sparse(grad_values.pow(2)))
                    std = state['sum']._sparse_mask(grad)
                    std_values = std._values().sqrt_().add_(1e-10)
                    p.data.add_(-clr, make_sparse(grad_values / std_values))
                else:
                    state['sum'].addcmul_(1, grad, grad)    #更新核心部分
                    std = state['sum'].sqrt().add_(1e-10)
                    p.data.addcdiv_(-clr, grad, std)

        return loss

Adadelta

为了避免AdaGrad存在的学习过早停止的问题，Adadelta不再保存过去所有时刻的梯度和，而是采用decaying average的方法平滑过去的梯度值和参数值。

算法描述

图3 Adadelta伪代码描述

其中E[g2]tE[g2]t存储的是历史梯度平方的平滑值，此外，这里还需要对历史的参数值的平方进行decaying average，也就是E[Δx2]t=ρE[Δx2]t−1+(1−ρ)Δx2tE[Δx2]t=ρE[Δx2]t−1+(1−ρ)Δxt2，然后分别进行开方处理得到RMS[Δx]t=√E[Δx2]t+ϵRMS[Δx]t=E[Δx2]t+ϵ和RMS[g]t=√E[g2]t+ϵRMS[g]t=E[g2]t+ϵ。最后按照xt+1=xt−RMS[Δx]t−1RMS[g]t⋅gtxt+1=xt−RMS[Δx]t−1RMS[g]t⋅gt更新参数。 从以上的算法描述可以看出，Adadelta并不需要是先定义一个全局的学习速率，而是可以根据参数自身的特点计算一个合理的学习速率进行更新。

PyTorch实现

代码的实现很简单，也是完全按照上图描述的流程进行计算和更新。 > class torch.optim.Adadelta(params, lr=1.0, rho=0.9, eps=1e-06, weight_decay=0)

def step(self, closure=None):
    """Performs a single optimization step.

    Arguments:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        loss = closure()

    for group in self.param_groups:
        for p in group['params']:
            if p.grad is None:
                continue
            grad = p.grad.data
            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state['step'] = 0
                state['square_avg'] = grad.new().resize_as_(grad).zero_()
                state['acc_delta'] = grad.new().resize_as_(grad).zero_()

            square_avg, acc_delta = state['square_avg'], state['acc_delta']
            rho, eps = group['rho'], group['eps']

            state['step'] += 1

            if group['weight_decay'] != 0:
                grad = grad.add(group['weight_decay'], p.data)

            square_avg.mul_(rho).addcmul_(1 - rho, grad, grad)  #更新核心部分
            std = square_avg.add(eps).sqrt_()
            delta = acc_delta.add(eps).sqrt_().div_(std).mul_(grad)
            p.data.add_(-group['lr'], delta)
            acc_delta.mul_(rho).addcmul_(1 - rho, delta, delta)

    return loss

RMSprop

算法描述

RMSprop算法并没有被正式的发表，而是Geoff Hinton在他的课程中提及。RMSprop是一种十分高效的算法，可以看作是对AdaGrad算法的改进，对历史的梯度信息使用decaying average的方式进行累计，在学习速率的处理上不再像AdaGrad那么激进。E[g2]t=ρE[g2]t−1+(1−ρ)g2tθt+1=θt−η√E[g2]t+ϵgt(6)(6)E[g2]t=ρE[g2]t−1+(1−ρ)gt2θt+1=θt−ηE[g2]t+ϵgtRMSprop也可以和传统的momentum方法结合，但是Hinton表示这样做的帮助不是很大(相对于其对传统SGD的帮助而言)，具体的原因需要更多的研究。当然如果和Nesterov momentum结合能够有更好的效果。

PyTorch实现

class torch.optim.RMSprop(params, lr=0.01, alpha=0.99, eps=1e-08, weight_decay=0, momentum=0, centered=False)

在PyTorch实现中，有一个centered的标志，即是否使用centered版本的RMSprop。 centered版本的的RMSprop按照E[g2]t=ρE[g2]t−1+(1−ρ)g2tE[g]t=ρE[g]t−1+(1−ρ)gtθt+1=θt−η√E[g2]t−E[g]2t+ϵgt(7)更新梯度。(7)E[g2]t=ρE[g2]t−1+(1−ρ)gt2E[g]t=ρE[g]t−1+(1−ρ)gtθt+1=θt−ηE[g2]t−E[g]t2+ϵgt

def step(self, closure=None):
    """Performs a single optimization step.

    Arguments:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        loss = closure()

    for group in self.param_groups:
        for p in group['params']:
            if p.grad is None:
                continue
            grad = p.grad.data
            state = self.state[p]

            # State initialization
            if len(state) == 0:
                state['step'] = 0
                state['square_avg'] = grad.new().resize_as_(grad).zero_()
                if group['momentum'] > 0:
                    state['momentum_buffer'] = grad.new().resize_as_(grad).zero_()
                if group['centered']:
                    state['grad_avg'] = grad.new().resize_as_(grad).zero_()

            square_avg = state['square_avg']
            alpha = group['alpha']

            state['step'] += 1

            if group['weight_decay'] != 0:
                grad = grad.add(group['weight_decay'], p.data)

            square_avg.mul_(alpha).addcmul_(1 - alpha, grad, grad)

            if group['centered']:   #使用centered RMSprop
                grad_avg = state['grad_avg']
                grad_avg.mul_(alpha).add_(1 - alpha, grad)
                avg = square_avg.addcmul(-1, grad_avg, grad_avg).sqrt().add_(group['eps'])
            else:
                avg = square_avg.sqrt().add_(group['eps'])

            if group['momentum'] > 0:   #添加动量项
                buf = state['momentum_buffer']
                buf.mul_(group['momentum']).addcdiv_(grad, avg)
                p.data.add_(-group['lr'], buf)
            else:
                p.data.addcdiv_(-group['lr'], grad, avg)

    return loss

ADAM

ADAM启发自AdaGrad和RMSProp两种优化算法，是一种利用一阶梯度的随机优化算法。它利用梯度的一阶和二阶矩估计为每个参数计算自适应的学习速率。这种优化算法具有容易实现，计算效率高，内存要求低的特点，尤其适合具有大量参数或数据维度较高的函数的优化。同时，ADAM也适合非平稳目标(non-stationary objectives)或者梯度非常嘈杂或者稀疏的目标的优化。优化器的超参数具有非常直观的解释，通常不需要过多的调试就可以获得一组比较好的参数设置。

算法描述

图4 ADAM算法伪代码描述

图4是ADAM算法的具体实现方式，其中αα为梯度更新初始化的步伐。β1,β2∈[0,1)β1,β2∈[0,1)为两个矩估计的衰减指数。f(θ)f(θ)为需要优化的目标函数，其中参数为θθ。 θ0θ0 为初始时的参数状态。 m0,v0m0,v0分别为一阶和二阶矩向量，初始为0。tt为迭代的步数，每次加1。 在具体计算的过程中，首先利用反向传播计算t时刻参数θθ的梯度gtgt，然后分别利用mt←β1⋅mt−1+(1−β1)⋅gtmt←β1⋅mt−1+(1−β1)⋅gt和vt←β2vt←β2⋅vt−1+(1−β2)⋅gt2m^t←mt/(1−β1t)v^t←vt/(1−β2t)β1t,β2t(8)θt←θt−1−α⋅mt^/(vt^+ϵ)⋅vt−1+(1−β2)⋅gt2分别更新梯度的一二阶矩估计。之后，利用^mt←mt/(1−β1t)和^vt←vt/(1−β2t)对估计值进行偏差校正（因为m和v的初始值为0，因此其矩估计是有偏的且偏向于0，尤其是在梯度下降更新开始的一段时间内，衰减系数βt1,βt2都趋向于1，这一步有助于对偏差进行校正）。最终，更新后的参数值由式(8)确定θt←θt−1−α⋅^mt/(√^vt+ϵ)(8)AdaMax是基于infinity norm的Adam算法的变体，在原始的Adam中，对梯度二阶矩按照

PyTorch实现

class torch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

其中具体实现利用了论文中提出的一种等效但是更高效的实现方式，即把图4的算法中最后三行替换为如下的计算方式。

(9)αt=α⋅1−β2t/(1−β1t)θt←θt−1−αt⋅mt/(vt+ϵ^)αt=α⋅√1−β2t/(1−β1t)θt←θt−1−αt⋅mt/(√vt+^ϵ)(9)

def step(self, closure=None): """Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data state = self.state[p] # State initialization if len(state) == 0: state['step'] = 0 # Exponential moving average of gradient values state['exp_avg'] = grad.new().resize_as_(grad).zero_() # Exponential moving average of squared gradient values state['exp_avg_sq'] = grad.new().resize_as_(grad).zero_() exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq'] beta1, beta2 = group['betas'] state['step'] += 1 if group['weight_decay'] != 0: grad = grad.add(group['weight_decay'], p.data) # Decay the first and second moment running average coefficient exp_avg.mul_(beta1).add_(1 - beta1, grad) #更新梯度的一阶矩 exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad) #更新梯度的二阶矩 denom = exp_avg_sq.sqrt().add_(group['eps']) bias_correction1 = 1 - beta1 ** state['step'] #偏差校正 bias_correction2 = 1 - beta2 ** state['step'] step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1 #step_size可以看作是计算自适应的学习速率 p.data.addcdiv_(-step_size, exp_avg, denom) return loss

AdaMax

算法描述

图5 AdaMax算法伪代码描述

l2(10)vt=β2vt−1+(1−β2)gt2lpvt(11)vt=β2pvt−1+(1−β2p)gtpl∞(12)vt=β2∞vt−1+(1−β2∞)|gt|∞=max(β2⋅vt−1,|gt|)l2更新。vt=β2vt−1+(1−β2)g2t(10)现在我们按照如下的lp法则更新vtvt=βp2vt−1+(1−βp2)gpt(11)使用大的p值进行规范化往往会导致数值不稳定(numerically unstable)，常用的p值为1或2。但是l∞也可以得到稳定结果。基于此就得到了AdaMax算法。vt=β∞2vt−1+(1−β∞2)|gt|∞=max(β2⋅vt−1,|gt|)(12)图6展示了从损失函数面上的任意一点到极值点的时间。可以看出，Adagrad,Adadelta和RMSprop三种方法几乎理科按照正确的方向到达了极值点，NAG和Momentum方法一开始在前进方向上出现了错误，但是NAG方法更快地摆脱了错误并到达极值点。而传统的SGD则最慢。图7显示了到算法位于一个鞍点时摆脱困境并到达极值点的速度。可以看出，传统的SGD极难摆脱这种情况，而NAG和Momentum虽然在鞍点位置对称震荡了比较长一段时间，但是最终还是到达了极值点。Adagrad,Adadelta和RMSprop三种方法则很快摆脱了鞍点并到达了极值点。

PyTorch实现

class torch.optim.Adamax(params, lr=0.002, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)

def step(self, closure=None): """Performs a single optimization step. Arguments: closure (callable, optional): A closure that reevaluates the model and returns the loss. """ loss = None if closure is not None: loss = closure() for group in self.param_groups: for p in group['params']: if p.grad is None: continue grad = p.grad.data state = self.state[p] # State initialization if len(state) == 0: state['step'] = 0 state['exp_avg'] = grad.new().resize_as_(grad).zero_() state['exp_inf'] = grad.new().resize_as_(grad).zero_() exp_avg, exp_inf = state['exp_avg'], state['exp_inf'] beta1, beta2 = group['betas'] eps = group['eps'] state['step'] += 1 if group['weight_decay'] != 0: grad = grad.add(group['weight_decay'], p.data) # Update biased first moment estimate. exp_avg.mul_(beta1).add_(1 - beta1, grad) # Update the exponentially weighted infinity norm. norm_buf = torch.cat([ exp_inf.mul_(beta2).unsqueeze(0), grad.abs().add_(eps).unsqueeze_(0) ], 0) torch.max(norm_buf, 0, keepdim=False, out=(exp_inf, exp_inf.new().long())) #与ADAM最主要的差别 bias_correction = 1 - beta1 ** state['step'] clr = group['lr'] / bias_correction #计算自适应学习速率 p.data.addcdiv_(-clr, exp_avg, exp_inf) return loss

算法可视化

图6和图7(Image credit: )展示了不同的算法在优化目标函数时的行为。

Alec Radford

图6 不同算法收敛速度比较

图7 不同算法摆脱鞍点的速度比较

Which optimizer to use?

这一段主要参考了Sebastian Ruder的。既然有这么多的优化算法，那么在实际使用时我们应该选择使用哪个呢？如果你的输入数据比较稀疏，那么最好选择一个具有自适应学习速率的算法，这样使用一个默认的学习速率往往也能取得一个较好的效果。 总的来说，RMSprop是AdaGrad的一个扩展，用来处理后期学习速率急剧下降的问题。Adadelta和RMSprop类似，不过历史参数值的引入使其甚至不需要设置一个初始的全局学习速率。最后Adam在RMSprop的基础上增加了偏差校正( bias-correction)和momentum。RMSprop、Adadelta和Adam是三种比较相似的算法，但是往往在优化的末期梯度十分稀疏的时候Adam的效果更好。因此，在一般使用时Adam是首选。

blog

参考文献

Optimization: Stochastic Gradient DescentAn overview of gradient descent optimization algorithmsOn the importance of initialization and momentum in deep learningCS231n Convolutional Neural Networks for Visual RecognitionAdaptive Subgradient Methods for Online Learning and Stochastic OptimizationBeyond SGD: Gradient Descent with Momentum and Adaptive Learning RateADADELTA: An Adaptive Learning Rate MethodLecture 6a Overview of mini-batch gradient descentAdam: A Method for Stochastic Optimization