[论文解读] How to Fine-Tune Vision Models with SGD

Ananya Kumar, Ruoqi Shen+2 more|arXiv (Cornell University)|Nov 17, 2022

Advanced Neural Network Applications被引用 5

一句话总结

该论文发现，SGD在分布外数据上的视觉变换器（vision transformers）和ConvNeXts模型上表现不如AdamW，原因是嵌入层（embedding layer）的梯度过大。通过冻结这一小规模参数层（参数量不足总参数的1%），SGD在显著降低显存占用的同时，性能达到或超过AdamW，在五个分布外数据基准测试中实现最先进准确率。

ABSTRACT

SGD and AdamW are the two most used optimizers for fine-tuning large neural networks in computer vision. When the two methods perform the same, SGD is preferable because it uses less memory (12 bytes/parameter with momentum and 8 bytes/parameter without) than AdamW (16 bytes/parameter). However, on a suite of downstream tasks, especially those with distribution shifts, we find that fine-tuning with AdamW performs substantially better than SGD on modern Vision Transformer and ConvNeXt models. We find that large gaps in performance between SGD and AdamW occur when the fine-tuning gradients in the first "embedding" layer are much larger than in the rest of the model. Our analysis suggests an easy fix that works consistently across datasets and models: freezing the embedding layer (less than 1% of the parameters) leads to SGD with or without momentum performing slightly better than AdamW while using less memory (e.g., on ViT-L, SGD uses 33% less GPU memory). Our insights result in state-of-the-art accuracies on five popular distribution shift benchmarks: WILDS-FMoW, WILDS-Camelyon, BREEDS-Living-17, Waterbirds, and DomainNet.

研究动机与目标

• 理解为何AdamW在分布外数据上的视觉模型微调中优于SGD。
• 识别现代视觉架构中SGD与AdamW性能差距的根本原因。
• 提出一种最小化、内存高效的SGD改进方法，使其性能达到或超过AdamW。
• 在不增加显存或计算成本的前提下，实现在分布外数据基准测试中的最先进结果。

提出的方法

• 在SGD微调过程中冻结首个嵌入层（参数量不足总参数的1%），以防止因初始高梯度更新导致的过拟合。
• 在多个模型和数据集上，对比SGD在有无动量、有无嵌入层冻结情况下的表现。
• 分析各层在初始化时的梯度范数，识别嵌入层中梯度异常增大的问题。
• 通过参数量的归一化和√(参数数量)进行归一化，评估各层间相对更新幅度。
• 对内存使用情况进行性能分析，量化所提方法相比AdamW的显存优势。
• 在五个分布外数据基准上评估性能：WILDS-FMoW、WILDS-Camelyon、BREEDS-Living-17、Waterbirds和DomainNet。

实验结果

研究问题

• RQ1 为何AdamW在分布外数据上的视觉变换器和ConvNeXts微调中优于SGD？
• RQ2 模型的何种结构或优化特性导致SGD在此类设置下相对于AdamW表现不佳？
• RQ3 尽管参数量很小，冻结嵌入层是否能为SGD微调带来一致的性能提升？
• RQ4 一种无动量的内存高效SGD变体能否在分布外泛化能力上超越标准SGD和AdamW？
• RQ5 冻结嵌入层带来的性能提升是由于减少了梯度爆炸，还是AdamW自适应学习率带来的独立改进？

主要发现

• 在CLIP ViT-B/16上，AdamW相比标准SGD在分布内准确率高出2.1%，在分布外准确率高出8.1%。
• 冻结嵌入层后，SGD性能达到或超过AdamW，在所有模型和数据集上实现76.7%的分布外准确率。
• 在ViT-B/16和ViT-L/14上，AdamW的显存开销分别比SGD（冻结嵌入层，无动量）高出18%和49%。
• SGD在冻结嵌入层且无动量时，实现76.9%的分布外准确率，略高于AdamW和基线SGD。
• 梯度分析显示，在SGD预训练模型中，嵌入层的梯度比其他层大2–3倍；在AdamW预训练模型中，梯度大逾10倍。
• 冻结嵌入层消除了SGD与AdamW之间的性能差距，表明这是SGD表现不佳的主要原因。