Apple:使用語言模型的自回歸方式進行圖像模型預訓練
1、背景
在GPT等大模型出現后,語言模型這種Transformer+自回歸建模的方式,也就是預測next token的預訓練任務,取得了非常大的成功。那么,這種自回歸建模方式能不能在視覺模型上取得比較好的效果呢?今天介紹的這篇文章,就是Apple近期發表的基于Transformer+自回歸預訓練的方式訓練視覺模型的文章,下面給大家展開介紹一下這篇工作。
圖片
論文標題:Scalable Pre-training of Large Autoregressive Image Models
下載地址:https://arxiv.org/pdf/2401.08541v1.pdf
開源代碼:https://github.com/apple/ml-aim
2、模型結構
模型結構整體采用Transformer,優化目標采用語言模型中的next token prediction。在模型結構上,主要的修改有3個方面。首先是相比ViT這種雙向attention,本文使用的是GPT這種單向attention,即每個位置的元素只能和前面的元素計算attention。
圖片
其次是引入了prefix,在Transformer前面的輸入加了多個prefix token,這部分token使用的是雙向attention。這個主要作用是預訓練和下游應用的一致性,下游都是類似ViT的雙向attention使用方法,預訓練增加prefix雙向attention,可以讓模型適應下游任務。
圖片
最后是在模型最終輸出MLP層的優化。原來的預訓練方法,一般會丟棄掉MLP層,在下游finetune使用一個新的MLP,防止預訓練的MLP過于傾向于預訓練任務導致下游任務效果下降。而本文中,作者采用的方法是每個patch都使用一個獨立的MLP,對于圖像整體的表征,也采用各個patch保證和attention融合的方式,代替一般的pooling,提升預訓練MLP head在下游任務的可用性。
在優化目標上,文中嘗試了兩種方法,第一種是直接擬合patch像素,用MSE進行預測。第二種是提前對圖像patch進行tokenize,轉換成分類任務,用交叉熵損失。不過在文中后續的消融實驗中發現,第二種方法雖然也可以讓模型正常訓練,但是效果并不如基于像素粒度MSE的效果更好。
3、實驗結果
文中的實驗部分詳細分析了這種基于自回歸的圖像模型的效果,以及各個部分對于效果的影響。
首先,隨著訓練的進行,下游的圖像分類任務效果越來越好了,說明這種預訓練方式確實能學到良好的圖像表征信息。
圖片
在訓練數據上,使用小數據集的訓練會導致overfitting,而使用DFN-2B雖然最開始驗證集loss較大,但是沒有明顯的過擬合問題。
圖片
對于模型各個模塊的設計方式,文中也進行了詳細的消融實驗分析。
圖片
在最終的效果對比上,AIM取得了非常不錯的效果,這也驗證了這種自回歸的預訓練方式在圖像上也是可用的,可能會成為后續圖像大模型預訓練的一種主要方式。