谷歌AI用30億數(shù)據(jù)訓(xùn)練了一個20億參數(shù)Vision Transformer模型，在ImageNet上達(dá)到新的SOTA！
點藍(lán)色字關(guān)注“機(jī)器學(xué)習(xí)算法工程師”
設(shè)為星標(biāo)，干貨直達(dá)！
近日，谷歌研究團(tuán)隊（ViT同名作者）最新的研究論文Scaling Vision Transformers（https://arxiv.org/pdf/2106.04560.pdf），發(fā)布了包含20億參數(shù)的vision transformer模型ViT-G/14，在ImageNet的Top-1可以達(dá)到90.45%，超過之前谷歌提出的Meta Pseduo Labels模型。
但其實谷歌這篇論文的重點是研究vision transformer模型的scaling laws，在NLP領(lǐng)域已經(jīng)有研究（Scaling laws for neural language models）給出了語言模型效果和 compute, data size, model size之間的指數(shù)定律，更有GPT-3這樣成功的模型。雖然已經(jīng)有論文研究（如EfficientNet）CNN模型的scaling strategy：模型增大提升效果。但是在CV領(lǐng)域還是缺少比較全面的研究，而且最近vision transformer的成功應(yīng)用更是需要這樣的工作。
在這篇論文中，實驗?zāi)Ｐ蛥?shù)從5M到2B，訓(xùn)練數(shù)據(jù)量從30M到3B，訓(xùn)練時長從1 TPUv3 core-day到10 000 core-days。這使得谷歌比較全面地研究ViT模型效果和model size，data size和compute之間的scaling laws。
論文中采用的是ViT模型，并增加了G/14這樣超大的模型（接近2B)，不同大小的模型如下所示，主要區(qū)別在patch size以及transformer layer的參數(shù)配置。
在數(shù)據(jù)方面，谷歌又拋出重磅炸彈：JFT-3B，這個是JFT-300M的超大版本，包含接近30億的圖像，標(biāo)注為包含30k類別的層級類別，由于采用半自動標(biāo)注，所以標(biāo)注是有噪音的。具體到訓(xùn)練模型，直接采用sigmoid cross-entropy損失訓(xùn)練多分類模型，忽略類別間層級關(guān)系。所有的測試數(shù)據(jù)均從JFT-3B中移除。
據(jù)此，實驗分別改變architecture size，number of training images和training duration來測試模型的representation quality，具體可以用(i) few-shot transfer via training a linear classifier on frozen weights, (ii) transfer via fine-tuning the whole model on all data, both to multiple benchmark tasks來作為評價指標(biāo)，下圖是在ImageNet數(shù)據(jù)集上的結(jié)果：
最終的結(jié)論主要有三點：
scaling up compute, model and data together improves representation quality：同時增大模型，訓(xùn)練數(shù)據(jù)和計算力是可以同步提升效果的，近似滿足指數(shù)定律（log-log曲線是線性的）；
representation quality can be bottlenecked by model size：模型大小會限制上限，小的模型即使用再大的數(shù)據(jù)集也難以繼續(xù)提升；
large models benefit from additional data, even beyond 1B images：對于大模型來說，訓(xùn)練數(shù)據(jù)會制約性能上限，對于最大的模型，訓(xùn)練數(shù)據(jù)量從1B提升至3B，效果還有提升。
雖然從實驗看來，ViT模型也滿足指數(shù)scaling定律，但是其實是double-saturating power law：對于最小的模型，其效果會比指數(shù)定律預(yù)測值要好，這是因為模型效果也有下限；而最大的模型即使給再多的訓(xùn)練數(shù)據(jù)和算力也不會達(dá)到0 error，模型也有上限。所以出現(xiàn)了雙飽和。
除了上述結(jié)論，論文還發(fā)現(xiàn)了額外的結(jié)論：bigger models are more sample efficient，即大的模型需要更少的unseen數(shù)據(jù)就可以達(dá)到和小模型類似的效果，如下圖所示：
最大的模型ViT-G/14，接近2B參數(shù)，在ImageNet上finetune后top-1 acc可達(dá)90.45%，而且在few-shot learning也表現(xiàn)優(yōu)異，每個類只用10個圖像訓(xùn)練就能在ImageNet上達(dá)到84.86%。
另外，在訓(xùn)練ViT模型，論文中還設(shè)計了一些訓(xùn)練策略來提升內(nèi)存利用和模型效果，這些策略也使得ViT-G/14可以采用數(shù)據(jù)并行訓(xùn)練策略，這意味著ViT-G/14模型可以放在一張TPUv3 core。具體策略包括：
Decoupled weight decay for the “head”：模型的head（分類的linear層）和模型的主體body部分采用不同的weight decay，論文中發(fā)現(xiàn)head采用較大的weight decay有助于遷移到下游任務(wù)（可能和SVM類似，拉大類間間距）；
Saving memory by removing the [class] token：采用multihead attention pooling (MAP) 來替換class token，class token會使得TPU需要padding而增加50%內(nèi)存使用；
Memory-efficient optimizers：對于Adam優(yōu)化器，采用half-precision momentum，并采用Adafactor優(yōu)化器（進(jìn)行了修改）來進(jìn)一步減少內(nèi)存使用；
Learning-rate schedule：學(xué)習(xí)速率learning-rate schedules引入cooldown階段（學(xué)習(xí)速率逐漸降為0），這樣可以一次訓(xùn)練就可以得到不同訓(xùn)練時長的模型。
谷歌做這個實驗的代價自然不必說，也有很多人在質(zhì)疑這個研究的意義，但是論文也在最后給出了這個工作的意義，首先就是這個scaling laws做出來后對后續(xù)研究是有啟發(fā)意義的：
First, such studies of scaling laws need only be performed once; future developers of ViT models may use our results to design models that can be trained with fewer compute resources.
此外，這個預(yù)訓(xùn)練模型可以遷移到其它任務(wù)：
Second, the models trained are designed primarily for transfer learning. Transfer of pre-trained weights is much less expensive than training from scratch on a downstream task, and typically reaches higher accuracy.
作為一個AI大廠，Google做這么大的work，我個人還是持肯定態(tài)度。
推薦閱讀
DETR：基于 Transformers 的目標(biāo)檢測
"未來"的經(jīng)典之作ViT：transformer is all you need！
PVT：可用于密集任務(wù)backbone的金字塔視覺transformer！
漲點神器FixRes：兩次超越ImageNet數(shù)據(jù)集上的SOTA
不妨試試MoCo，來替換ImageNet上pretrain模型！
機(jī)器學(xué)習(xí)算法工程師
一個用心的公眾號
