1、背景

在GPT等大模型出現(xiàn)后，語言模型這種Transformer+自回歸建模的方式，也就是預(yù)測next token的預(yù)訓(xùn)練任務(wù)，取得了非常大的成功。那么，這種自回歸建模方式能不能在視覺模型上取得比較好的效果呢？今天介紹的這篇文章，就是Apple近期發(fā)表的基于Transformer+自回歸預(yù)訓(xùn)練的方式訓(xùn)練視覺模型的文章，下面給大家展開介紹一下這篇工作。

圖片

論文標(biāo)題：Scalable Pre-training of Large Autoregressive Image Models

下載地址：https://arxiv.org/pdf/2401.08541v1.pdf

開源代碼：https://github.com/apple/ml-aim

2、模型結(jié)構(gòu)

模型結(jié)構(gòu)整體采用Transformer，優(yōu)化目標(biāo)采用語言模型中的next token prediction。在模型結(jié)構(gòu)上，主要的修改有3個(gè)方面。首先是相比ViT這種雙向attention，本文使用的是GPT這種單向attention，即每個(gè)位置的元素只能和前面的元素計(jì)算attention。

圖片

其次是引入了prefix，在Transformer前面的輸入加了多個(gè)prefix token，這部分token使用的是雙向attention。這個(gè)主要作用是預(yù)訓(xùn)練和下游應(yīng)用的一致性，下游都是類似ViT的雙向attention使用方法，預(yù)訓(xùn)練增加prefix雙向attention，可以讓模型適應(yīng)下游任務(wù)。

圖片

最后是在模型最終輸出MLP層的優(yōu)化。原來的預(yù)訓(xùn)練方法，一般會(huì)丟棄掉MLP層，在下游finetune使用一個(gè)新的MLP，防止預(yù)訓(xùn)練的MLP過于傾向于預(yù)訓(xùn)練任務(wù)導(dǎo)致下游任務(wù)效果下降。而本文中，作者采用的方法是每個(gè)patch都使用一個(gè)獨(dú)立的MLP，對(duì)于圖像整體的表征，也采用各個(gè)patch保證和attention融合的方式，代替一般的pooling，提升預(yù)訓(xùn)練MLP head在下游任務(wù)的可用性。

在優(yōu)化目標(biāo)上，文中嘗試了兩種方法，第一種是直接擬合patch像素，用MSE進(jìn)行預(yù)測。第二種是提前對(duì)圖像patch進(jìn)行tokenize，轉(zhuǎn)換成分類任務(wù)，用交叉熵?fù)p失。不過在文中后續(xù)的消融實(shí)驗(yàn)中發(fā)現(xiàn)，第二種方法雖然也可以讓模型正常訓(xùn)練，但是效果并不如基于像素粒度MSE的效果更好。

3、實(shí)驗(yàn)結(jié)果

文中的實(shí)驗(yàn)部分詳細(xì)分析了這種基于自回歸的圖像模型的效果，以及各個(gè)部分對(duì)于效果的影響。

首先，隨著訓(xùn)練的進(jìn)行，下游的圖像分類任務(wù)效果越來越好了，說明這種預(yù)訓(xùn)練方式確實(shí)能學(xué)到良好的圖像表征信息。

圖片

在訓(xùn)練數(shù)據(jù)上，使用小數(shù)據(jù)集的訓(xùn)練會(huì)導(dǎo)致overfitting，而使用DFN-2B雖然最開始驗(yàn)證集loss較大，但是沒有明顯的過擬合問題。

圖片

對(duì)于模型各個(gè)模塊的設(shè)計(jì)方式，文中也進(jìn)行了詳細(xì)的消融實(shí)驗(yàn)分析。

圖片

在最終的效果對(duì)比上，AIM取得了非常不錯(cuò)的效果，這也驗(yàn)證了這種自回歸的預(yù)訓(xùn)練方式在圖像上也是可用的，可能會(huì)成為后續(xù)圖像大模型預(yù)訓(xùn)練的一種主要方式。

圖片