当前位置: 首页 > 科技观察

训练ViT和MAE可以减少一半的计算量!Sea和北大联合提出高效优化器Adan,深度模型可以使用

时间:2023-03-22 13:37:01 科技观察

自从谷歌提出VisionTransformer(ViT)以来,ViT逐渐成为很多视觉任务的默认backbone。有了ViT结构,很多视觉任务的SoTA得到进一步提升,包括图像分类、分割、检测、识别等。然而,训练ViT并不是一件容易的事。除了需要更复杂的训练技术外,模型训练的计算量往往比以前的CNN大得多。近日,新加坡SeaAILAB(SAIL)与北京大学零实验室的研究团队联合提出了一种新型深度模型优化器Adan,只需一半的计算量即可完成ViT的训练。论文链接:https://arxiv.org/pdf/2208.06677.pdf代码链接:https://github.com/sail-sg/Adan另外,在相同的计算量下,Adan可以用于多种场景(涉及CV、NLP、RL)、多种训练方式(监督和自监督)、多种网络结构/算法(Swin、ViT、ResNet、ConvNext、MAE、LSTM、BERT、Transformer-XL、PPO算法),两者都有获得了性能改进。代码、配置文件和训练日志都是开源的。深度模型的训练范式和优化器随着ViT的引入,深度模型的训练方法变得越来越复杂。常见的训练技术包括复杂数据增强(如MixUp、CutMix、AutoRand)、标签处理(如标签平滑和噪声标签)、模型参数的移动平均、随机网络深度、dropout等。这些技术的混合使用,模型的泛化性和鲁棒性得到了提升,但随之而来的是模型训练的计算量越来越大。在ImageNet1k上,训练epoch的数量从ResNet刚刚提出的90个增加到300个,这是训练ViT常用的。即使对于一些自监督学习模型,如MAE和ViT,预训练的epoch数量也达到了1.6k。trainingepoch的增加意味着训练时间的大大延长,这会急剧增加学术研究或工业实施的成本。目前常见的解决方案是增加训练的batchsize,辅助并行训练以减少训练时间,但随之而来的问题是大batchsize往往意味着性能下降,batchsize越大越明显情况。这主要是因为随着批量大小的增加,模型参数的更新次数急剧减少。目前的优化器无法实现复杂训练范式下更新次数少的模型的快速训练,这进一步加剧了模型训练epoch数量的增加。因此,是否有一种新的优化器可以用更少的参数更新更快更好地训练深度模型?在减少训练epoch数量的同时,是否也能缓解增加batchsize带来的负面影响?被忽视的冲量加快优化器收敛速度的最直接的方法就是引入冲量。近年来提出的深度模型优化器都遵循Adam中使用的脉冲范式——重球法:其中g_k为随机噪声,m_k为矩,eta为学习率。Adam将m_k的更新由累积形式替换为移动平均形式,并引入二阶矩(n_k)来缩放学习率,即:然而,由于Adam无法训练原始ViT,其改进版本AdamW逐渐成为训练ViT乃至ConvNext的首选。但是,AdamW并没有改变Adam中的脉冲范式,所以当batchsize超过4,096时,AdamW训练的ViT的性能会急剧下降。在传统的凸优化领域,有一种与加权球法齐名的脉冲技术——Nesterov脉冲算法:Nesterov脉冲算法在光滑和一般凸问题上比加权球法具有更快的理论收敛速度,并且可以理论上也可以接受更大的批量。与重球法不同的是,Nesterov算法不计算当前点的梯度,而是利用冲量找到一个外推点,然后在该点计算出梯度后累加冲量。外插点可以帮助Nesterov算法提前感知当前点周围的几何信息。这个特性使得Nesterovimpulse更适合复杂的训练范式和模型结构(比如ViT),因为它不是简单地依靠过去的impulse来绕过sharplocalminima,而是通过提前观察周围的梯度,调整更新的方向。Nesterov脉冲算法虽然具有一定的优势,但在深度优化器中应用和探索较少。主要原因之一是Nesterov算法需要在外推点计算梯度并在当前点更新。在此期间,需要多次重新加载模型参数,并在外推点人为地进行反向传播(BP)。这些不便极大地限制了Nesterov脉冲算法在深度模型优化器中的应用。Adan优化器结合了重写的Nesterov脉冲和自适应优化算法,并引入解耦权重衰减得到最终的Adan优化器。利用外插点,Adan可以提前感知周围的梯度信息,从而有效地逃离尖锐的局部最小区域,增加模型的泛化能力。1)自适应Nesterov冲激为了解决Nesterov冲激算法中多模型参数过载的问题,研究人员首先对Nesterov进行了改写:可以证明改写后的Nesterov冲激算法与原算法等价,迭代点两者可以相互变换,最终收敛点相同。可以看出,通过引入梯度的微分项,可以避免外推点处的人工参数重载和人工BP。将重写的Nesterov脉冲算法与自适应类优化器相结合——将m_k的更新从累积形式替换为移动平均形式,并利用二阶矩对学习率进行缩放:至此,Adan算法的基础已经完成获得的版本。2)梯度差的冲动可以发现,m_k的更新将梯度和梯度的差耦合在一起,但是在实际场景中,往往需要分别处理两个物理意义不同的item,所以研究人员介绍theimpulsev_kofthegradientdifference:这里对梯度的冲量和其差的冲量设置了不同的冲量/平均系数。梯度差项可以在相邻梯度不一致时减慢优化器的更新速度,反之在梯度方向一致时加快更新速度。3)Decoupledweightdecay对于带有L2weightregularization的目标函数,目前流行的AdamW优化器通过解耦L2regularization和trainingloss,在ViT和ConvNext上取得了更好的性能。但是AdamW使用的解耦方法偏向启发式,目前无法得到其收敛的理论保证。基于解耦L2正则的思想,Adan也引入了一种解耦的权重衰减策略。目前Adan的每一次迭代都可以看作是对优化目标F的最小化的一阶近似:因为F中的L2权重正则化过于简单和平滑,所以不需要一阶近似。因此,只能对训练损失做一阶近似而忽略L2权重正则化,则Adan的最后一次迭代将变为:有趣的是,可以发现AdamW的更新准则是Adan的更新准则,当学习率eta接近于0一阶近似。因此,可以从近端算子的角度而不是原来的启发式改进的角度对Adan甚至AdamW进行合理的解释。4)Adan优化器将2)和3)的两个改进结合到Adan的基础版本中,可以得到如下Adan优化器。Adan结合了自适应优化器、Nesterov脉冲和解耦权重衰减策略的优点,可以承受更大的学习率和batchsize,可以实现模型参数的动态L2正则化。5)收敛性分析这里略过复杂的数学分析过程,只给出结论:定理:在给定和不给定Hessian-smooth条件的两种情况下,Adan优化器的收敛速度在非凸随机优化问题上面可以达到已知的理论下界,当存在解耦的权重衰减策略时,结论仍然成立。实验结果1.CV场景1)监督学习——ViT模型对于ViT模型,研究人员测试了Adan在ViT和Swin结构上的表现。可以看出,例如在ViT-small、ViT-base、Swin-tiny和Swin-base上,Adan只消耗了一半的计算资源就获得了接近SoTA优化器的结果,并且在相同的计算量下,Adan在两种ViT模型上都表现出了很大的优势。另外,Adan的性能也在大batchsize下进行了测试:可以看出Adan在各种batchsize下的表现都不错,相对于为大batchsize设计的优化器(LAMB)有一定的优势。.2)监督学习——CNN模型除了难以训练的ViT模型外,研究人员还测试了Adan在sharplocalminima相对较少的CNN模型上的表现——包括经典的ResNet和更先进的ConvNext。结果如下:可以观察到,无论是ResNet还是ConvNext,Adan都可以在大约2/3的训练epoch内实现超越SoTA的性能。3)无监督学习在无监督训练框架下,研究人员测试了Adan在新提出的MAE上的表现。结果如下:与监督学习的结论一致,Adan只消耗了一半的计算量,等于甚至超过了原来的SoTA优化器,而且训练epoch越小,Adan的优势越明显。2.NLP场景1)监督学习在NLP的监督学习任务中,分别观察Adan在经典LSTM和进阶Transformer-XL上的表现。Adan在上述两个网络上表现出一致的优势。而对于Transformer-XL,Adan将默认的Adam优化器与一半的训练步骤相提并论。2)无监督学习为了测试Adan在NLP场景下对无监督任务的模型训练。研究人员从头开始训练BERT:经过1000k次预训练迭代后,在GLUE数据集的7个子任务上测试了Adan训练的模型的性能。结果如下:Adan在测试的7个词分类任务上表现出色。大优势。值得一提的是,Adan训练的BERT-base模型在某些子任务(如RTE、CoLA和SST-2)上的结果甚至超过了Adam训练的BERT-large。将常用的PPO算法中的优化器替换为Adan,并在MuJoCo引擎的4个游戏上测试了Adan的性能。在四场比赛中,使用Adan作为网络优化器的PPO算法总能获得更高的奖励。Adan在RL网络训练方面也展现出巨大的潜力。结论与展望Adan优化器为当前的深度模型优化器引入了一种新的脉冲范式。在复杂的训练范式下以更少的更新实现模型的快速训练。实验表明,Adan仅需1/2-2/3的计算量就可以匹配现有的SoTA优化器。Adan在多场景(涉及CV、NLP、RL)、多种训练方式(监督和自监督)、多种网络结构(ViT、CNN、LSTM、Transformer等)中表现出了极大的性能优势。.此外,Adan优化器的收敛速度已经达到了非凸随机优化的理论下限。